Merge pull request #108 from scaleapi/add_classification_type_to_groundtruth

Category Annotation API Support

Author: cmpajot

nucleus/__init__.py

@@ -25,6 +25,7 @@
     Point,
     Point3D,
     PolygonAnnotation,
+    CategoryAnnotation,
     Segment,
     SegmentationAnnotation,
 )
@@ -581,6 +582,7 @@ def annotate_dataset(
                 BoxAnnotation,
                 PolygonAnnotation,
                 CuboidAnnotation,
+                CategoryAnnotation,
                 SegmentationAnnotation,
             ]
         ],

nucleus/annotation.py

@@ -8,6 +8,7 @@
     ANNOTATION_ID_KEY,
     ANNOTATIONS_KEY,
     BOX_TYPE,
+    CATEGORY_TYPE,
     CUBOID_TYPE,
     DIMENSIONS_KEY,
     GEOMETRY_KEY,
@@ -20,6 +21,7 @@
     POLYGON_TYPE,
     POSITION_KEY,
     REFERENCE_ID_KEY,
+    TAXONOMY_NAME_KEY,
     TYPE_KEY,
     VERTICES_KEY,
     WIDTH_KEY,
@@ -41,6 +43,8 @@ def from_json(cls, payload: dict):
             return PolygonAnnotation.from_json(payload)
         elif payload.get(TYPE_KEY, None) == CUBOID_TYPE:
             return CuboidAnnotation.from_json(payload)
+        elif payload.get(TYPE_KEY, None) == CATEGORY_TYPE:
+            return CategoryAnnotation.from_json(payload)
         else:
             return SegmentationAnnotation.from_json(payload)
 
@@ -120,6 +124,7 @@ class AnnotationTypes(Enum):
     BOX = BOX_TYPE
     POLYGON = POLYGON_TYPE
     CUBOID = CUBOID_TYPE
+    CATEGORY = CATEGORY_TYPE
 
 
 @dataclass  # pylint: disable=R0902
@@ -291,6 +296,36 @@ def to_payload(self) -> dict:
         return payload
 
 
+@dataclass
+class CategoryAnnotation(Annotation):
+    label: 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(
+            label=payload[LABEL_KEY],
+            taxonomy_name=payload[TAXONOMY_NAME_KEY],
+            reference_id=payload[REFERENCE_ID_KEY],
+            metadata=payload.get(METADATA_KEY, {}),
+        )
+
+    def to_payload(self) -> dict:
+        return {
+            LABEL_KEY: self.label,
+            TAXONOMY_NAME_KEY: self.taxonomy_name,
+            TYPE_KEY: CATEGORY_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

@@ -10,7 +10,13 @@
 CUBOID_TYPE = "cuboid"
 CATEGORY_TYPE = "category"
 MULTICATEGORY_TYPE = "multicategory"
-ANNOTATION_TYPES = (BOX_TYPE, POLYGON_TYPE, SEGMENTATION_TYPE, CUBOID_TYPE)
+ANNOTATION_TYPES = (
+    BOX_TYPE,
+    POLYGON_TYPE,
+    SEGMENTATION_TYPE,
+    CUBOID_TYPE,
+    CATEGORY_TYPE,
+)
 ANNOTATION_UPDATE_KEY = "update"
 AUTOTAGS_KEY = "autotags"
 AUTOTAG_SCORE_THRESHOLD = "score_threshold"

nucleus/dataset.py

@@ -217,7 +217,6 @@ def annotate(
                 route=f"dataset/{self.id}/annotate?async=1",
             )
             return AsyncJob.from_json(response, self._client)
-
         return self._client.annotate_dataset(
             self.id, annotations, update=update, batch_size=batch_size
         )

nucleus/payload_constructor.py

@@ -5,6 +5,7 @@
     BoxAnnotation,
     CuboidAnnotation,
     PolygonAnnotation,
+    CategoryAnnotation,
     SegmentationAnnotation,
 )
 from .prediction import (
@@ -60,6 +61,7 @@ def construct_annotation_payload(
             BoxAnnotation,
             PolygonAnnotation,
             CuboidAnnotation,
+            CategoryAnnotation,
             SegmentationAnnotation,
         ]
     ],

nucleus/utils.py

@@ -14,6 +14,7 @@
     BoxAnnotation,
     CuboidAnnotation,
     PolygonAnnotation,
+    CategoryAnnotation,
     SegmentationAnnotation,
 )
 
@@ -22,6 +23,7 @@
     ANNOTATIONS_KEY,
     BOX_TYPE,
     CUBOID_TYPE,
+    CATEGORY_TYPE,
     ITEM_KEY,
     POLYGON_TYPE,
     REFERENCE_ID_KEY,
@@ -116,6 +118,11 @@ def convert_export_payload(api_payload):
         for cuboid in row[CUBOID_TYPE]:
             cuboid[REFERENCE_ID_KEY] = row[ITEM_KEY][REFERENCE_ID_KEY]
             annotations[CUBOID_TYPE].append(CuboidAnnotation.from_json(cuboid))
+        for category in row[CATEGORY_TYPE]:
+            category[REFERENCE_ID_KEY] = row[ITEM_KEY][REFERENCE_ID_KEY]
+            annotations[CATEGORY_TYPE].append(
+                CategoryAnnotation.from_json(category)
+            )
         return_payload_row[ANNOTATIONS_KEY] = annotations
         return_payload.append(return_payload_row)
     return return_payload

tests/helpers.py

@@ -158,6 +158,14 @@ def reference_id_from_url(url):
     for i in range(len(TEST_POINTCLOUD_URLS))
 ]
 
+TEST_CATEGORY_ANNOTATIONS = [
+    {
+        "label": f"[Pytest] Category Label ${i}",
+        "reference_id": reference_id_from_url(TEST_IMG_URLS[i]),
+        "taxonomy_name": "[Pytest] Category 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"
 
@@ -266,6 +274,15 @@ def assert_cuboid_annotation_matches_dict(
     assert annotation_instance.yaw == annotation_dict["geometry"]["yaw"]
 
 
+def assert_category_annotation_matches_dict(
+    annotation_instance, annotation_dict
+):
+    assert annotation_instance.label == annotation_dict["label"]
+    assert (
+        annotation_instance.taxonomy_name == annotation_dict["taxonomy_name"]
+    )
+
+
 def assert_segmentation_annotation_matches_dict(
     annotation_instance, annotation_dict
 ):

tests/test_annotation.py

@@ -5,16 +5,19 @@
     TEST_IMG_URLS,
     TEST_BOX_ANNOTATIONS,
     TEST_POLYGON_ANNOTATIONS,
+    TEST_CATEGORY_ANNOTATIONS,
     TEST_SEGMENTATION_ANNOTATIONS,
     reference_id_from_url,
     assert_box_annotation_matches_dict,
     assert_polygon_annotation_matches_dict,
+    assert_category_annotation_matches_dict,
     assert_segmentation_annotation_matches_dict,
 )
 
 from nucleus import (
     BoxAnnotation,
     PolygonAnnotation,
+    CategoryAnnotation,
     SegmentationAnnotation,
     DatasetItem,
     Segment,
@@ -55,6 +58,12 @@ def dataset(CLIENT):
 
     response = ds.append(ds_items)
     assert ERROR_PAYLOAD not in response.json()
+
+    response = ds.add_taxonomy(
+        "[Pytest] Category Taxonomy 1",
+        "category",
+        [f"[Pytest] Category Label ${i}" for i in range((len(TEST_IMG_URLS)))],
+    )
     yield ds
 
     response = CLIENT.delete_dataset(ds.id)
@@ -100,6 +109,24 @@ def test_polygon_gt_upload(dataset):
     )
 
 
+def test_category_gt_upload(dataset):
+    annotation = CategoryAnnotation.from_json(TEST_CATEGORY_ANNOTATIONS[0])
+    response = dataset.annotate(annotations=[annotation])
+
+    assert response["dataset_id"] == dataset.id
+    assert response["annotations_processed"] == 1
+    assert response["annotations_ignored"] == 0
+
+    response = dataset.refloc(annotation.reference_id)["annotations"][
+        "category"
+    ]
+    assert len(response) == 1
+    response_annotation = response[0]
+    assert_category_annotation_matches_dict(
+        response_annotation, TEST_CATEGORY_ANNOTATIONS[0]
+    )
+
+
 def test_single_semseg_gt_upload(dataset):
     annotation = SegmentationAnnotation.from_json(
         TEST_SEGMENTATION_ANNOTATIONS[0]
@@ -308,6 +335,63 @@ def test_polygon_gt_upload_ignore(dataset):
         response_annotation, TEST_POLYGON_ANNOTATIONS[0]
     )
 
+
+def test_category_gt_upload_update(dataset):
+    annotation = CategoryAnnotation.from_json(TEST_CATEGORY_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_CATEGORY_ANNOTATIONS[1])
+    annotation_update_params["reference_id"] = TEST_CATEGORY_ANNOTATIONS[0][
+        "reference_id"
+    ]
+
+    annotation_update = CategoryAnnotation.from_json(annotation_update_params)
+    response = dataset.annotate(annotations=[annotation_update], update=True)
+
+    assert response["annotations_processed"] == 1
+    assert response["annotations_ignored"] == 0
+
+    response = dataset.refloc(annotation.reference_id)["annotations"][
+        "category"
+    ]
+    assert len(response) == 1
+    response_annotation = response[0]
+    assert_category_annotation_matches_dict(
+        response_annotation, annotation_update_params
+    )
+
+
+def test_category_gt_upload_ignore(dataset):
+    annotation = CategoryAnnotation.from_json(TEST_CATEGORY_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_CATEGORY_ANNOTATIONS[1])
+    annotation_update_params["reference_id"] = TEST_CATEGORY_ANNOTATIONS[0][
+        "reference_id"
+    ]
+
+    annotation_update = CategoryAnnotation.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"][
+        "category"
+    ]
+    assert len(response) == 1
+    response_annotation = response[0]
+    assert_category_annotation_matches_dict(
+        response_annotation, TEST_CATEGORY_ANNOTATIONS[0]
+    )
+
     @pytest.mark.integration
     def test_box_gt_deletion(dataset):
         annotation = BoxAnnotation(**TEST_BOX_ANNOTATIONS[0])
@@ -323,3 +407,19 @@ def test_box_gt_deletion(dataset):
         job_status = job.status()
         assert job_status["status"] == "Completed"
         assert job_status["job_id"] == 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.id