incorporate additional tests

Author: sasha-scale

nucleus/annotation.py

@@ -21,6 +21,8 @@
     ANNOTATIONS_KEY,
 )
 
+ORIGINAL_URL_KEY = "original_url"
+
 
 class Segment:
     def __init__(
@@ -74,7 +76,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, None),
             annotations=[
                 Segment.from_json(ann)
                 for ann in payload.get(ANNOTATIONS_KEY, [])

nucleus/constants.py

@@ -1,5 +1,4 @@
-# NUCLEUS_ENDPOINT = "https://api.scale.com/v1/nucleus"
-NUCLEUS_ENDPOINT = "http://localhost:3000/v1/nucleus"
+NUCLEUS_ENDPOINT = "https://api.scale.com/v1/nucleus"
 ITEMS_KEY = "items"
 ITEM_KEY = "item"
 REFERENCE_ID_KEY = "reference_id"

nucleus/dataset.py

@@ -253,16 +253,11 @@ def _format_dataset_item_response(self, response: dict) -> dict:
                 for ann in annotation_payload[POLYGON_TYPE]
             ]
         if SEGMENTATION_TYPE in annotation_payload:
-            annotation_response[SEGMENTATION_TYPE] = [
-                SegmentationAnnotation.from_json(ann)
-                for ann in annotation_payload[SEGMENTATION_TYPE]
-            ]
-        # annotations = [
-        #     BoxAnnotation.from_json(ann)
-        #     if ann["type"] == "box"
-        #     else PolygonAnnotation.from_json(ann)
-        #     for ann in annotation_payload
-        # ]
+            annotation_response[
+                SEGMENTATION_TYPE
+            ] = SegmentationAnnotation.from_json(
+                annotation_payload[SEGMENTATION_TYPE]
+            )
         return {
             ITEM_KEY: DatasetItem.from_json(item),
             ANNOTATIONS_KEY: annotation_response,

nucleus/model_run.py

@@ -1,5 +1,11 @@
 from typing import Optional, List, Dict, Any, Union
-from .constants import ANNOTATIONS_KEY, DEFAULT_ANNOTATION_UPDATE_MODE
+from .constants import (
+    ANNOTATIONS_KEY,
+    DEFAULT_ANNOTATION_UPDATE_MODE,
+    BOX_TYPE,
+    POLYGON_TYPE,
+    SEGMENTATION_TYPE,
+)
 from .prediction import (
     BoxPrediction,
     PolygonPrediction,
@@ -122,12 +128,26 @@ def _format_prediction_response(
         self, response: dict
     ) -> Union[dict, List[Union[BoxPrediction, PolygonPrediction]]]:
         annotations = response.get(ANNOTATIONS_KEY, None)
+        print(annotations)
         if annotations:
-            return [
-                BoxPrediction.from_json(ann)
-                if ann["type"] == "box"
-                else PolygonPrediction.from_json(ann)
-                for ann in annotations
-            ]
+            annotation_response = {}
+            if BOX_TYPE in annotations:
+                annotation_response[BOX_TYPE] = [
+                    BoxPrediction.from_json(ann)
+                    for ann in annotations[BOX_TYPE]
+                ]
+            if POLYGON_TYPE in annotations:
+                annotation_response[POLYGON_TYPE] = [
+                    PolygonPrediction.from_json(ann)
+                    for ann in annotations[POLYGON_TYPE]
+                ]
+            if SEGMENTATION_TYPE in annotations:
+                annotation_response[
+                    SEGMENTATION_TYPE
+                ] = SegmentationPrediction.from_json(
+                    annotations[SEGMENTATION_TYPE]
+                )
+            print(annotation_response)
+            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,16 @@ 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"]
+    )
+
+
 # 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

@@ -1,6 +1,6 @@
 import pytest
 
-from helpers import (
+from .helpers import (
     TEST_DATASET_NAME,
     TEST_IMG_URLS,
     TEST_BOX_ANNOTATIONS,
@@ -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"]
+    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(

tests/test_dataset.py

@@ -1,6 +1,6 @@
 import pytest
 
-from helpers import (
+from .helpers import (
     TEST_SLICE_NAME,
     TEST_DATASET_NAME,
     TEST_IMG_URLS,
@@ -17,6 +17,7 @@
     DATASET_ID_KEY,
 )
 
+
 @pytest.fixture()
 def dataset(CLIENT):
     ds = CLIENT.create_dataset(TEST_DATASET_NAME)
@@ -25,6 +26,7 @@ def dataset(CLIENT):
     response = CLIENT.delete_dataset(ds.id)
     assert response == {}
 
+
 def test_dataset_create_and_delete(CLIENT):
     # Creation
     ds = CLIENT.create_dataset(TEST_DATASET_NAME)
@@ -88,14 +90,17 @@ def test_dataset_list_autotags(CLIENT, dataset):
     autotag_response = CLIENT.list_autotags(dataset.id)
     assert autotag_response == []
 
+
 def test_slice_create_and_delete(dataset):
     # Dataset upload
     ds_items = []
     for url in TEST_IMG_URLS:
-        ds_items.append(DatasetItem(
-            image_location=url,
-            reference_id=reference_id_from_url(url),
-        ))
+        ds_items.append(
+            DatasetItem(
+                image_location=url,
+                reference_id=reference_id_from_url(url),
+            )
+        )
     response = dataset.append(ds_items)
     assert ERROR_PAYLOAD not in response.json()
 
@@ -121,10 +126,12 @@ def test_slice_append(dataset):
     # Dataset upload
     ds_items = []
     for url in TEST_IMG_URLS:
-        ds_items.append(DatasetItem(
-            image_location=url,
-            reference_id=reference_id_from_url(url),
-        ))
+        ds_items.append(
+            DatasetItem(
+                image_location=url,
+                reference_id=reference_id_from_url(url),
+            )
+        )
     response = dataset.append(ds_items)
     assert ERROR_PAYLOAD not in response.json()