Model-only endpoints

Author: Ubuntu

nucleus/__init__.py

@@ -734,7 +734,9 @@ def create_model_run(self, dataset_id: str, payload: dict) -> ModelRun:
 
     def predict(
         self,
-        model_run_id: str,
+        model_run_id: Optional[str],
+        model_id: Optional[str],
+        dataset_id: Optional[str],
         annotations: List[
             Union[
                 BoxPrediction,
@@ -758,6 +760,16 @@ def predict(
             "predictions_ignored": int,
         }
         """
+        if model_run_id is not None:
+            assert model_id is None and dataset_id is None
+            endpoint = f"modelRun/{model_run_id}/predict"
+        else:
+            assert (
+                model_id is not None and dataset_id is not None
+            ), "Model ID and dataset ID are required if not using model run id."
+            endpoint = (
+                f"dataset/{dataset_id}/model/{model_id}/uploadPredictions"
+            )
         segmentations = [
             ann
             for ann in annotations
@@ -793,9 +805,7 @@ def predict(
                 batch,
                 update,
             )
-            response = self.make_request(
-                batch_payload, f"modelRun/{model_run_id}/predict"
-            )
+            response = self.make_request(batch_payload, endpoint)
             if STATUS_CODE_KEY in response:
                 agg_response[ERRORS_KEY] = response
             else:
@@ -808,9 +818,7 @@ def predict(
 
         for s_batch in s_batches:
             payload = construct_segmentation_payload(s_batch, update)
-            response = self.make_request(
-                payload, f"modelRun/{model_run_id}/predict_segmentation"
-            )
+            response = self.make_request(payload, endpoint)
             # pbar.update(1)
             if STATUS_CODE_KEY in response:
                 agg_response[ERRORS_KEY] = response

nucleus/dataset.py

@@ -3,18 +3,19 @@
 import requests
 
 from nucleus.job import AsyncJob
+from nucleus.prediction import from_json
 from nucleus.url_utils import sanitize_string_args
 from nucleus.utils import (
     convert_export_payload,
     format_dataset_item_response,
+    format_prediction_response,
     serialize_and_write_to_presigned_url,
 )
 
-from .annotation import (
-    Annotation,
-    check_all_mask_paths_remote,
-)
+from .annotation import Annotation, check_all_mask_paths_remote
 from .constants import (
+    ANNOTATIONS_KEY,
+    AUTOTAG_SCORE_THRESHOLD,
     DATASET_LENGTH_KEY,
     DATASET_MODEL_RUNS_KEY,
     DATASET_NAME_KEY,
@@ -24,20 +25,19 @@
     NAME_KEY,
     REFERENCE_IDS_KEY,
     REQUEST_ID_KEY,
-    AUTOTAG_SCORE_THRESHOLD,
     UPDATE_KEY,
 )
 from .dataset_item import (
     DatasetItem,
     check_all_paths_remote,
     check_for_duplicate_reference_ids,
 )
-from .scene import LidarScene, Scene, check_all_scene_paths_remote
 from .payload_constructor import (
     construct_append_scenes_payload,
     construct_model_run_creation_payload,
     construct_taxonomy_payload,
 )
+from .scene import LidarScene, Scene, check_all_scene_paths_remote
 
 WARN_FOR_LARGE_UPLOAD = 50000
 WARN_FOR_LARGE_SCENES_UPLOAD = 5
@@ -525,3 +525,89 @@ def get_scene(self, reference_id) -> Scene:
                 requests_command=requests.get,
             )
         )
+
+    def export_predictions(self, model):
+        json_response = self._client.make_request(
+            payload=None,
+            route=f"dataset/{self.id}/model/{model.id}/export",
+            requests_command=requests.get,
+        )
+        return format_prediction_response({ANNOTATIONS_KEY: json_response})
+
+    def calculate_evaluation_metrics(self, model, options=None):
+        """
+        class_agnostic -- A flag to specify if matching algorithm should be class-agnostic or not.
+                          Default value: True
+
+        allowed_label_matches -- An optional list of AllowedMatch objects to specify allowed matches
+                                 for ground truth and model predictions.
+                                 If specified, 'class_agnostic' flag is assumed to be False
+
+        Type 'AllowedMatch':
+        {
+            ground_truth_label: string,       # A label for ground truth annotation.
+            model_prediction_label: string,   # A label for model prediction that can be matched with
+                                              # corresponding ground truth label.
+        }
+
+        payload:
+        {
+            "class_agnostic": boolean,
+            "allowed_label_matches": List[AllowedMatch],
+        }"""
+        if options == None:
+            options = {}
+        return self._client.make_request(
+            payload=options,
+            route=f"dataset/{self.id}/model/{model.id}/calculateEvaluationMetrics",
+        )
+
+    def upload_predictions(
+        self, model, predictions, update=False, asynchronous=False
+    ):
+        if asynchronous:
+            check_all_mask_paths_remote(predictions)
+
+            request_id = serialize_and_write_to_presigned_url(
+                predictions, self.id, self._client
+            )
+            response = self._client.make_request(
+                payload={REQUEST_ID_KEY: request_id, UPDATE_KEY: update},
+                route=f"dataset/{self.id}/model/{model.id}/uploadPredictions?async=1",
+            )
+            return AsyncJob.from_json(response, self._client)
+        else:
+            return self._client.predict(
+                model_run_id=None,
+                dataset_id=self.id,
+                model_id=model.id,
+                annotations=predictions,
+                update=update,
+            )
+
+    def predictions_iloc(self, model, index):
+        return format_prediction_response(
+            self._client.make_request(
+                payload=None,
+                route=f"dataset/{self.id}/model/{model.id}/iloc/{index}",
+                requests_command=requests.get,
+            )
+        )
+
+    def predictions_refloc(self, model, reference_id):
+        return format_prediction_response(
+            self._client.make_request(
+                payload=None,
+                route=f"dataset/{self.id}/model/{model.id}/referenceId/{reference_id}",
+                requests_command=requests.get,
+            )
+        )
+
+    def prediction_loc(self, model, reference_id, annotation_id):
+        return from_json(
+            self._client.make_request(
+                payload=None,
+                route=f"dataset/{self.id}/model/{model.id}/loc/{reference_id}/{annotation_id}",
+                requests_command=requests.get,
+            )
+        )

nucleus/model.py

@@ -1,4 +1,4 @@
-from typing import List, Optional, Dict, Union
+from typing import List, Optional, Dict, Type, Union
 from .dataset import Dataset
 from .prediction import (
     BoxPrediction,
@@ -8,9 +8,14 @@
 )
 from .model_run import ModelRun
 from .constants import (
+    ANNOTATIONS_KEY,
+    BOX_TYPE,
+    CUBOID_TYPE,
     NAME_KEY,
+    POLYGON_TYPE,
     REFERENCE_ID_KEY,
     METADATA_KEY,
+    SEGMENTATION_TYPE,
 )
 
 

nucleus/model_run.py

@@ -1,17 +1,18 @@
-from typing import Dict, List, Optional, Type, Union
+from typing import List, Optional, Union
+
 import requests
+
 from nucleus.annotation import check_all_mask_paths_remote
 from nucleus.job import AsyncJob
-from nucleus.utils import serialize_and_write_to_presigned_url
+from nucleus.utils import (
+    format_prediction_response,
+    serialize_and_write_to_presigned_url,
+)
 
 from .constants import (
     ANNOTATIONS_KEY,
-    BOX_TYPE,
-    CUBOID_TYPE,
     DEFAULT_ANNOTATION_UPDATE_MODE,
-    POLYGON_TYPE,
     REQUEST_ID_KEY,
-    SEGMENTATION_TYPE,
     UPDATE_KEY,
 )
 from .prediction import (
@@ -124,7 +125,11 @@ def predict(
             )
             return AsyncJob.from_json(response, self._client)
         else:
-            return self._client.predict(self.model_run_id, annotations, update)
+            return self._client.predict(
+                model_run_id=self.model_run_id,
+                annotation=annotations,
+                update=update,
+            )
 
     def iloc(self, i: int):
         """
@@ -134,7 +139,7 @@ def iloc(self, i: int):
         }
         """
         response = self._client.predictions_iloc(self.model_run_id, i)
-        return self._format_prediction_response(response)
+        return format_prediction_response(response)
 
     def refloc(self, reference_id: str):
         """
@@ -145,7 +150,7 @@ def refloc(self, reference_id: str):
         response = self._client.predictions_ref_id(
             self.model_run_id, reference_id
         )
-        return self._format_prediction_response(response)
+        return format_prediction_response(response)
 
     def loc(self, dataset_item_id: str):
         """
@@ -159,7 +164,7 @@ def loc(self, dataset_item_id: str):
         response = self._client.predictions_loc(
             self.model_run_id, dataset_item_id
         )
-        return self._format_prediction_response(response)
+        return format_prediction_response(response)
 
     def prediction_loc(self, reference_id: str, annotation_id: str):
         """
@@ -184,46 +189,4 @@ def ungrouped_export(self):
             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[
-        dict,
-        List[
-            Union[
-                BoxPrediction,
-                PolygonPrediction,
-                CuboidPrediction,
-                SegmentationPrediction,
-            ]
-        ],
-    ]:
-        annotation_payload = response.get(ANNOTATIONS_KEY, None)
-        if not annotation_payload:
-            # An error occurred
-            return response
-        annotation_response = {}
-        type_key_to_class: Dict[
-            str,
-            Union[
-                Type[BoxPrediction],
-                Type[PolygonPrediction],
-                Type[CuboidPrediction],
-                Type[SegmentationPrediction],
-            ],
-        ] = {
-            BOX_TYPE: BoxPrediction,
-            POLYGON_TYPE: PolygonPrediction,
-            CUBOID_TYPE: CuboidPrediction,
-            SEGMENTATION_TYPE: SegmentationPrediction,
-        }
-        for type_key in annotation_payload:
-            type_class = type_key_to_class[type_key]
-            annotation_response[type_key] = [
-                type_class.from_json(annotation)
-                for annotation in annotation_payload[type_key]
-            ]
-        return annotation_response
+        return format_prediction_response({ANNOTATIONS_KEY: json_response})

nucleus/utils.py

@@ -4,7 +4,7 @@
 import io
 import uuid
 import json
-from typing import IO, Dict, List, Sequence, Union
+from typing import IO, Dict, List, Sequence, Type, Union
 
 import requests
 from requests.models import HTTPError
@@ -30,10 +30,56 @@
     SEGMENTATION_TYPE,
 )
 from .dataset_item import DatasetItem
-from .prediction import BoxPrediction, CuboidPrediction, PolygonPrediction
+from .prediction import (
+    BoxPrediction,
+    CuboidPrediction,
+    PolygonPrediction,
+    SegmentationPrediction,
+)
 from .scene import LidarScene
 
 
+def format_prediction_response(
+    response: dict,
+) -> Union[
+    dict,
+    List[
+        Union[
+            BoxPrediction,
+            PolygonPrediction,
+            CuboidPrediction,
+            SegmentationPrediction,
+        ]
+    ],
+]:
+    annotation_payload = response.get(ANNOTATIONS_KEY, None)
+    if not annotation_payload:
+        # An error occurred
+        return response
+    annotation_response = {}
+    type_key_to_class: Dict[
+        str,
+        Union[
+            Type[BoxPrediction],
+            Type[PolygonPrediction],
+            Type[CuboidPrediction],
+            Type[SegmentationPrediction],
+        ],
+    ] = {
+        BOX_TYPE: BoxPrediction,
+        POLYGON_TYPE: PolygonPrediction,
+        CUBOID_TYPE: CuboidPrediction,
+        SEGMENTATION_TYPE: SegmentationPrediction,
+    }
+    for type_key in annotation_payload:
+        type_class = type_key_to_class[type_key]
+        annotation_response[type_key] = [
+            type_class.from_json(annotation)
+            for annotation in annotation_payload[type_key]
+        ]
+    return annotation_response
+
+
 def _get_all_field_values(metadata_list: List[dict], key: str):
     return {metadata[key] for metadata in metadata_list if key in metadata}