Initial classification support

Author: Claire Pajot

nucleus/__init__.py

@@ -73,6 +73,7 @@
     Point,
     Point3D,
     PolygonAnnotation,
+    CategoryAnnotation,
     Segment,
     SegmentationAnnotation,
 )
@@ -622,6 +623,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,
     DATASET_ITEM_ID_KEY,
     DIMENSIONS_KEY,
@@ -135,6 +136,7 @@ class AnnotationTypes(Enum):
     BOX = BOX_TYPE
     POLYGON = POLYGON_TYPE
     CUBOID = CUBOID_TYPE
+    CATEGORY = CATEGORY_TYPE
 
 
 @dataclass  # pylint: disable=R0902
@@ -314,6 +316,40 @@ def to_payload(self) -> dict:
         return payload
 
 
+@dataclass
+class CategoryAnnotation(Annotation):
+    label: str
+    reference_id: Optional[str] = None
+    item_id: Optional[str] = None
+    annotation_id: Optional[str] = None
+    metadata: Optional[Dict] = None
+
+    def __post_init__(self):
+        self._check_ids()
+        self.metadata = self.metadata if self.metadata else {}
+
+    @classmethod
+    def from_json(cls, payload: dict):
+        # TODO: Remove? geometry = payload.get(GEOMETRY_KEY, {})
+        return cls(
+            label=payload.get(LABEL_KEY, 0),
+            reference_id=payload.get(REFERENCE_ID_KEY, None),
+            item_id=payload.get(DATASET_ITEM_ID_KEY, None),
+            annotation_id=payload.get(ANNOTATION_ID_KEY, None),
+            metadata=payload.get(METADATA_KEY, {}),
+        )
+
+    def to_payload(self) -> dict:
+        return {
+            LABEL_KEY: self.label,
+            TYPE_KEY: CATEGORY_TYPE,
+            GEOMETRY_KEY: {},
+            REFERENCE_ID_KEY: self.reference_id,
+            ANNOTATION_ID_KEY: self.annotation_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

@@ -8,6 +8,7 @@
 MASK_TYPE = "mask"
 SEGMENTATION_TYPE = "segmentation"
 CUBOID_TYPE = "cuboid"
+CATEGORY_TYPE = "category"
 ANNOTATION_TYPES = (BOX_TYPE, POLYGON_TYPE, SEGMENTATION_TYPE, CUBOID_TYPE)
 ANNOTATION_UPDATE_KEY = "update"
 AUTOTAGS_KEY = "autotags"

nucleus/payload_constructor.py

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

tests/helpers.py

@@ -155,6 +155,14 @@ def reference_id_from_url(url):
     for i in range(len(TEST_POINTCLOUD_URLS))
 ]
 
+TEST_CATEGORY_ANNOTATIONS = [
+    {
+        "label": f"[Pytest] Category Annotation ${i}",
+        "reference_id": reference_id_from_url(TEST_IMG_URLS[i]),
+        "annotation_id": f"[Pytest] Category Annotation Annotation Id{i}",
+    }
+    for i in range(len(TEST_IMG_URLS))
+]
 
 TEST_MASK_URL = "https://raw.githubusercontent.com/scaleapi/nucleus-python-client/master/tests/testdata/000000000285.png"
 
@@ -263,6 +271,16 @@ 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.annotation_id == annotation_dict["annotation_id"]
+    )
+
+
 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,
@@ -95,6 +98,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]