Merge pull request #37 from superannotateai/SAS-3383

Fix multi template coco conversion SAS-3383 - SAS-3401

Author: dshabin

superannotate/db/annotation_classes.py

@@ -12,6 +12,7 @@
 )
 from .project_api import get_project_metadata_bare
 from ..mixp.decorators import Trackable
+from .utils import get_templates_mapping
 
 logger = logging.getLogger("superannotate-python-sdk")
 
@@ -344,7 +345,11 @@ def fill_class_and_attribute_ids(annotation_json, annotation_classes_dict):
                     'attribute_groups': {}
                 }
     annotation_classes_dict = {**annotation_classes_dict, **unknown_classes}
-
+    templates_map = get_templates_mapping()
+    for ann in (
+        i for i in annotation_json["instances"] if i['type'] == 'template'
+    ):
+        ann['templateId'] = templates_map.get(ann['templateName'], -1)
     for ann in annotation_json["instances"]:
         if "className" not in ann:
             logger.warning("No className in annotation instance")

superannotate/db/utils.py

@@ -747,3 +747,19 @@ def __attach_image_urls_to_project_thread(
         logger.warning(e)
     else:
         uploaded[thread_id] += uploaded_imgs
+
+
+def get_templates_mapping():
+    response = _api.send_request(
+        req_type='GET', path=f'/templates', params={"team_id": _api.team_id}
+    )
+    if not response.ok:
+        raise SABaseException(
+            response.status_code, "Couldn't get templates " + response.text
+        )
+    res = response.json()
+    templates = res['data']
+    templates_map = {}
+    for template in templates:
+        templates_map[template['name']] = template['id']
+    return templates_map

superannotate/input_converters/converters/coco_converters/coco_to_sa_vector.py

@@ -171,7 +171,8 @@ def coco_keypoint_detection_to_sa_vector(coco_path, output_dir):
         cat_id_to_cat[cat["id"]] = {
             "name": cat["name"],
             "keypoints": cat["keypoints"],
-            "skeleton": cat["skeleton"]
+            "skeleton": cat["skeleton"],
+            "supercategory": cat["supercategory"]
         }
 
     image_id_to_annotations = {}
@@ -218,11 +219,8 @@ def coco_keypoint_detection_to_sa_vector(coco_path, output_dir):
                 for connection in cat_id_to_cat[annot["category_id"]
                                                ]['skeleton']:
 
-                    index = cat['skeleton'].index(connection)
-                    from_point = cat_id_to_cat[annot["category_id"]
-                                              ]['skeleton'][index][0]
-                    to_point = cat_id_to_cat[annot["category_id"]
-                                            ]['skeleton'][index][1]
+                    from_point = connection[0]
+                    to_point = connection[1]
 
                     if from_point in bad_points or to_point in bad_points:
                         continue
@@ -242,8 +240,12 @@ def coco_keypoint_detection_to_sa_vector(coco_path, output_dir):
                     pointLabels[id_mapping[kp_index + 1] - 1] = kp_name
 
                 sa_obj = _create_vector_instance(
-                    'template', points, pointLabels, [],
-                    cat_id_to_cat[annot["category_id"]]["name"], connections
+                    'template',
+                    points,
+                    pointLabels, [],
+                    cat_id_to_cat[annot["category_id"]]["supercategory"],
+                    connections,
+                    template_name=cat_id_to_cat[annot["category_id"]]["name"]
                 )
                 if str(annot['image_id']) not in image_id_to_annotations:
                     image_id_to_annotations[str(annot['image_id'])] = [sa_obj]

superannotate/input_converters/converters/sa_json_helper.py

@@ -8,7 +8,8 @@ def _create_vector_instance(
     pointLabels,
     attributes,
     class_name='',
-    connections=[]
+    connections=[],
+    template_name=''
 ):
     sa_instance = {
         'type': instance_type,
@@ -20,6 +21,8 @@ def _create_vector_instance(
     if instance_type == 'template':
         sa_instance['points'] = points
         sa_instance['connections'] = connections
+        sa_instance['className'] = class_name
+        sa_instance['templateName'] = template_name
     elif instance_type == 'point':
         sa_instance['x'] = points[0]
         sa_instance['y'] = points[1]

tests/converter_test/COCO/input/toSuperAnnotate/keypoint_detection_multi_template/keypoint_multi_template_test.json

@@ -0,0 +1,172 @@
+{
+  "info": {
+    "description": "This is dataset.",
+    "url": "https://superannotate.ai",
+    "version": "1.0",
+    "year": 2021,
+    "contributor": "Superannotate AI",
+    "date_created": "01/06/2021"
+  },
+  "licenses": [
+    {
+      "url": "https://superannotate.ai",
+      "id": 1,
+      "name": "Superannotate AI"
+    }
+  ],
+  "images": [
+    {
+      "id": 1,
+      "file_name": "68307_47130_68308_47130_68307_47131_68308_47131_0.png",
+      "height": 1000,
+      "width": 1000,
+      "license": 1
+    }
+  ],
+  "annotations": [
+    {
+      "id": 1,
+      "image_id": 1,
+      "iscrowd": 0,
+      "bbox": [
+        271,
+        114,
+        182,
+        248
+      ],
+      "area": 45136,
+      "num_keypoints": 5,
+      "keypoints": [
+        271.99,
+        114.9,
+        2,
+        419.26,
+        137.22,
+        2,
+        279.01,
+        351.5,
+        2,
+        453.32,
+        362.66,
+        2,
+        349.13,
+        233.2,
+        2
+      ],
+      "category_id": 1
+    },
+    {
+      "id": 2,
+      "image_id": 1,
+      "iscrowd": 0,
+      "bbox": [
+        668,
+        208,
+        131,
+        208
+      ],
+      "area": 27248,
+      "num_keypoints": 4,
+      "keypoints": [
+        668.76,
+        219.62,
+        2,
+        763.61,
+        208.26,
+        2,
+        799.82,
+        377.39,
+        2,
+        713.64,
+        416.52,
+        2
+      ],
+      "category_id": 2
+    },
+    {
+      "id": 3,
+      "image_id": 1,
+      "iscrowd": 0,
+      "bbox": [
+        455,
+        522,
+        273,
+        194
+      ],
+      "area": 52962,
+      "num_keypoints": 4,
+      "keypoints": [
+        455.12,
+        533.02,
+        2,
+        652.62,
+        522.44,
+        2,
+        728.01,
+        679.91,
+        2,
+        548.56,
+        716.34,
+        2
+      ],
+      "category_id": 2
+    }
+  ],
+  "categories": [
+    {
+      "name": "custom",
+      "supercategory": "class1",
+      "skeleton": [
+        [
+          1,
+          5
+        ],
+        [
+          2,
+          5
+        ],
+        [
+          5,
+          3
+        ],
+        [
+          5,
+          4
+        ]
+      ],
+      "keypoints": [
+        "ca",
+        "cb",
+        "cd",
+        "cc",
+        "ce"
+      ],
+      "id": 1
+    },
+    {
+      "name": "template2",
+      "supercategory": "class1",
+      "skeleton": [
+        [
+          1,
+          4
+        ],
+        [
+          2,
+          4
+        ],
+        [
+          3,
+          4
+        ]
+      ],
+      "keypoints": [
+        "ta",
+        "tb",
+        "tc",
+        "td"
+      ],
+      "id": 2
+    }
+  ]
+}

tests/converter_test/COCO/input/toSuperAnnotate/keypoint_detection_multi_template/truth.json

@@ -0,0 +1 @@
+{"instances": [{"type": "template", "pointLabels": {"0": "ca", "1": "cb", "2": "cd", "3": "cc", "4": "ce"}, "attributes": [], "creationType": "Pre-annotation", "points": [{"id": 1, "x": 271.99, "y": 114.9}, {"id": 2, "x": 419.26, "y": 137.22}, {"id": 3, "x": 279.01, "y": 351.5}, {"id": 4, "x": 453.32, "y": 362.66}, {"id": 5, "x": 349.13, "y": 233.2}], "connections": [{"id": 7, "from": 1, "to": 5}, {"id": 7, "from": 2, "to": 5}, {"id": 7, "from": 5, "to": 3}, {"id": 7, "from": 5, "to": 4}], "className": "class1", "templateName": "custom"}, {"type": "template", "pointLabels": {"0": "ta", "1": "tb", "2": "tc", "3": "td"}, "attributes": [], "creationType": "Pre-annotation", "points": [{"id": 1, "x": 668.76, "y": 219.62}, {"id": 2, "x": 763.61, "y": 208.26}, {"id": 3, "x": 799.82, "y": 377.39}, {"id": 4, "x": 713.64, "y": 416.52}], "connections": [{"id": 6, "from": 1, "to": 4}, {"id": 6, "from": 2, "to": 4}, {"id": 6, "from": 3, "to": 4}], "className": "class1", "templateName": "template2"}, {"type": "template", "pointLabels": {"0": "ta", "1": "tb", "2": "tc", "3": "td"}, "attributes": [], "creationType": "Pre-annotation", "points": [{"id": 1, "x": 455.12, "y": 533.02}, {"id": 2, "x": 652.62, "y": 522.44}, {"id": 3, "x": 728.01, "y": 679.91}, {"id": 4, "x": 548.56, "y": 716.34}], "connections": [{"id": 6, "from": 1, "to": 4}, {"id": 6, "from": 2, "to": 4}, {"id": 6, "from": 3, "to": 4}], "className": "class1", "templateName": "template2"}], "metadata": {"name": "68307_47130_68308_47130_68307_47131_68308_47131_0.png", "width": 1000, "height": 1000}, "tags": [], "comments": []}

tests/converter_test/test_conversion.py

@@ -29,6 +29,32 @@ def test_keypoint_detection_coco2sa(tmpdir):
     )
 
 
+def test_keypoint_detection_coco2sa_multi_template(tmpdir):
+    input_dir = Path(
+        "tests"
+    ) / "converter_test" / "COCO" / "input" / "toSuperAnnotate" / "keypoint_detection_multi_template"
+    out_path = Path(
+        tmpdir
+    ) / "toSuperAnnotate" / "keypoint_detection_multi_template"
+
+    sa.import_annotation(
+        input_dir, out_path, "COCO", "keypoint_multi_template_test", "Vector",
+        "keypoint_detection"
+    )
+    import json
+    with open(str(Path(input_dir) / "truth.json")) as f:
+        truth = json.loads(f.read())
+
+    with open(
+        str(
+            Path(out_path) /
+            "68307_47130_68308_47130_68307_47131_68308_47131_0.png___objects.json"
+        )
+    ) as f:
+        data = json.loads(f.read())
+    assert data == truth
+
+
 # test instance segmentation
 def test_instance_segmentation_coco2sa(tmpdir):
     input_dir = Path(
@@ -121,3 +147,26 @@ def test_instance_segmentation_sa2coco_vector_empty_name(tmpdir):
         input_dir, out_path, "COCO", "instance_test_vector", "Vector",
         "instance_segmentation"
     )
+
+
+def test_upload_annotations_with_template_id(tmpdir):
+    tmpdir = Path(tmpdir)
+    project_name = "test_templates"
+    for project in sa.search_projects(project_name):
+        sa.delete_project(project)
+    project = sa.create_project(project_name, "test", "Vector")
+    sa.upload_images_from_folder_to_project(
+        project, "./tests/sample_coco_with_templates"
+    )
+    input_dir = Path("tests") / "sample_coco_with_templates"
+    out_path = Path(
+        tmpdir
+    ) / "toSuperAnnotate" / "keypoint_detection_multi_template"
+
+    sa.import_annotation(
+        input_dir, out_path, "COCO", "sample_coco", "Vector",
+        "keypoint_detection"
+    )
+    sa.upload_annotations_from_folder_to_project(project, out_path)
+    image_metadata = sa.get_image_annotations(project_name, "t.png")
+    assert image_metadata['annotation_json']['instances'][0]['templateId'] == -1