Merge pull request #39 from scaleapi/sasha/refloc_response_change

adjust refloc payload format

Author: sasha-scale

nucleus/annotation.py

@@ -22,6 +22,19 @@
 )
 
 
+class Annotation:
+    @classmethod
+    def from_json(cls, payload: dict):
+        if payload.get(TYPE_KEY, None) == BOX_TYPE:
+            geometry = payload.get(GEOMETRY_KEY, {})
+            return BoxAnnotation.from_json(payload)
+        elif payload.get(TYPE_KEY, None) == POLYGON_TYPE:
+            geometry = payload.get(GEOMETRY_KEY, {})
+            return PolygonAnnotation.from_json(payload)
+        else:
+            return SegmentationAnnotation.from_json(payload)
+
+
 class Segment:
     def __init__(
         self, label: str, index: int, metadata: Optional[dict] = None
@@ -51,14 +64,17 @@ def to_payload(self) -> dict:
         return payload
 
 
-class SegmentationAnnotation:
+class SegmentationAnnotation(Annotation):
     def __init__(
         self,
         mask_url: str,
         annotations: List[Segment],
         reference_id: Optional[str] = None,
         item_id: Optional[str] = None,
     ):
+        super().__init__()
+        if not mask_url:
+            raise Exception("You must specify a mask_url.")
         if bool(reference_id) == bool(item_id):
             raise Exception(
                 "You must specify either a reference_id or an item_id for an annotation."
@@ -74,7 +90,7 @@ def __str__(self):
     @classmethod
     def from_json(cls, payload: dict):
         return cls(
-            mask_url=payload[MASK_URL_KEY],
+            mask_url=payload.get(MASK_URL_KEY),
             annotations=[
                 Segment.from_json(ann)
                 for ann in payload.get(ANNOTATIONS_KEY, [])
@@ -101,7 +117,7 @@ class AnnotationTypes(Enum):
 
 
 # TODO: Add base annotation class to reduce repeated code here
-class BoxAnnotation:
+class BoxAnnotation(Annotation):
     # pylint: disable=too-many-instance-attributes
     def __init__(
         self,
@@ -115,6 +131,7 @@ def __init__(
         annotation_id: Optional[str] = None,
         metadata: Optional[Dict] = None,
     ):
+        super().__init__()
         if bool(reference_id) == bool(item_id):
             raise Exception(
                 "You must specify either a reference_id or an item_id for an annotation."
@@ -164,7 +181,7 @@ def __str__(self):
 
 
 # TODO: Add Generic type for 2D point
-class PolygonAnnotation:
+class PolygonAnnotation(Annotation):
     def __init__(
         self,
         label: str,
@@ -174,6 +191,7 @@ def __init__(
         annotation_id: Optional[str] = None,
         metadata: Optional[Dict] = None,
     ):
+        super().__init__()
         if bool(reference_id) == bool(item_id):
             raise Exception(
                 "You must specify either a reference_id or an item_id for an annotation."

nucleus/constants.py

@@ -47,6 +47,8 @@
 VERTICES_KEY = "vertices"
 BOX_TYPE = "box"
 POLYGON_TYPE = "polygon"
+SEGMENTATION_TYPE = "segmentation"
+ANNOTATION_TYPES = (BOX_TYPE, POLYGON_TYPE, SEGMENTATION_TYPE)
 GEOMETRY_KEY = "geometry"
 AUTOTAGS_KEY = "autotags"
 MASK_URL_KEY = "mask_url"

nucleus/dataset.py

@@ -1,6 +1,11 @@
 from typing import List, Dict, Any, Optional, Union
 from .dataset_item import DatasetItem
-from .annotation import BoxAnnotation, PolygonAnnotation
+from .annotation import (
+    Annotation,
+    BoxAnnotation,
+    PolygonAnnotation,
+    SegmentationAnnotation,
+)
 from .constants import (
     DATASET_NAME_KEY,
     DATASET_MODEL_RUNS_KEY,
@@ -12,9 +17,14 @@
     ITEM_KEY,
     DEFAULT_ANNOTATION_UPDATE_MODE,
     ANNOTATIONS_KEY,
+    BOX_TYPE,
+    POLYGON_TYPE,
+    SEGMENTATION_TYPE,
+    ANNOTATION_TYPES,
 )
 from .payload_constructor import construct_model_run_creation_payload
 
+
 class Dataset:
     """
     Nucleus Dataset. You can append images with metadata to your dataset,
@@ -227,17 +237,19 @@ def list_autotags(self):
 
     def _format_dataset_item_response(self, response: dict) -> dict:
         item = response.get(ITEM_KEY, None)
-        annotation_payload = response.get(ANNOTATIONS_KEY, [])
+        annotation_payload = response.get(ANNOTATIONS_KEY, {})
         if not item or not annotation_payload:
             # An error occured
             return response
-        annotations = [
-            BoxAnnotation.from_json(ann)
-            if ann["type"] == "box"
-            else PolygonAnnotation.from_json(ann)
-            for ann in annotation_payload
-        ]
+
+        annotation_response = {}
+        for annotation_type in ANNOTATION_TYPES:
+            if annotation_type in annotation_payload:
+                annotation_response[annotation_type] = [
+                    Annotation.from_json(ann)
+                    for ann in annotation_payload[annotation_type]
+                ]
         return {
             ITEM_KEY: DatasetItem.from_json(item),
-            ANNOTATIONS_KEY: annotations,
+            ANNOTATIONS_KEY: annotation_response,
         }

nucleus/model_run.py

@@ -1,11 +1,16 @@
-from typing import Optional, List, Dict, Any, Union
-from .constants import ANNOTATIONS_KEY, DEFAULT_ANNOTATION_UPDATE_MODE
+from typing import Optional, List, Union
+from .constants import (
+    ANNOTATIONS_KEY,
+    DEFAULT_ANNOTATION_UPDATE_MODE,
+    BOX_TYPE,
+    POLYGON_TYPE,
+    SEGMENTATION_TYPE,
+)
 from .prediction import (
     BoxPrediction,
     PolygonPrediction,
     SegmentationPrediction,
 )
-from .payload_constructor import construct_box_predictions_payload
 
 
 class ModelRun:
@@ -121,13 +126,18 @@ def loc(self, dataset_item_id: str):
     def _format_prediction_response(
         self, response: dict
     ) -> Union[dict, List[Union[BoxPrediction, PolygonPrediction]]]:
-        annotations = response.get(ANNOTATIONS_KEY, None)
-        if annotations:
-            return [
-                BoxPrediction.from_json(ann)
-                if ann["type"] == "box"
-                else PolygonPrediction.from_json(ann)
-                for ann in annotations
-            ]
+        annotation_payload = response.get(ANNOTATIONS_KEY, None)
+        if annotation_payload:
+            annotation_response = {}
+            for (type_key, type_cls) in zip(
+                [BOX_TYPE, POLYGON_TYPE, SEGMENTATION_TYPE],
+                [BoxPrediction, PolygonPrediction, SegmentationPrediction],
+            ):
+                if type_key in annotation_payload:
+                    annotation_response[type_key] = [
+                        type_cls.from_json(ann)
+                        for ann in annotation_payload[type_key]
+                    ]
+            return annotation_response
         else:  # An error occurred
             return response

tests/helpers.py

@@ -11,11 +11,11 @@
 TEST_DATASET_NAME = "[PyTest] Test Dataset"
 TEST_SLICE_NAME = "[PyTest] Test Slice"
 
-TEST_MODEL_NAME = '[PyTest] Test Model Name'
-TEST_MODEL_REFERENCE = '[PyTest] Test Model Reference'
-TEST_MODEL_RUN = '[PyTest] Test Model Run Reference'
-TEST_DATASET_NAME = '[PyTest] Test Dataset'
-TEST_SLICE_NAME = '[PyTest] Test Slice'
+TEST_MODEL_NAME = "[PyTest] Test Model Name"
+TEST_MODEL_REFERENCE = "[PyTest] Test Model Reference"
+TEST_MODEL_RUN = "[PyTest] Test Model Run Reference"
+TEST_DATASET_NAME = "[PyTest] Test Dataset"
+TEST_SLICE_NAME = "[PyTest] Test Slice"
 TEST_IMG_URLS = [
     "s3://scaleapi-attachments/BDD/BDD/bdd100k/images/100k/train/6dd63871-831611a6.jpg",
     "s3://scaleapi-attachments/BDD/BDD/bdd100k/images/100k/train/82c1005c-e2d1d94f.jpg",
@@ -24,16 +24,16 @@
     "s3://scaleapi-attachments/BDD/BDD/bdd100k/images/100k/train/89b42832-10d662f4.jpg",
 ]
 TEST_DATASET_ITEMS = [
-    DatasetItem(TEST_IMG_URLS[0], '1'),
-    DatasetItem(TEST_IMG_URLS[1], '2'),
-    DatasetItem(TEST_IMG_URLS[2], '3'),
-    DatasetItem(TEST_IMG_URLS[3], '4')
+    DatasetItem(TEST_IMG_URLS[0], "1"),
+    DatasetItem(TEST_IMG_URLS[1], "2"),
+    DatasetItem(TEST_IMG_URLS[2], "3"),
+    DatasetItem(TEST_IMG_URLS[3], "4"),
 ]
 TEST_PREDS = [
-    BoxPrediction('[Pytest Box Prediction 1]', 0, 0, 100, 100, '1'),
-    BoxPrediction('[Pytest Box Prediction 2]', 0, 0, 100, 100, '2'),
-    BoxPrediction('[Pytest Box Prediction 3]', 0, 0, 100, 100, '3'),
-    BoxPrediction('[Pytest Box Prediction 4]', 0, 0, 100, 100, '4')
+    BoxPrediction("[Pytest Box Prediction 1]", 0, 0, 100, 100, "1"),
+    BoxPrediction("[Pytest Box Prediction 2]", 0, 0, 100, 100, "2"),
+    BoxPrediction("[Pytest Box Prediction 3]", 0, 0, 100, 100, "3"),
+    BoxPrediction("[Pytest Box Prediction 4]", 0, 0, 100, 100, "4"),
 ]
 
 
@@ -147,6 +147,23 @@ def assert_polygon_annotation_matches_dict(
         assert instance_pt["y"] == dict_pt["y"]
 
 
+def assert_segmentation_annotation_matches_dict(
+    annotation_instance, annotation_dict
+):
+    assert annotation_instance.mask_url == annotation_dict["mask_url"]
+    # Cannot guarantee segments are in same order
+    assert len(annotation_instance.annotations) == len(
+        annotation_dict["annotations"]
+    )
+
+    for instance_segment, dict_segment in zip(
+        sorted(annotation_instance.annotations, key=lambda i: i.index),
+        sorted(annotation_dict["annotations"], key=lambda i: i["index"]),
+    ):
+        assert instance_segment.index == dict_segment["index"]
+        assert instance_segment.label == dict_segment["label"]
+
+
 # Asserts that a box prediction instance matches a dict representing its properties.
 # Useful to check prediction uploads/updates match.
 def assert_box_prediction_matches_dict(prediction_instance, prediction_dict):

tests/test_annotation.py

@@ -9,6 +9,7 @@
     reference_id_from_url,
     assert_box_annotation_matches_dict,
     assert_polygon_annotation_matches_dict,
+    assert_segmentation_annotation_matches_dict,
 )
 
 from nucleus import (
@@ -48,7 +49,7 @@ def test_box_gt_upload(dataset):
     assert response["annotations_processed"] == 1
     assert response["annotations_ignored"] == 0
 
-    response = dataset.refloc(annotation.reference_id)["annotations"]
+    response = dataset.refloc(annotation.reference_id)["annotations"]["box"]
     assert len(response) == 1
     response_annotation = response[0]
     assert_box_annotation_matches_dict(
@@ -64,7 +65,9 @@ def test_polygon_gt_upload(dataset):
     assert response["annotations_processed"] == 1
     assert response["annotations_ignored"] == 0
 
-    response = dataset.refloc(annotation.reference_id)["annotations"]
+    response = dataset.refloc(annotation.reference_id)["annotations"][
+        "polygon"
+    ]
     assert len(response) == 1
     response_annotation = response[0]
     assert_polygon_annotation_matches_dict(
@@ -80,7 +83,13 @@ def test_single_semseg_gt_upload(dataset):
     assert response["dataset_id"] == dataset.id
     assert response["annotations_processed"] == 1
     assert response["annotations_ignored"] == 0
-    # assert_box_annotation_matches_dict(response_annotation, TEST_BOX_ANNOTATIONS[0])
+
+    response_annotation = dataset.refloc(annotation.reference_id)[
+        "annotations"
+    ]["segmentation"][0]
+    assert_segmentation_annotation_matches_dict(
+        response_annotation, TEST_SEGMENTATION_ANNOTATIONS[0]
+    )
 
 
 def test_batch_semseg_gt_upload(dataset):
@@ -142,6 +151,12 @@ def test_mixed_annotation_upload(dataset):
     assert response["dataset_id"] == dataset.id
     assert response["annotations_processed"] == 10
     assert response["annotations_ignored"] == 0
+    response_annotations = dataset.refloc(bbox_annotations[0].reference_id)[
+        "annotations"
+    ]
+    assert len(response_annotations) == 2
+    assert len(response_annotations["box"]) == 1
+    assert "segmentation" in response_annotations
 
 
 def test_box_gt_upload_update(dataset):
@@ -165,7 +180,7 @@ def test_box_gt_upload_update(dataset):
     assert response["annotations_processed"] == 1
     assert response["annotations_ignored"] == 0
 
-    response = dataset.refloc(annotation.reference_id)["annotations"]
+    response = dataset.refloc(annotation.reference_id)["annotations"]["box"]
     assert len(response) == 1
     response_annotation = response[0]
     assert_box_annotation_matches_dict(
@@ -194,7 +209,7 @@ def test_box_gt_upload_ignore(dataset):
     assert response["annotations_processed"] == 1
     assert response["annotations_ignored"] == 1
 
-    response = dataset.refloc(annotation.reference_id)["annotations"]
+    response = dataset.refloc(annotation.reference_id)["annotations"]["box"]
     assert len(response) == 1
     response_annotation = response[0]
     assert_box_annotation_matches_dict(
@@ -223,7 +238,9 @@ def test_polygon_gt_upload_update(dataset):
     assert response["annotations_processed"] == 1
     assert response["annotations_ignored"] == 0
 
-    response = dataset.refloc(annotation.reference_id)["annotations"]
+    response = dataset.refloc(annotation.reference_id)["annotations"][
+        "polygon"
+    ]
     assert len(response) == 1
     response_annotation = response[0]
     assert_polygon_annotation_matches_dict(
@@ -253,7 +270,9 @@ def test_polygon_gt_upload_ignore(dataset):
     assert response["annotations_processed"] == 1
     assert response["annotations_ignored"] == 1
 
-    response = dataset.refloc(annotation.reference_id)["annotations"]
+    response = dataset.refloc(annotation.reference_id)["annotations"][
+        "polygon"
+    ]
     assert len(response) == 1
     response_annotation = response[0]
     assert_polygon_annotation_matches_dict(