Merge pull request #131 from scaleapi/da-model-only

Model-only endpoints

Author: ardila

nucleus/__init__.py

@@ -734,7 +734,6 @@ def create_model_run(self, dataset_id: str, payload: dict) -> ModelRun:
 
     def predict(
         self,
-        model_run_id: str,
         annotations: List[
             Union[
                 BoxPrediction,
@@ -743,7 +742,10 @@ def predict(
                 SegmentationPrediction,
             ]
         ],
-        update: bool,
+        model_run_id: Optional[str] = None,
+        model_id: Optional[str] = None,
+        dataset_id: Optional[str] = None,
+        update: bool = False,
         batch_size: int = 5000,
     ):
         """
@@ -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
@@ -780,11 +792,9 @@ def predict(
             for i in range(0, len(other_predictions), batch_size)
         ]
 
-        agg_response = {
-            MODEL_RUN_ID_KEY: model_run_id,
-            PREDICTIONS_PROCESSED_KEY: 0,
-            PREDICTIONS_IGNORED_KEY: 0,
-        }
+        errors = []
+        predictions_processed = 0
+        predictions_ignored = 0
 
         tqdm_batches = self.tqdm_bar(batches)
 
@@ -793,36 +803,29 @@ 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
+                errors.append(response)
             else:
-                agg_response[PREDICTIONS_PROCESSED_KEY] += response[
-                    PREDICTIONS_PROCESSED_KEY
-                ]
-                agg_response[PREDICTIONS_IGNORED_KEY] += response[
-                    PREDICTIONS_IGNORED_KEY
-                ]
+                predictions_processed += response[PREDICTIONS_PROCESSED_KEY]
+                predictions_ignored += response[PREDICTIONS_IGNORED_KEY]
 
         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
+                errors.append(response)
             else:
-                agg_response[PREDICTIONS_PROCESSED_KEY] += response[
-                    PREDICTIONS_PROCESSED_KEY
-                ]
-                agg_response[PREDICTIONS_IGNORED_KEY] += response[
-                    PREDICTIONS_IGNORED_KEY
-                ]
+                predictions_processed += response[PREDICTIONS_PROCESSED_KEY]
+                predictions_ignored += response[PREDICTIONS_IGNORED_KEY]
 
-        return agg_response
+        return {
+            MODEL_RUN_ID_KEY: model_run_id,
+            PREDICTIONS_PROCESSED_KEY: predictions_processed,
+            PREDICTIONS_IGNORED_KEY: predictions_ignored,
+            ERRORS_KEY: errors,
+        }
 
     def commit_model_run(
         self, model_run_id: str, payload: Optional[dict] = None

nucleus/dataset.py

@@ -3,18 +3,25 @@
 import requests
 
 from nucleus.job import AsyncJob
+from nucleus.prediction import (
+    BoxPrediction,
+    CuboidPrediction,
+    PolygonPrediction,
+    SegmentationPrediction,
+    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 +31,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 +531,137 @@ def get_scene(self, reference_id) -> Scene:
                 requests_command=requests.get,
             )
         )
+
+    def export_predictions(self, model):
+        """Exports all predictions from a model on this dataset"""
+        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):
+        """
+
+        :param model: the model to calculate eval metrics for
+        :param options: Dict with keys:
+            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 is None:
+            options = {}
+        return self._client.make_request(
+            payload=options,
+            route=f"dataset/{self.id}/model/{model.id}/calculateEvaluationMetrics",
+        )
+
+    def upload_predictions(
+        self,
+        model,
+        predictions: List[
+            Union[
+                BoxPrediction,
+                PolygonPrediction,
+                CuboidPrediction,
+                SegmentationPrediction,
+            ]
+        ],
+        update=False,
+        asynchronous=False,
+    ):
+        """
+        Uploads model outputs as predictions for a model_run. Returns info about the upload.
+        :param predictions: List of prediction objects to ingest
+        :param update: Whether to update (if true) or ignore (if false) on conflicting reference_id/annotation_id
+        :param asynchronous: If true, return launch and then return a reference to an asynchronous job object. This is recommended for large ingests.
+        :return:
+        If synchronoius
+        {
+            "model_run_id": str,
+            "predictions_processed": int,
+            "predictions_ignored": int,
+        }
+        """
+        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):
+        """
+        Returns predictions For Dataset Item by index.
+        :param model: model object to get predictions from.
+        :param index: absolute number of Dataset Item for a dataset corresponding to the model run.
+        :return: List[Union[BoxPrediction, PolygonPrediction, CuboidPrediction, SegmentationPrediction]],
+        }
+        """
+        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):
+        """
+        Returns predictions for dataset Item by its reference_id.
+        :param model: model object to get predictions from.
+        :param reference_id: reference_id of a dataset item.
+        :return: List[Union[BoxPrediction, PolygonPrediction, CuboidPrediction, SegmentationPrediction]],
+        """
+        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):
+        """
+        Returns info for single Prediction by its reference id and annotation id. Not supported for segmentation predictions yet.
+        :param reference_id: the user specified id for the image
+        :param annotation_id: the user specified id for the prediction, or if one was not provided, the Scale internally generated id for the prediction
+        :return:
+         BoxPrediction | PolygonPrediction | CuboidPrediction
+        """
+        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,
+            )
+        )