Merge branch 'master' into support_cuboid_annotations

Author: drew-kaul

nucleus/__init__.py

@@ -181,11 +181,11 @@ def list_models(self) -> List[Model]:
 
         return [
             Model(
-                model["id"],
-                model["name"],
-                model["ref_id"],
-                model["metadata"],
-                self,
+                model_id=model["id"],
+                name=model["name"],
+                reference_id=model["ref_id"],
+                metadata=model["metadata"] or None,
+                client=self,
             )
             for model in model_objects["models"]
         ]
@@ -231,6 +231,19 @@ def get_dataset(self, dataset_id: str) -> Dataset:
         """
         return Dataset(dataset_id, self)
 
+    def get_model(self, model_id: str) -> Model:
+        """
+        Fetched a model for a given id
+        :param model_id: internally controlled dataset_id
+        :return: model
+        """
+        payload = self.make_request(
+            payload={},
+            route=f"model/{model_id}",
+            requests_command=requests.get,
+        )
+        return Model.from_json(payload=payload, client=self)
+
     def get_model_run(self, model_run_id: str, dataset_id: str) -> ModelRun:
         """
         Fetches a model_run for given id

nucleus/annotation.py

@@ -37,7 +37,9 @@ class Annotation:
     item_id: Optional[str] = None
 
     def _check_ids(self):
-        if bool(self.reference_id) == bool(self.item_id):
+        if self.reference_id and self.item_id:
+            self.item_id = None  # Prefer reference id to item id.
+        if not (self.reference_id or self.item_id):
             raise Exception(
                 "You must specify either a reference_id or an item_id for an annotation."
             )
@@ -307,12 +309,12 @@ def to_payload(self) -> dict:
         }
 
 
-def check_all_annotation_paths_remote(
+def check_all_mask_paths_remote(
     annotations: Sequence[Union[Annotation]],
 ):
     for annotation in annotations:
         if hasattr(annotation, MASK_URL_KEY):
             if is_local_path(getattr(annotation, MASK_URL_KEY)):
                 raise ValueError(
-                    f"Found an annotation with a local path, which cannot be uploaded asynchronously. Use a remote path instead. {annotation}"
+                    f"Found an annotation with a local path, which is not currently supported. Use a remote path instead. {annotation}"
                 )

nucleus/dataset.py

@@ -12,7 +12,8 @@
 
 from .annotation import (
     Annotation,
-    check_all_annotation_paths_remote,
+    CuboidAnnotation,
+    check_all_mask_paths_remote,
 )
 from .constants import (
     DATASET_ITEM_IDS_KEY,
@@ -167,9 +168,9 @@ def annotate(
             "ignored_items": int,
         }
         """
-        if asynchronous:
-            check_all_annotation_paths_remote(annotations)
+        check_all_mask_paths_remote(annotations)
 
+        if asynchronous:
             request_id = serialize_and_write_to_presigned_url(
                 annotations, self.id, self._client
             )

nucleus/model.py

@@ -33,11 +33,26 @@ def __repr__(self):
         return f"Model(model_id='{self.id}', name='{self.name}', reference_id='{self.reference_id}', metadata={self.metadata}, client={self._client})"
 
     def __eq__(self, other):
-        return self.id == other.id
+        return (
+            (self.id == other.id)
+            and (self.name == other.name)
+            and (self.metadata == other.metadata)
+            and (self._client == other._client)
+        )
 
     def __hash__(self):
         return hash(self.id)
 
+    @classmethod
+    def from_json(cls, payload: dict, client):
+        return cls(
+            model_id=payload["id"],
+            name=payload["name"],
+            reference_id=payload["ref_id"],
+            metadata=payload["metadata"] or None,
+            client=client,
+        )
+
     def create_run(
         self,
         name: str,

nucleus/model_run.py

@@ -1,6 +1,6 @@
 from typing import Dict, List, Optional, Type, Union
-
-from nucleus.annotation import check_all_annotation_paths_remote
+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
 
@@ -113,7 +113,7 @@ def predict(
         }
         """
         if asynchronous:
-            check_all_annotation_paths_remote(annotations)
+            check_all_mask_paths_remote(annotations)
 
             request_id = serialize_and_write_to_presigned_url(
                 annotations, self._dataset_id, self._client
@@ -162,6 +162,16 @@ def loc(self, dataset_item_id: str):
         )
         return self._format_prediction_response(response)
 
+    def ungrouped_export(self):
+        json_response = self._client.make_request(
+            payload={},
+            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[

pyproject.toml

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

tests/test_models.py

@@ -64,6 +64,10 @@ def test_model_creation_and_listing(CLIENT, dataset):
     # List the models
     ms = CLIENT.list_models()
 
+    # Get a model
+    m = CLIENT.get_model(model.id)
+    assert m == model
+
     assert model in ms
     assert list(set(ms) - set(models_before))[0] == model
 

tests/test_prediction.py

@@ -50,7 +50,7 @@ def test_repr(test_object: any):
 def model_run(CLIENT):
     ds = CLIENT.create_dataset(TEST_DATASET_NAME)
     ds_items = []
-    for url in TEST_IMG_URLS:
+    for url in TEST_IMG_URLS[:2]:
         ds_items.append(
             DatasetItem(
                 image_location=url,
@@ -246,7 +246,7 @@ def test_polygon_pred_upload_ignore(model_run):
     )
 
 
-def test_mixed_pred_upload(model_run):
+def test_mixed_pred_upload(model_run: ModelRun):
     prediction_semseg = SegmentationPrediction.from_json(
         TEST_SEGMENTATION_PREDICTIONS[0]
     )
@@ -262,15 +262,15 @@ def test_mixed_pred_upload(model_run):
     assert response["predictions_processed"] == 3
     assert response["predictions_ignored"] == 0
 
-    response_refloc = model_run.refloc(prediction_polygon.reference_id)
+    all_predictions = model_run.ungrouped_export()
     assert_box_prediction_matches_dict(
-        response_refloc["box"][0], TEST_BOX_PREDICTIONS[0]
+        all_predictions["box"][0], TEST_BOX_PREDICTIONS[0]
     )
     assert_polygon_prediction_matches_dict(
-        response_refloc["polygon"][0], TEST_POLYGON_PREDICTIONS[0]
+        all_predictions["polygon"][0], TEST_POLYGON_PREDICTIONS[0]
     )
     assert_segmentation_annotation_matches_dict(
-        response_refloc["segmentation"][0], TEST_SEGMENTATION_PREDICTIONS[0]
+        all_predictions["segmentation"][0], TEST_SEGMENTATION_PREDICTIONS[0]
     )