scale-nucleus 0.1.10__py3-none-any.whl → 0.1.24__py3-none-any.whl

nucleus/job.py

@@ -1,47 +1,68 @@
 from dataclasses import dataclass
 import time
 from typing import Dict, List
-
 import requests
+from nucleus.constants import (
+    JOB_CREATION_TIME_KEY,
+    JOB_ID_KEY,
+    JOB_LAST_KNOWN_STATUS_KEY,
+    JOB_TYPE_KEY,
+    STATUS_KEY,
+)
 
 JOB_POLLING_INTERVAL = 5
 
 
 @dataclass
 class AsyncJob:
-    id: str
+    job_id: str
+    job_last_known_status: str
+    job_type: str
+    job_creation_time: str
     client: "NucleusClient"  # type: ignore # noqa: F821
 
     def status(self) -> Dict[str, str]:
-        return self.client.make_request(
+        response = self.client.make_request(
             payload={},
-            route=f"job/{self.id}",
+            route=f"job/{self.job_id}",
             requests_command=requests.get,
         )
+        self.job_last_known_status = response[STATUS_KEY]
+        return response
 
     def errors(self) -> List[str]:
         return self.client.make_request(
             payload={},
-            route=f"job/{self.id}/errors",
+            route=f"job/{self.job_id}/errors",
             requests_command=requests.get,
         )
 
     def sleep_until_complete(self, verbose_std_out=True):
         while 1:
             status = self.status()
-
             time.sleep(JOB_POLLING_INTERVAL)
 
             if verbose_std_out:
                 print(f"Status at {time.ctime()}: {status}")
             if status["status"] == "Running":
                 continue
+
             break
 
         final_status = status
         if final_status["status"] == "Errored":
             raise JobError(final_status, self)
 
+    @classmethod
+    def from_json(cls, payload: dict, client):
+        return cls(
+            job_id=payload[JOB_ID_KEY],
+            job_last_known_status=payload[JOB_LAST_KNOWN_STATUS_KEY],
+            job_type=payload[JOB_TYPE_KEY],
+            job_creation_time=payload[JOB_CREATION_TIME_KEY],
+            client=client,
+        )
+
 
 class JobError(Exception):
     def __init__(self, job_status: Dict[str, str], job: AsyncJob):

nucleus/model.py

@@ -2,6 +2,7 @@ from typing import List, Optional, Dict, Union
 from .dataset import Dataset
 from .prediction import (
     BoxPrediction,
+    CuboidPrediction,
     PolygonPrediction,
     SegmentationPrediction,
 )
@@ -32,17 +33,37 @@ class Model:
         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,
         dataset: Dataset,
         predictions: List[
-            Union[BoxPrediction, PolygonPrediction, SegmentationPrediction]
+            Union[
+                BoxPrediction,
+                PolygonPrediction,
+                CuboidPrediction,
+                SegmentationPrediction,
+            ]
         ],
         metadata: Optional[Dict] = None,
         asynchronous: bool = False,

nucleus/model_run.py

@@ -1,14 +1,14 @@
 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
 
 from .constants import (
     ANNOTATIONS_KEY,
     BOX_TYPE,
+    CUBOID_TYPE,
     DEFAULT_ANNOTATION_UPDATE_MODE,
-    JOB_ID_KEY,
     POLYGON_TYPE,
     REQUEST_ID_KEY,
     SEGMENTATION_TYPE,
@@ -16,8 +16,10 @@ from .constants import (
 )
 from .prediction import (
     BoxPrediction,
+    CuboidPrediction,
     PolygonPrediction,
     SegmentationPrediction,
+    from_json,
 )
 
 
@@ -30,10 +32,10 @@ class ModelRun:
     def __init__(self, model_run_id: str, dataset_id: str, client):
         self.model_run_id = model_run_id
         self._client = client
-        self._dataset_id = dataset_id
+        self.dataset_id = dataset_id
 
     def __repr__(self):
-        return f"ModelRun(model_run_id='{self.model_run_id}', dataset_id='{self._dataset_id}', client={self._client})"
+        return f"ModelRun(model_run_id='{self.model_run_id}', dataset_id='{self.dataset_id}', client={self._client})"
 
     def __eq__(self, other):
         if self.model_run_id == other.model_run_id:
@@ -90,14 +92,19 @@ class ModelRun:
     def predict(
         self,
         annotations: List[
-            Union[BoxPrediction, PolygonPrediction, SegmentationPrediction]
+            Union[
+                BoxPrediction,
+                PolygonPrediction,
+                CuboidPrediction,
+                SegmentationPrediction,
+            ]
         ],
         update: Optional[bool] = DEFAULT_ANNOTATION_UPDATE_MODE,
         asynchronous: bool = False,
     ) -> Union[dict, AsyncJob]:
         """
         Uploads model outputs as predictions for a model_run. Returns info about the upload.
-        :param annotations: List[Union[BoxPrediction, PolygonPrediction]],
+        :param annotations: List[Union[BoxPrediction, PolygonPrediction, CuboidPrediction, SegmentationPrediction]],
         :return:
         {
             "model_run_id": str,
@@ -106,17 +113,16 @@ class ModelRun:
         }
         """
         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
+                annotations, self.dataset_id, self._client
             )
             response = self._client.make_request(
                 payload={REQUEST_ID_KEY: request_id, UPDATE_KEY: update},
                 route=f"modelRun/{self.model_run_id}/predict?async=1",
             )
-
-            return AsyncJob(response[JOB_ID_KEY], self._client)
+            return AsyncJob.from_json(response, self._client)
         else:
             return self._client.predict(self.model_run_id, annotations, update)
 
@@ -124,7 +130,7 @@ class ModelRun:
         """
         Returns Model Run Info For Dataset Item by its number.
         :param i: absolute number of Dataset Item for a dataset corresponding to the model run.
-        :return: List[Union[BoxPrediction, PolygonPrediction]],
+        :return: List[Union[BoxPrediction, PolygonPrediction, CuboidPrediction, SegmentationPrediction]],
         }
         """
         response = self._client.predictions_iloc(self.model_run_id, i)
@@ -134,7 +140,7 @@ class ModelRun:
         """
         Returns Model Run Info For Dataset Item by its reference_id.
         :param reference_id: reference_id of a dataset item.
-        :return: List[Union[BoxPrediction, PolygonPrediction]],
+        :return: List[Union[BoxPrediction, PolygonPrediction, CuboidPrediction, SegmentationPrediction]],
         """
         response = self._client.predictions_ref_id(
             self.model_run_id, reference_id
@@ -155,11 +161,45 @@ class ModelRun:
         )
         return self._format_prediction_response(response)
 
+    def prediction_loc(self, reference_id: str, annotation_id: str):
+        """
+        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, or if one was not provided, the Scale internally generated id for the prediction
+        :return:
+         BoxPrediction | PolygonPrediction | CuboidPrediction
+        """
+
+        response = self._client.make_request(
+            {},
+            f"modelRun/{self.model_run_id}/prediction/loc/{reference_id}/{annotation_id}",
+            requests.get,
+        )
+
+        return from_json(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[
         dict,
-        List[Union[BoxPrediction, PolygonPrediction, SegmentationPrediction]],
+        List[
+            Union[
+                BoxPrediction,
+                PolygonPrediction,
+                CuboidPrediction,
+                SegmentationPrediction,
+            ]
+        ],
     ]:
         annotation_payload = response.get(ANNOTATIONS_KEY, None)
         if not annotation_payload:
@@ -171,11 +211,13 @@ class ModelRun:
             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:

nucleus/payload_constructor.py

@@ -1,12 +1,16 @@
 from typing import List, Optional, Dict, Union
 from .dataset_item import DatasetItem
+from .scene import LidarScene
 from .annotation import (
     BoxAnnotation,
+    CuboidAnnotation,
     PolygonAnnotation,
+    CategoryAnnotation,
     SegmentationAnnotation,
 )
 from .prediction import (
     BoxPrediction,
+    CuboidPrediction,
     PolygonPrediction,
     SegmentationPrediction,
 )
@@ -17,10 +21,14 @@ from .constants import (
     REFERENCE_ID_KEY,
     ANNOTATIONS_KEY,
     ITEMS_KEY,
+    SCENES_KEY,
     UPDATE_KEY,
     MODEL_ID_KEY,
     ANNOTATION_METADATA_SCHEMA_KEY,
     SEGMENTATIONS_KEY,
+    TAXONOMY_NAME_KEY,
+    TYPE_KEY,
+    LABELS_KEY,
 )
 
 
@@ -38,8 +46,25 @@ def construct_append_payload(
     )
 
 
+def construct_append_scenes_payload(
+    scene_list: List[LidarScene], update: Optional[bool] = False
+) -> dict:
+    scenes = []
+    for scene in scene_list:
+        scenes.append(scene.to_payload())
+    return {SCENES_KEY: scenes, UPDATE_KEY: update}
+
+
 def construct_annotation_payload(
-    annotation_items: List[Union[BoxAnnotation, PolygonAnnotation]],
+    annotation_items: List[
+        Union[
+            BoxAnnotation,
+            PolygonAnnotation,
+            CuboidAnnotation,
+            CategoryAnnotation,
+            SegmentationAnnotation,
+        ]
+    ],
     update: bool,
 ) -> dict:
     annotations = []
@@ -63,7 +88,9 @@ def construct_segmentation_payload(
 
 
 def construct_box_predictions_payload(
-    box_predictions: List[Union[BoxPrediction, PolygonPrediction]],
+    box_predictions: List[
+        Union[BoxPrediction, PolygonPrediction, CuboidPrediction]
+    ],
     update: bool,
 ) -> dict:
     predictions = []
@@ -104,3 +131,13 @@ def construct_model_run_creation_payload(
         METADATA_KEY: metadata if metadata else {},
         ANNOTATION_METADATA_SCHEMA_KEY: annotation_metadata_schema,
     }
+
+
+def construct_taxonomy_payload(
+    taxonomy_name: str, taxonomy_type: str, labels: List[str]
+) -> dict:
+    return {
+        TAXONOMY_NAME_KEY: taxonomy_name,
+        TYPE_KEY: taxonomy_type,
+        LABELS_KEY: labels,
+    }

nucleus/prediction.py

@@ -5,14 +5,19 @@ from .annotation import (
     PolygonAnnotation,
     Segment,
     SegmentationAnnotation,
+    CuboidAnnotation,
+    Point3D,
 )
 from .constants import (
     ANNOTATION_ID_KEY,
-    DATASET_ITEM_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,
@@ -21,11 +26,24 @@ from .constants import (
     CONFIDENCE_KEY,
     VERTICES_KEY,
     ANNOTATIONS_KEY,
-    ITEM_ID_KEY,
     MASK_URL_KEY,
+    POSITION_KEY,
+    DIMENSIONS_KEY,
+    YAW_KEY,
 )
 
 
+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
@@ -37,8 +55,7 @@ class SegmentationPrediction(SegmentationAnnotation):
                 Segment.from_json(ann)
                 for ann in payload.get(ANNOTATIONS_KEY, [])
             ],
-            reference_id=payload.get(REFERENCE_ID_KEY, None),
-            item_id=payload.get(ITEM_ID_KEY, None),
+            reference_id=payload[REFERENCE_ID_KEY],
             annotation_id=payload.get(ANNOTATION_ID_KEY, None),
         )
 
@@ -51,8 +68,7 @@ class BoxPrediction(BoxAnnotation):
         y: int,
         width: int,
         height: int,
-        reference_id: Optional[str] = None,
-        item_id: Optional[str] = None,
+        reference_id: str,
         confidence: Optional[float] = None,
         annotation_id: Optional[str] = None,
         metadata: Optional[Dict] = None,
@@ -65,7 +81,6 @@ class BoxPrediction(BoxAnnotation):
             width=width,
             height=height,
             reference_id=reference_id,
-            item_id=item_id,
             annotation_id=annotation_id,
             metadata=metadata,
         )
@@ -90,8 +105,7 @@ class BoxPrediction(BoxAnnotation):
             y=geometry.get(Y_KEY, 0),
             width=geometry.get(WIDTH_KEY, 0),
             height=geometry.get(HEIGHT_KEY, 0),
-            reference_id=payload.get(REFERENCE_ID_KEY, None),
-            item_id=payload.get(DATASET_ITEM_ID_KEY, None),
+            reference_id=payload[REFERENCE_ID_KEY],
             confidence=payload.get(CONFIDENCE_KEY, None),
             annotation_id=payload.get(ANNOTATION_ID_KEY, None),
             metadata=payload.get(METADATA_KEY, {}),
@@ -104,8 +118,7 @@ class PolygonPrediction(PolygonAnnotation):
         self,
         label: str,
         vertices: List[Point],
-        reference_id: Optional[str] = None,
-        item_id: Optional[str] = None,
+        reference_id: str,
         confidence: Optional[float] = None,
         annotation_id: Optional[str] = None,
         metadata: Optional[Dict] = None,
@@ -115,7 +128,6 @@ class PolygonPrediction(PolygonAnnotation):
             label=label,
             vertices=vertices,
             reference_id=reference_id,
-            item_id=item_id,
             annotation_id=annotation_id,
             metadata=metadata,
         )
@@ -139,8 +151,57 @@ class PolygonPrediction(PolygonAnnotation):
             vertices=[
                 Point.from_json(_) for _ in geometry.get(VERTICES_KEY, [])
             ],
-            reference_id=payload.get(REFERENCE_ID_KEY, None),
-            item_id=payload.get(DATASET_ITEM_ID_KEY, None),
+            reference_id=payload[REFERENCE_ID_KEY],
+            confidence=payload.get(CONFIDENCE_KEY, None),
+            annotation_id=payload.get(ANNOTATION_ID_KEY, None),
+            metadata=payload.get(METADATA_KEY, {}),
+            class_pdf=payload.get(CLASS_PDF_KEY, None),
+        )
+
+
+class CuboidPrediction(CuboidAnnotation):
+    def __init__(
+        self,
+        label: str,
+        position: Point3D,
+        dimensions: Point3D,
+        yaw: float,
+        reference_id: str,
+        confidence: Optional[float] = None,
+        annotation_id: Optional[str] = None,
+        metadata: Optional[Dict] = None,
+        class_pdf: Optional[Dict] = None,
+    ):
+        super().__init__(
+            label=label,
+            position=position,
+            dimensions=dimensions,
+            yaw=yaw,
+            reference_id=reference_id,
+            annotation_id=annotation_id,
+            metadata=metadata,
+        )
+        self.confidence = confidence
+        self.class_pdf = class_pdf
+
+    def to_payload(self) -> dict:
+        payload = super().to_payload()
+        if self.confidence is not None:
+            payload[CONFIDENCE_KEY] = self.confidence
+        if self.class_pdf is not None:
+            payload[CLASS_PDF_KEY] = self.class_pdf
+
+        return payload
+
+    @classmethod
+    def from_json(cls, payload: dict):
+        geometry = payload.get(GEOMETRY_KEY, {})
+        return cls(
+            label=payload.get(LABEL_KEY, 0),
+            position=Point3D.from_json(geometry.get(POSITION_KEY, {})),
+            dimensions=Point3D.from_json(geometry.get(DIMENSIONS_KEY, {})),
+            yaw=geometry.get(YAW_KEY, 0),
+            reference_id=payload[REFERENCE_ID_KEY],
             confidence=payload.get(CONFIDENCE_KEY, None),
             annotation_id=payload.get(ANNOTATION_ID_KEY, None),
             metadata=payload.get(METADATA_KEY, {}),