Fix df filtering

Author: hovnatan

superannotate/analytics/common.py

@@ -1,10 +1,12 @@
+import glob
 import json
 import logging
 from pathlib import Path
-import glob
 
 import pandas as pd
+
 from ..exceptions import SABaseException
+
 logger = logging.getLogger("superannotate-python-sdk")
 
 
@@ -29,7 +31,12 @@ def df_to_annotations(df, output_dir):
         image_height = None
         image_width = None
         image_df = df[df["imageName"] == image]
-        image_annotation = []
+        image_annotation = {
+            "instances": [],
+            "metadata": {},
+            "tags": [],
+            "comments": []
+        }
         instances = image_df["instanceId"].dropna().unique()
         for instance in instances:
             instance_df = image_df[image_df["instanceId"] == instance]
@@ -64,27 +71,31 @@ def df_to_annotations(df, output_dir):
                             "name": row["attributeName"]
                         }
                     )
-            image_annotation.append(instance_annotation)
+            image_annotation["instances"].append(instance_annotation)
             image_width = image_width or instance_df.iloc[0]["imageWidth"]
             image_height = image_height or instance_df.iloc[0]["imageHeight"]
             image_pinned = image_pinned or instance_df.iloc[0]["imagePinned"]
             image_status = image_status or instance_df.iloc[0]["imageStatus"]
 
         comments = image_df[image_df["type"] == "comment"]
         for _, comment in comments.iterrows():
-            comment_json = {"type": "comment"}
+            comment_json = {}
             comment_json.update(comment["meta"])
+            comment_json["correspondence"] = comment_json["comments"]
+            del comment_json["comments"]
             comment_json["resolved"] = comment["commentResolved"]
-            image_annotation.append(comment_json)
+            image_annotation["comments"].append(comment_json)
 
-        meta = {
-            "type": "meta",
+        tags = image_df[image_df["type"] == "tag"]
+        for _, tag in tags.iterrows():
+            image_annotation["tags"].append(tag["tag"])
+
+        image_annotation["metadata"] = {
             "width": int(image_width),
             "height": int(image_height),
             "status": image_status,
             "pinned": bool(image_pinned)
         }
-        image_annotation.append(meta)
         json.dump(
             image_annotation,
             open(output_dir / f"{image}___{project_suffix}", "w"),
@@ -281,10 +292,10 @@ def __get_user_metadata(annotation):
                 comment_meta = {
                     "x": annotation["x"],
                     "y": annotation["y"],
-                    "comments": annotation["comments"]
+                    "comments": annotation["correspondence"]
                 }
                 annotation_dict = {
-                    "type": annotation_type,
+                    "type": "comment",
                     "meta": comment_meta,
                     "commentResolved": comment_resolved,
                 }
@@ -294,11 +305,7 @@ def __get_user_metadata(annotation):
                 __append_annotation(annotation_dict)
         if include_tags:
             for annotation in annotation_json["tags"]:
-                annotation_tag = annotation["name"]
-                annotation_dict = {
-                    "type": annotation_type,
-                    "tag": annotation_tag
-                }
+                annotation_dict = {"type": "tag", "tag": annotation}
                 annotation_dict.update(image_metadata)
                 __append_annotation(annotation_dict)
         for annotation in annotation_json["instances"]:

tests/test_filter_instances.py

@@ -215,11 +215,17 @@ def test_df_to_annotations_full(tmpdir):
     tmpdir = Path(tmpdir)
 
     df = sa.aggregate_annotations_as_df(
-        PROJECT_DIR, include_classes_wo_annotations=True, include_comments=True
+        PROJECT_DIR,
+        include_classes_wo_annotations=True,
+        include_comments=True,
+        include_tags=True
     )
     sa.df_to_annotations(df, tmpdir)
     df_new = sa.aggregate_annotations_as_df(
-        tmpdir, include_classes_wo_annotations=True, include_comments=True
+        tmpdir,
+        include_classes_wo_annotations=True,
+        include_comments=True,
+        include_tags=True
     )
     for project in sa.search_projects("test df to annotations 4"):
         sa.delete_project(project)