add filter by tag

Author: rcmangasaryan

docs/source/superannotate.sdk.rst

@@ -341,6 +341,7 @@ ________________________
 .. _ref_filter_annotation_instances:
 .. autofunction:: superannotate.filter_annotation_instances
 .. autofunction:: superannotate.filter_images_by_comments
+.. autofunction:: superannotate.filter_images_by_tags
 
 ----------
 

superannotate/__init__.py

@@ -22,7 +22,8 @@
     project_type_str_to_int, user_role_str_to_int
 )
 from .dataframe_filtering import (
-    filter_annotation_instances, filter_images_by_comments
+    filter_annotation_instances, filter_images_by_comments,
+    filter_images_by_tags
 )
 from .db.annotation_classes import (
     create_annotation_class, create_annotation_classes_from_classes_json,

superannotate/analytics/common.py

@@ -9,7 +9,7 @@
 
 def df_to_annotations(df, output_dir):
     """Converts and saves pandas DataFrame annotation info (see aggregate_annotations_as_df) in output_dir
-    The DataFrame should have columns: "imageName", "classNmae", "attributeGroupName", "attributeName", "type", "error", "locked", "visible", trackingId", "probability", "pointLabels", "meta", "commentResolved", "classColor", "groupId"
+    The DataFrame should have columns: "imageName", "className", "attributeGroupName", "attributeName", "type", "error", "locked", "visible", trackingId", "probability", "pointLabels", "meta", "commentResolved", "classColor", "groupId"
 
     Currently only works for Vector projects.
 
@@ -120,6 +120,7 @@ def aggregate_annotations_as_df(
     project_root,
     include_classes_wo_annotations=False,
     include_comments=False,
+    include_tags=False,
     verbose=True
 ):
     """Aggregate annotations as pandas dataframe from project root.
@@ -154,7 +155,6 @@ def aggregate_annotations_as_df(
         "error": [],
         "locked": [],
         "visible": [],
-        "tag": [],
         "trackingId": [],
         "probability": [],
         "pointLabels": [],
@@ -172,6 +172,8 @@ def aggregate_annotations_as_df(
 
     if include_comments:
         annotation_data["commentResolved"] = []
+    if include_tags:
+        annotation_data["tag"] = []
 
     classes_path = Path(project_root) / "classes" / "classes.json"
     if not classes_path.is_file():
@@ -245,13 +247,14 @@ def __get_image_metadata(image_name, annotations):
                     __append_annotation(annotation_dict)
                 continue
             if annotation_type == "tag":
-                annotation_tag = annotation["name"]
-                annotation_dict = {
-                    "type": annotation_type,
-                    "tag": annotation_tag
-                }
-                annotation_dict.update(image_metadata)
-                __append_annotation(annotation_dict)
+                if include_tags:
+                    annotation_tag = annotation["name"]
+                    annotation_dict = {
+                        "type": annotation_type,
+                        "tag": annotation_tag
+                    }
+                    annotation_dict.update(image_metadata)
+                    __append_annotation(annotation_dict)
                 continue
             annotation_instance_id += 1
             annotation_class_name = annotation.get("className")

superannotate/dataframe_filtering.py

@@ -42,6 +42,40 @@ def filter_images_by_comments(
     return list(images)
 
 
+def filter_images_by_tags(annotations_df, include=None, exclude=None):
+    """Filter images on tags
+
+    :param annotations_df: pandas DataFrame of project annotations
+    :type annotations_df: pandas.DataFrame
+    :param include: include images with given tags
+    :type include: list of strs
+    :param exclude: exclude images with given tags
+    :type exclude: list of strs
+
+    :return: filtered image names
+    :rtype: list of strs
+
+    """
+
+    df = annotations_df[annotations_df["type"] == "tag"]
+    images = set(df["imageName"].dropna().unique())
+
+    if include:
+        include_images = set(
+            df[df["tag"].isin(include)]["imageName"].dropna().unique()
+        )
+        images = images.intersection(include_images)
+
+    if exclude:
+        exclude_images = set(
+            df[df["tag"].isin(exclude)]["imageName"].dropna().unique()
+        )
+
+        images = images.difference(exclude_images)
+
+    return list(images)
+
+
 def filter_annotation_instances(annotations_df, include=None, exclude=None):
     """Filter annotation instances from project annotations pandas DataFrame.