relevant tests pass

Author: Ubuntu

conftest.py

@@ -28,6 +28,13 @@ def dataset(CLIENT):
     CLIENT.delete_dataset(ds.id)
 
 
+@pytest.fixture()
+def model(CLIENT):
+    model = CLIENT.create_model(TEST_DATASET_NAME, "fake_reference_id")
+    yield model
+    CLIENT.delete_model(model.id)
+
+
 if __name__ == "__main__":
     client = nucleus.NucleusClient(API_KEY)
     # ds = client.create_dataset("Test Dataset With Autotags")

nucleus/annotation_uploader.py

@@ -24,15 +24,20 @@ def accumulate_dict_values(dicts: Iterable[dict]):
     result = {}
     for d in dicts:
         for key, value in d.items():
-            if key not in result or key == "dataset_id":
+            if (
+                key not in result
+                or key == "dataset_id"
+                or key == "model_run_id"
+            ):
                 result[key] = value
             else:
                 result[key] += value
     return result
 
 
 class AnnotationUploader:
-    """This is a helper class not intended for direct use. Please use dataset.annotate.
+    """This is a helper class not intended for direct use. Please use dataset.annotate
+    or dataset.upload_predictions.
 
     This class is purely a helper class for implementing dataset.annotate/dataset.predict.
     """

tests/test_prediction.py

@@ -194,39 +194,6 @@ def test_non_existent_taxonomy_category_gt_upload(model_run):
     )
 
 
-def test_segmentation_pred_upload(model_run):
-    prediction = SegmentationPrediction.from_json(
-        TEST_SEGMENTATION_PREDICTIONS[0]
-    )
-    response = model_run.predict(annotations=[prediction])
-
-    assert response["model_run_id"] == model_run.model_run_id
-    assert response["predictions_processed"] == 1
-    assert response["predictions_ignored"] == 0
-
-    response = model_run.refloc(prediction.reference_id)["segmentation"]
-    assert isinstance(response[0], SegmentationPrediction)
-
-    assert_segmentation_annotation_matches_dict(
-        response[0], TEST_SEGMENTATION_PREDICTIONS[0]
-    )
-
-
-def test_segmentation_pred_upload_ignore(model_run):
-    prediction = SegmentationPrediction.from_json(
-        TEST_SEGMENTATION_PREDICTIONS[0]
-    )
-    response1 = model_run.predict(annotations=[prediction])
-
-    assert response1["predictions_processed"] == 1
-
-    # Upload Duplicate annotation
-    response = model_run.predict(annotations=[prediction])
-    assert response["model_run_id"] == model_run.model_run_id
-    assert response["predictions_processed"] == 0
-    assert response["predictions_ignored"] == 1
-
-
 def test_box_pred_upload_update(model_run):
     prediction = BoxPrediction(**TEST_BOX_PREDICTIONS[0])
     response = model_run.predict(annotations=[prediction])

tests/test_segmentation.py

@@ -1,13 +1,65 @@
 from nucleus.annotation import SegmentationAnnotation
 from nucleus.dataset import Dataset
+from nucleus.model import Model
+from nucleus.prediction import SegmentationPrediction
 from tests.helpers import (
     NUM_VALID_SEGMENTATIONS_IN_MAIN_DATASET,
     TEST_LOCAL_MASK_URL,
     TEST_SEGMENTATION_ANNOTATIONS,
+    TEST_SEGMENTATION_PREDICTIONS,
     assert_segmentation_annotation_matches_dict,
 )
 
 
+def test_segmentation_pred_upload_local(dataset: Dataset, model: Model):
+    prediction = SegmentationPrediction.from_json(
+        TEST_SEGMENTATION_PREDICTIONS[0]
+    )
+    prediction.mask_url = TEST_LOCAL_MASK_URL
+    response = dataset.upload_predictions(model, [prediction])
+
+    assert response["predictions_processed"] == 1
+
+    response = dataset.predictions_refloc(model, prediction.reference_id)[
+        "segmentation"
+    ][0]
+    assert isinstance(response, SegmentationPrediction)
+    assert response == prediction
+
+
+def test_segmentation_pred_upload(dataset: Dataset, model: Model):
+    prediction = SegmentationPrediction.from_json(
+        TEST_SEGMENTATION_PREDICTIONS[0]
+    )
+    response = dataset.upload_predictions(model, [prediction])
+
+    assert response["predictions_processed"] == 1
+    assert response["predictions_ignored"] == 0
+
+    response = dataset.predictions_refloc(model, prediction.reference_id)[
+        "segmentation"
+    ]
+    assert isinstance(response[0], SegmentationPrediction)
+
+    assert_segmentation_annotation_matches_dict(
+        response[0], TEST_SEGMENTATION_PREDICTIONS[0]
+    )
+
+
+def test_segmentation_pred_upload_ignore(dataset: Dataset, model: Model):
+    prediction = SegmentationPrediction.from_json(
+        TEST_SEGMENTATION_PREDICTIONS[0]
+    )
+    response1 = dataset.upload_predictions(model, [prediction])
+
+    assert response1["predictions_processed"] == 1
+
+    # Upload Duplicate annotation
+    response = dataset.upload_predictions(model, [prediction])
+    assert response["predictions_processed"] == 0
+    assert response["predictions_ignored"] == 1
+
+
 def test_single_local_semseg_gt_upload(dataset: Dataset):
     request_annotation = SegmentationAnnotation.from_json(
         TEST_SEGMENTATION_ANNOTATIONS[0]