Merge pull request #81 from scaleapi/da-batch-model-export

Batch export for model predictions

Author: ardila

nucleus/annotation.py

@@ -37,7 +37,9 @@ class Annotation:
     item_id: Optional[str] = None
 
     def _check_ids(self):
-        if bool(self.reference_id) == bool(self.item_id):
+        if self.reference_id and self.item_id:
+            self.item_id = None  # Prefer reference id to item id.
+        if not (self.reference_id or self.item_id):
             raise Exception(
                 "You must specify either a reference_id or an item_id for an annotation."
             )

nucleus/model_run.py

@@ -1,5 +1,7 @@
 from typing import Dict, List, Optional, Type, Union
 
+import requests
+
 from nucleus.annotation import check_all_annotation_paths_remote
 from nucleus.job import AsyncJob
 from nucleus.utils import serialize_and_write_to_presigned_url
@@ -155,6 +157,16 @@ def loc(self, dataset_item_id: str):
         )
         return self._format_prediction_response(response)
 
+    def ungrouped_export(self):
+        json_response = self._client.make_request(
+            payload={},
+            route=f"modelRun/{self.model_run_id}/ungrouped",
+            requests_command=requests.get,
+        )
+        return self._format_prediction_response(
+            {ANNOTATIONS_KEY: json_response}
+        )
+
     def _format_prediction_response(
         self, response: dict
     ) -> Union[

tests/test_prediction.py

@@ -50,7 +50,7 @@ def test_repr(test_object: any):
 def model_run(CLIENT):
     ds = CLIENT.create_dataset(TEST_DATASET_NAME)
     ds_items = []
-    for url in TEST_IMG_URLS:
+    for url in TEST_IMG_URLS[:2]:
         ds_items.append(
             DatasetItem(
                 image_location=url,
@@ -246,7 +246,7 @@ def test_polygon_pred_upload_ignore(model_run):
     )
 
 
-def test_mixed_pred_upload(model_run):
+def test_mixed_pred_upload(model_run: ModelRun):
     prediction_semseg = SegmentationPrediction.from_json(
         TEST_SEGMENTATION_PREDICTIONS[0]
     )
@@ -262,15 +262,15 @@ def test_mixed_pred_upload(model_run):
     assert response["predictions_processed"] == 3
     assert response["predictions_ignored"] == 0
 
-    response_refloc = model_run.refloc(prediction_polygon.reference_id)
+    all_predictions = model_run.ungrouped_export()
     assert_box_prediction_matches_dict(
-        response_refloc["box"][0], TEST_BOX_PREDICTIONS[0]
+        all_predictions["box"][0], TEST_BOX_PREDICTIONS[0]
     )
     assert_polygon_prediction_matches_dict(
-        response_refloc["polygon"][0], TEST_POLYGON_PREDICTIONS[0]
+        all_predictions["polygon"][0], TEST_POLYGON_PREDICTIONS[0]
     )
     assert_segmentation_annotation_matches_dict(
-        response_refloc["segmentation"][0], TEST_SEGMENTATION_PREDICTIONS[0]
+        all_predictions["segmentation"][0], TEST_SEGMENTATION_PREDICTIONS[0]
     )