Switch to using annotation id

Author: Ubuntu

nucleus/__init__.py

@@ -951,40 +951,6 @@ def predictions_loc(self, model_run_id: str, dataset_item_id: str):
             {}, f"modelRun/{model_run_id}/loc/{dataset_item_id}", requests.get
         )
 
-    def prediction_loc(self, model_run_id: str, prediction_id: str):
-        """
-        Returns info for single Prediction by its id.
-        :param model_run_id: id of the model run
-        :param prediction_id: internally controlled id for model predictions.
-        :return:
-        {
-            "ref_id": Reference id of dataset item associated with this prediction
-            "prediction": BoxPrediction | PolygonPrediction | CuboidPrediction
-        }
-        """
-        return self.make_request(
-            {},
-            f"modelRun/{model_run_id}/prediction/loc/{prediction_id}",
-            requests.get,
-        )
-
-    def ground_truth_loc(self, dataset_id: str, ground_truth_id: str):
-        """
-        Returns info for single Prediction by its id.
-        :param dataset_id: id of the dataset
-        :param ground_truth_id: internally controlled id for ground truth annotations.
-        :return:
-        {
-            "ref_id": Reference id of dataset item associated with this prediction
-            "ground_truth": BoxAnnotation | PolygonAnnotation | CuboidAnnotation
-        }
-        """
-        return self.make_request(
-            {},
-            f"dataset/{dataset_id}/groundTruth/loc/{ground_truth_id}",
-            requests.get,
-        )
-
     def create_slice(self, dataset_id: str, payload: dict) -> Slice:
         """
         Creates a slice from items already present in a dataset.

nucleus/dataset.py

@@ -386,27 +386,19 @@ def loc(self, dataset_item_id: str) -> dict:
         response = self._client.dataitem_loc(self.id, dataset_item_id)
         return format_dataset_item_response(response)
 
-    def ground_truth_loc(self, ground_truth_id: str):
+    def ground_truth_loc(self, reference_id: str, annotation_id: str):
         """
         Returns info for single ground truth Annotation by its id.
-        :param prediction_id: internally controlled id for ground truth annotations.
+        :param annotation_id: internally controlled id for ground truth annotations.
         :return:
-        {
-            "ref_id": Reference id of dataset item associated with this annotation
-            "groundTruth": BoxAnnotation | PolygonAnnotation | CuboidAnnotation
-        }
+        BoxAnnotation | PolygonAnnotation | CuboidAnnotation
         """
-        response = self._client.ground_truth_loc(self.id, ground_truth_id)
-        annotation_type = response["groundTruth"].keys()[0]
-        return {
-            "ref_id": response["ref_id"],
-            "ground_truth": Annotation.from_json(
-                {
-                    REFERENCE_ID_KEY: response["ref_id"],
-                    **response["groundTruth"][annotation_type],
-                }
-            ),
-        }
+        response = self._client.make_request(
+            {},
+            f"dataset/{self.id}/groundTruth/loc/{reference_id}/{annotation_id}",
+            requests.get,
+        )
+        return Annotation.from_json(response)
 
     def create_slice(
         self,

nucleus/model_run.py

@@ -19,6 +19,7 @@
     CuboidPrediction,
     PolygonPrediction,
     SegmentationPrediction,
+    from_json,
 )
 
 
@@ -160,25 +161,22 @@ def loc(self, dataset_item_id: str):
         )
         return self._format_prediction_response(response)
 
-    def prediction_loc(self, prediction_id: str):
+    def prediction_loc(self, reference_id: str, annotation_id: str):
         """
-        Returns info for single Prediction by its id.
-        :param prediction_id: internally controlled id for model predictions.
+        Returns info for single Prediction by its reference id and annotation id.
+        :param reference_id: the user specified id for the image
+        :param annotation_id: the user specified id for the prediction
         :return:
-        {
-            "ref_id": Reference id of dataset item associated with this prediction
-            "prediction": BoxPrediction | PolygonPrediction | CuboidPrediction
-        }
+         BoxPrediction | PolygonPrediction | CuboidPrediction
         """
-        response = self._client.prediction_loc(
-            self.model_run_id, prediction_id
+
+        response = self._client.make_request(
+            {},
+            f"modelRun/{self.model_run_id}/prediction/loc/{reference_id}/{annotation_id}",
+            requests.get,
         )
-        return {
-            "ref_id": response["ref_id"],
-            "prediction": self._format_prediction_response(
-                {"predictions": [response["prediction"]]}
-            ),
-        }
+
+        return from_json(response)
 
     def ungrouped_export(self):
         json_response = self._client.make_request(

nucleus/prediction.py

@@ -10,10 +10,14 @@
 )
 from .constants import (
     ANNOTATION_ID_KEY,
+    BOX_TYPE,
+    CUBOID_TYPE,
+    POLYGON_TYPE,
     REFERENCE_ID_KEY,
     METADATA_KEY,
     GEOMETRY_KEY,
     LABEL_KEY,
+    TYPE_KEY,
     X_KEY,
     Y_KEY,
     WIDTH_KEY,
@@ -29,6 +33,17 @@
 )
 
 
+def from_json(payload: dict):
+    if payload.get(TYPE_KEY, None) == BOX_TYPE:
+        return BoxPrediction.from_json(payload)
+    elif payload.get(TYPE_KEY, None) == POLYGON_TYPE:
+        return PolygonPrediction.from_json(payload)
+    elif payload.get(TYPE_KEY, None) == CUBOID_TYPE:
+        return CuboidPrediction.from_json(payload)
+    else:
+        return SegmentationPrediction.from_json(payload)
+
+
 class SegmentationPrediction(SegmentationAnnotation):
     # No need to define init or to_payload methods because
     # we default to functions defined in the parent class

tests/helpers.py

@@ -14,11 +14,6 @@
 TEST_PROJECT_ID = "60b699d70f139e002dd31bfc"
 
 DATASET_WITH_AUTOTAG = "ds_c4dgj702e2vjft7m9xa0"
-DATASET_WITH_PREDICTIONS = "ds_c4283fxxjfvg05rmm5p0"
-DATASET_WITH_PREDICTIONS_MODEL_RUN_ID = "run_c4283gpxjfvg0ehz6gjg"
-DATASET_WITH_PREDICTIONS_SAMPLE_PREDICTION_ID = "pred_c4283gp2s2b00esdarq0"
-DATASET_WITH_GROUND_TRUTH = "ds_c3pz51mjy60g09hfzhfg"
-DATASET_WITH_GROUND_TRUTH_SAMPLE_GROUND_TRUTH_ID = "ann_c3pz524k221g0css4ns0"
 NUCLEUS_PYTEST_USER_ID = "60ad648c85db770026e9bf77"
 
 

tests/test_annotation.py

@@ -70,6 +70,11 @@ def test_box_gt_upload(dataset):
     assert response["annotations_ignored"] == 0
 
     response = dataset.refloc(annotation.reference_id)["annotations"]["box"]
+    single_annotation_response = dataset.ground_truth_loc(
+        annotation.reference_id, annotation.annotation_id
+    )
+
+    assert response[0] == single_annotation_response
     assert len(response) == 1
     response_annotation = response[0]
     assert_box_annotation_matches_dict(

tests/test_dataset.py

@@ -33,8 +33,6 @@
 from nucleus.job import AsyncJob, JobError
 
 from .helpers import (
-    DATASET_WITH_GROUND_TRUTH,
-    DATASET_WITH_GROUND_TRUTH_SAMPLE_GROUND_TRUTH_ID,
     LOCAL_FILENAME,
     TEST_BOX_ANNOTATIONS,
     TEST_DATASET_NAME,
@@ -102,15 +100,6 @@ def make_dataset_items():
     return ds_items_with_metadata
 
 
-def test_dataset_ground_truth_loc(CLIENT):
-    ds = CLIENT.get_dataset(DATASET_WITH_GROUND_TRUTH)
-    response = ds.ground_truth_loc(
-        DATASET_WITH_GROUND_TRUTH_SAMPLE_GROUND_TRUTH_ID
-    )
-    assert response["ref_id"] is not None
-    assert response["ground_truth"] is not None
-
-
 def test_dataset_create_and_delete(CLIENT):
     # Creation
     ds = CLIENT.create_dataset(TEST_DATASET_NAME)

tests/test_prediction.py

@@ -3,10 +3,6 @@
 import pytest
 import time
 from .helpers import (
-    DATASET_WITH_PREDICTIONS,
-    DATASET_WITH_PREDICTIONS_MODEL_RUN_ID,
-    DATASET_WITH_PREDICTIONS_SAMPLE_PREDICTION_ID,
-    NUCLEUS_PYTEST_USER_ID,
     TEST_DATASET_NAME,
     TEST_MODEL_NAME,
     TEST_MODEL_RUN,
@@ -25,14 +21,10 @@
     PolygonPrediction,
     SegmentationPrediction,
     DatasetItem,
-    Segment,
     ModelRun,
-    Point,
 )
 from nucleus.constants import ERROR_PAYLOAD
 
-from nucleus import utils
-
 
 def test_reprs():
     def test_repr(test_object: any):
@@ -81,19 +73,6 @@ def model_run(CLIENT):
     assert response == {}
 
 
-def test_pred_loc(CLIENT):
-    if NUCLEUS_PYTEST_USER_ID in os.environ["NUCLEUS_PYTEST_USER_ID"]:
-        run = CLIENT.get_model_run(
-            DATASET_WITH_PREDICTIONS_MODEL_RUN_ID, DATASET_WITH_PREDICTIONS
-        )
-        response = run.prediction_loc(
-            DATASET_WITH_PREDICTIONS_SAMPLE_PREDICTION_ID
-        )
-
-        assert response["ref_id"] is not None
-        assert response["prediction"] is not None
-
-
 def test_box_pred_upload(model_run):
     prediction = BoxPrediction(**TEST_BOX_PREDICTIONS[0])
     response = model_run.predict(annotations=[prediction])
@@ -103,6 +82,10 @@ def test_box_pred_upload(model_run):
     assert response["predictions_ignored"] == 0
 
     response = model_run.refloc(prediction.reference_id)["box"]
+    single_prediction = model_run.prediction_loc(
+        prediction.reference_id, prediction.annotation_id
+    )
+    assert response[0] == single_prediction
     assert len(response) == 1
     assert_box_prediction_matches_dict(response[0], TEST_BOX_PREDICTIONS[0])