address PR comments 2

Author: sasha-scale

nucleus/model_run.py

@@ -126,25 +126,13 @@ 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:
+        annotation_payload = response.get(ANNOTATIONS_KEY, None)
+        if annotation_payload:
             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]
-                )
+            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

@@ -156,6 +156,10 @@ def assert_segmentation_annotation_matches_dict(
         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.

tests/test_annotation.py

@@ -86,7 +86,7 @@ def test_single_semseg_gt_upload(dataset):
 
     response_annotation = dataset.refloc(annotation.reference_id)[
         "annotations"
-    ]["segmentation"]
+    ]["segmentation"][0]
     assert_segmentation_annotation_matches_dict(
         response_annotation, TEST_SEGMENTATION_ANNOTATIONS[0]
     )

tests/test_prediction.py

@@ -92,10 +92,10 @@ def test_segmentation_pred_upload(model_run):
     assert response["predictions_ignored"] == 0
 
     response = model_run.refloc(prediction.reference_id)["segmentation"]
-    assert isinstance(response, SegmentationPrediction)
+    assert isinstance(response[0], SegmentationPrediction)
 
     assert_segmentation_annotation_matches_dict(
-        response, TEST_SEGMENTATION_PREDICTIONS[0]
+        response[0], TEST_SEGMENTATION_PREDICTIONS[0]
     )
 
 
@@ -238,3 +238,12 @@ def test_mixed_pred_upload(model_run):
     assert response["predictions_ignored"] == 0
 
     response_refloc = model_run.refloc(prediction_polygon.reference_id)
+    assert_box_prediction_matches_dict(
+        response_refloc["box"][0], TEST_BOX_PREDICTIONS[0]
+    )
+    assert_polygon_prediction_matches_dict(
+        response_refloc["polygon"][0], TEST_POLYGON_PREDICTIONS[0]
+    )
+    assert_segmentation_annotation_matches_dict(
+        response_refloc["segmentation"][0], TEST_SEGMENTATION_PREDICTIONS[0]
+    )