lint and up version

Author: jihan-yin

nucleus/__init__.py

@@ -951,6 +951,40 @@ 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

@@ -22,6 +22,7 @@
     DEFAULT_ANNOTATION_UPDATE_MODE,
     EXPORTED_ROWS,
     NAME_KEY,
+    REFERENCE_ID_KEY,
     REFERENCE_IDS_KEY,
     REQUEST_ID_KEY,
     AUTOTAG_SCORE_THRESHOLD,
@@ -385,6 +386,28 @@ 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):
+        """
+        Returns info for single ground truth Annotation by its id.
+        :param prediction_id: internally controlled id for ground truth annotations.
+        :return:
+        {
+            "ref_id": Reference id of dataset item associated with this annotation
+            "groundTruth": 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],
+                }
+            ),
+        }
+
     def create_slice(
         self,
         name: str,

nucleus/model_run.py

@@ -160,6 +160,26 @@ def loc(self, dataset_item_id: str):
         )
         return self._format_prediction_response(response)
 
+    def prediction_loc(self, prediction_id: str):
+        """
+        Returns info for single Prediction by its id.
+        :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
+        }
+        """
+        response = self._client.prediction_loc(
+            self.model_run_id, prediction_id
+        )
+        return {
+            "ref_id": response["ref_id"],
+            "prediction": self._format_prediction_response(
+                {"predictions": [response["prediction"]]}
+            ),
+        }
+
     def ungrouped_export(self):
         json_response = self._client.make_request(
             payload={},

pyproject.toml

@@ -21,7 +21,7 @@ exclude = '''
 
 [tool.poetry]
 name = "scale-nucleus"
-version = "0.1.21"
+version = "0.1.22"
 description = "The official Python client library for Nucleus, the Data Platform for AI"
 license =  "MIT"
 authors = ["Scale AI Nucleus Team <nucleusapi@scaleapi.com>"]

tests/helpers.py

@@ -14,6 +14,11 @@
 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_dataset.py

@@ -33,6 +33,8 @@
 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,
@@ -100,6 +102,15 @@ 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

@@ -1,7 +1,12 @@
 from nucleus.job import AsyncJob
+import os
 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,
@@ -76,6 +81,19 @@ 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])