Initial multicategory support

Author: Claire Pajot

nucleus/__init__.py

@@ -26,6 +26,7 @@
     Point3D,
     PolygonAnnotation,
     CategoryAnnotation,
+    MultiCategoryAnnotation,
     Segment,
     SegmentationAnnotation,
 )
@@ -583,6 +584,7 @@ def annotate_dataset(
                 PolygonAnnotation,
                 CuboidAnnotation,
                 CategoryAnnotation,
+                MultiCategoryAnnotation,
                 SegmentationAnnotation,
             ]
         ],

nucleus/annotation.py

@@ -9,12 +9,14 @@
     ANNOTATIONS_KEY,
     BOX_TYPE,
     CATEGORY_TYPE,
+    MULTICATEGORY_TYPE,
     CUBOID_TYPE,
     DIMENSIONS_KEY,
     GEOMETRY_KEY,
     HEIGHT_KEY,
     INDEX_KEY,
     LABEL_KEY,
+    LABELS_KEY,
     MASK_TYPE,
     MASK_URL_KEY,
     METADATA_KEY,
@@ -45,6 +47,8 @@ def from_json(cls, payload: dict):
             return CuboidAnnotation.from_json(payload)
         elif payload.get(TYPE_KEY, None) == CATEGORY_TYPE:
             return CategoryAnnotation.from_json(payload)
+        elif payload.get(TYPE_KEY, None) == MULTICATEGORY_TYPE:
+            return MultiCategoryAnnotation.from_json(payload)
         else:
             return SegmentationAnnotation.from_json(payload)
 
@@ -125,6 +129,7 @@ class AnnotationTypes(Enum):
     POLYGON = POLYGON_TYPE
     CUBOID = CUBOID_TYPE
     CATEGORY = CATEGORY_TYPE
+    MULTICATEGORY = MULTICATEGORY_TYPE
 
 
 @dataclass  # pylint: disable=R0902
@@ -326,6 +331,36 @@ def to_payload(self) -> dict:
         }
 
 
+@dataclass
+class MultiCategoryAnnotation(Annotation):
+    labels: List[str]
+    taxonomy_name: str
+    reference_id: str
+    metadata: Optional[Dict] = None
+
+    def __post_init__(self):
+        self.metadata = self.metadata if self.metadata else {}
+
+    @classmethod
+    def from_json(cls, payload: dict):
+        return cls(
+            labels=payload[LABELS_KEY],
+            taxonomy_name=payload[TAXONOMY_NAME_KEY],
+            reference_id=payload[REFERENCE_ID_KEY],
+            metadata=payload.get(METADATA_KEY, {}),
+        )
+
+    def to_payload(self) -> dict:
+        return {
+            LABELS_KEY: self.labels,
+            TAXONOMY_NAME_KEY: self.taxonomy_name,
+            TYPE_KEY: MULTICATEGORY_TYPE,
+            GEOMETRY_KEY: {},
+            REFERENCE_ID_KEY: self.reference_id,
+            METADATA_KEY: self.metadata,
+        }
+
+
 def is_local_path(path: str) -> bool:
     return urlparse(path).scheme not in {"https", "http", "s3", "gs"}
 

nucleus/constants.py

@@ -16,6 +16,7 @@
     SEGMENTATION_TYPE,
     CUBOID_TYPE,
     CATEGORY_TYPE,
+    MULTICATEGORY_TYPE,
 )
 ANNOTATION_UPDATE_KEY = "update"
 AUTOTAGS_KEY = "autotags"

nucleus/payload_constructor.py

@@ -6,6 +6,7 @@
     CuboidAnnotation,
     PolygonAnnotation,
     CategoryAnnotation,
+    MultiCategoryAnnotation,
     SegmentationAnnotation,
 )
 from .prediction import (
@@ -62,6 +63,7 @@ def construct_annotation_payload(
             PolygonAnnotation,
             CuboidAnnotation,
             CategoryAnnotation,
+            MultiCategoryAnnotation,
             SegmentationAnnotation,
         ]
     ],

nucleus/utils.py

@@ -15,6 +15,7 @@
     CuboidAnnotation,
     PolygonAnnotation,
     CategoryAnnotation,
+    MultiCategoryAnnotation,
     SegmentationAnnotation,
 )
 
@@ -24,6 +25,7 @@
     BOX_TYPE,
     CUBOID_TYPE,
     CATEGORY_TYPE,
+    MULTICATEGORY_TYPE,
     ITEM_KEY,
     POLYGON_TYPE,
     REFERENCE_ID_KEY,
@@ -123,6 +125,11 @@ def convert_export_payload(api_payload):
             annotations[CATEGORY_TYPE].append(
                 CategoryAnnotation.from_json(category)
             )
+        for multicategory in row[MULTICATEGORY_TYPE]:
+            multicategory[REFERENCE_ID_KEY] = row[ITEM_KEY][REFERENCE_ID_KEY]
+            annotations[MULTICATEGORY_TYPE].append(
+                MultiCategoryAnnotation.from_json(multicategory)
+            )
         return_payload_row[ANNOTATIONS_KEY] = annotations
         return_payload.append(return_payload_row)
     return return_payload

tests/helpers.py

@@ -167,6 +167,18 @@ def reference_id_from_url(url):
     for i in range(len(TEST_IMG_URLS))
 ]
 
+TEST_MULTICATEGORY_ANNOTATIONS = [
+    {
+        "labels": [
+            f"[Pytest] MultiCategory Label ${i}",
+            f"[Pytest] MultiCategory Label ${i+1}",
+        ],
+        "reference_id": reference_id_from_url(TEST_IMG_URLS[i]),
+        "taxonomy_name": "[Pytest] MultiCategory Taxonomy 1",
+    }
+    for i in range(len(TEST_IMG_URLS))
+]
+
 TEST_MASK_URL = "https://raw.githubusercontent.com/scaleapi/nucleus-python-client/master/tests/testdata/000000000285.png"
 
 TEST_SEGMENTATION_ANNOTATIONS = [
@@ -283,6 +295,15 @@ def assert_category_annotation_matches_dict(
     )
 
 
+def assert_multicategory_annotation_matches_dict(
+    annotation_instance, annotation_dict
+):
+    assert set(annotation_instance.labels) == set(annotation_dict["labels"])
+    assert (
+        annotation_instance.taxonomy_name == annotation_dict["taxonomy_name"]
+    )
+
+
 def assert_segmentation_annotation_matches_dict(
     annotation_instance, annotation_dict
 ):

tests/test_annotation.py

@@ -6,18 +6,21 @@
     TEST_BOX_ANNOTATIONS,
     TEST_POLYGON_ANNOTATIONS,
     TEST_CATEGORY_ANNOTATIONS,
+    TEST_MULTICATEGORY_ANNOTATIONS,
     TEST_SEGMENTATION_ANNOTATIONS,
     reference_id_from_url,
     assert_box_annotation_matches_dict,
     assert_polygon_annotation_matches_dict,
     assert_category_annotation_matches_dict,
+    assert_multicategory_annotation_matches_dict,
     assert_segmentation_annotation_matches_dict,
 )
 
 from nucleus import (
     BoxAnnotation,
     PolygonAnnotation,
     CategoryAnnotation,
+    MultiCategoryAnnotation,
     SegmentationAnnotation,
     DatasetItem,
     Segment,
@@ -64,6 +67,14 @@ def dataset(CLIENT):
         "category",
         [f"[Pytest] Category Label ${i}" for i in range((len(TEST_IMG_URLS)))],
     )
+    response = ds.add_taxonomy(
+        "[Pytest] MultiCategory Taxonomy 1",
+        "multicategory",
+        [
+            f"[Pytest] MultiCategory Label ${i}"
+            for i in range((len(TEST_IMG_URLS) + 1))
+        ],
+    )
     yield ds
 
     response = CLIENT.delete_dataset(ds.id)
@@ -73,6 +84,7 @@ def dataset(CLIENT):
 def test_box_gt_upload(dataset):
     annotation = BoxAnnotation(**TEST_BOX_ANNOTATIONS[0])
     response = dataset.annotate(annotations=[annotation])
+    print(response)
 
     assert response["dataset_id"] == dataset.id
     assert response["annotations_processed"] == 1
@@ -127,6 +139,28 @@ def test_category_gt_upload(dataset):
     )
 
 
+def test_multicategory_gt_upload(dataset):
+    annotation = MultiCategoryAnnotation.from_json(
+        TEST_MULTICATEGORY_ANNOTATIONS[0]
+    )
+    response = dataset.annotate(annotations=[annotation])
+
+    assert response["dataset_id"] == dataset.id
+    assert response["annotations_processed"] == 1
+    assert response["annotations_ignored"] == 0
+    print("HERE")
+    response = dataset.refloc(annotation.reference_id)["annotations"][
+        "multicategory"
+    ]
+    print("RESPONSE: ", response)
+    # TODO: Weirdness here also in the response from refloc?
+    assert len(response) == 1
+    response_annotation = response[0]
+    assert_multicategory_annotation_matches_dict(
+        response_annotation, TEST_MULTICATEGORY_ANNOTATIONS[0]
+    )
+
+
 def test_single_semseg_gt_upload(dataset):
     annotation = SegmentationAnnotation.from_json(
         TEST_SEGMENTATION_ANNOTATIONS[0]
@@ -206,6 +240,7 @@ def test_mixed_annotation_upload(dataset):
     response_annotations = dataset.refloc(bbox_annotations[0].reference_id)[
         "annotations"
     ]
+    print(response_annotations)
     assert len(response_annotations) == 2
     assert len(response_annotations["box"]) == 1
     assert "segmentation" in response_annotations
@@ -392,34 +427,120 @@ def test_category_gt_upload_ignore(dataset):
         response_annotation, TEST_CATEGORY_ANNOTATIONS[0]
     )
 
-    @pytest.mark.integration
-    def test_box_gt_deletion(dataset):
-        annotation = BoxAnnotation(**TEST_BOX_ANNOTATIONS[0])
 
-        print(annotation)
+def test_multicategory_gt_upload_update(dataset):
+    annotation = MultiCategoryAnnotation.from_json(
+        TEST_MULTICATEGORY_ANNOTATIONS[0]
+    )
+    response = dataset.annotate(annotations=[annotation])
 
-        response = dataset.annotate(annotations=[annotation])
+    assert response["annotations_processed"] == 1
 
-        assert response["annotations_processed"] == 1
+    # Copy so we don't modify the original.
+    annotation_update_params = dict(TEST_MULTICATEGORY_ANNOTATIONS[1])
+    annotation_update_params["reference_id"] = TEST_MULTICATEGORY_ANNOTATIONS[
+        0
+    ]["reference_id"]
 
-        job = dataset.delete_annotations()
-        job.sleep_until_complete()
-        job_status = job.status()
-        assert job_status["status"] == "Completed"
-        assert job_status["job_id"] == job.id
+    annotation_update = MultiCategoryAnnotation.from_json(
+        annotation_update_params
+    )
+    response = dataset.annotate(annotations=[annotation_update], update=True)
 
-    @pytest.mark.integration
-    def test_category_gt_deletion(dataset):
-        annotation = CategoryAnnotation.from_json(TEST_CATEGORY_ANNOTATIONS[0])
+    assert response["annotations_processed"] == 1
+    assert response["annotations_ignored"] == 0
 
-        print(annotation)
+    response = dataset.refloc(annotation.reference_id)["annotations"][
+        "multicategory"
+    ]
+    assert len(response) == 1
+    response_annotation = response[0]
+    assert_multicategory_annotation_matches_dict(
+        response_annotation, annotation_update_params
+    )
 
-        response = dataset.annotate(annotations=[annotation])
 
-        assert response["annotations_processed"] == 1
+def test_multicategory_gt_upload_ignore(dataset):
+    annotation = MultiCategoryAnnotation.from_json(
+        TEST_MULTICATEGORY_ANNOTATIONS[0]
+    )
+    response = dataset.annotate(annotations=[annotation])
+
+    assert response["annotations_processed"] == 1
+
+    # Copy so we don't modify the original.
+    annotation_update_params = dict(TEST_MULTICATEGORY_ANNOTATIONS[1])
+    annotation_update_params["reference_id"] = TEST_MULTICATEGORY_ANNOTATIONS[
+        0
+    ]["reference_id"]
+
+    annotation_update = MultiCategoryAnnotation.from_json(
+        annotation_update_params
+    )
+    # Default behavior is ignore.
+    response = dataset.annotate(annotations=[annotation_update])
+
+    assert response["annotations_processed"] == 0
+    assert response["annotations_ignored"] == 1
+
+    response = dataset.refloc(annotation.reference_id)["annotations"][
+        "multicategory"
+    ]
+    assert len(response) == 1
+    response_annotation = response[0]
+    assert_multicategory_annotation_matches_dict(
+        response_annotation, TEST_MULTICATEGORY_ANNOTATIONS[0]
+    )
+
+
+@pytest.mark.integration
+def test_box_gt_deletion(dataset):
+    annotation = BoxAnnotation(**TEST_BOX_ANNOTATIONS[0])
+
+    print(annotation)
+
+    response = dataset.annotate(annotations=[annotation])
+
+    assert response["annotations_processed"] == 1
+
+    job = dataset.delete_annotations()
+    job.sleep_until_complete()
+    job_status = job.status()
+    assert job_status["status"] == "Completed"
+    assert job_status["job_id"] == job.job_id
+
+
+@pytest.mark.integration
+def test_category_gt_deletion(dataset):
+    annotation = CategoryAnnotation.from_json(TEST_CATEGORY_ANNOTATIONS[0])
+
+    print(annotation)
+
+    response = dataset.annotate(annotations=[annotation])
+
+    assert response["annotations_processed"] == 1
+
+    job = dataset.delete_annotations()
+    job.sleep_until_complete()
+    job_status = job.status()
+    assert job_status["status"] == "Completed"
+    assert job_status["job_id"] == job.job_id
+
+
+@pytest.mark.integration
+def test_multicategory_gt_deletion(dataset):
+    annotation = MultiCategoryAnnotation.from_json(
+        TEST_MULTICATEGORY_ANNOTATIONS[0]
+    )
+
+    print(annotation)
+
+    response = dataset.annotate(annotations=[annotation])
+
+    assert response["annotations_processed"] == 1
 
-        job = dataset.delete_annotations()
-        job.sleep_until_complete()
-        job_status = job.status()
-        assert job_status["status"] == "Completed"
-        assert job_status["job_id"] == job.id
+    job = dataset.delete_annotations()
+    job.sleep_until_complete()
+    job_status = job.status()
+    assert job_status["status"] == "Completed"
+    assert job_status["job_id"] == job.job_id