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

nucleus/dataset.py

@@ -1,27 +1,30 @@
-from typing import Any, Dict, List, Optional, Union
+from typing import Any, Dict, List, Optional, Sequence, Union
 
 import requests
 
 from nucleus.job import AsyncJob
+from nucleus.url_utils import sanitize_string_args
 from nucleus.utils import (
     convert_export_payload,
     format_dataset_item_response,
     serialize_and_write_to_presigned_url,
 )
 
-from .annotation import Annotation, check_all_annotation_paths_remote
+from .annotation import (
+    Annotation,
+    check_all_mask_paths_remote,
+)
 from .constants import (
-    DATASET_ITEM_IDS_KEY,
     DATASET_LENGTH_KEY,
     DATASET_MODEL_RUNS_KEY,
     DATASET_NAME_KEY,
     DATASET_SLICES_KEY,
     DEFAULT_ANNOTATION_UPDATE_MODE,
     EXPORTED_ROWS,
-    JOB_ID_KEY,
     NAME_KEY,
     REFERENCE_IDS_KEY,
     REQUEST_ID_KEY,
+    AUTOTAG_SCORE_THRESHOLD,
     UPDATE_KEY,
 )
 from .dataset_item import (
@@ -29,10 +32,15 @@ from .dataset_item import (
     check_all_paths_remote,
     check_for_duplicate_reference_ids,
 )
-from .payload_constructor import construct_model_run_creation_payload
-
+from .scene import LidarScene, check_all_scene_paths_remote
+from .payload_constructor import (
+    construct_append_scenes_payload,
+    construct_model_run_creation_payload,
+    construct_taxonomy_payload,
+)
 
 WARN_FOR_LARGE_UPLOAD = 50000
+WARN_FOR_LARGE_SCENES_UPLOAD = 5
 
 
 class Dataset:
@@ -79,21 +87,56 @@ class Dataset:
     def items(self) -> List[DatasetItem]:
         return self._client.get_dataset_items(self.id)
 
-    def autotag_scores(self, autotag_name, for_scores_greater_than=0):
-        """Export the autotag scores above a threshold, largest scores first.
+    @sanitize_string_args
+    def autotag_items(self, autotag_name, for_scores_greater_than=0):
+        """For a given Autotag of this dataset, export its tagged items with scores above a threshold, largest scores first.
 
-        If you have pandas installed, you can create a pandas dataframe using
+        :return: dictionary of the form
+            {
+                'autotagItems': {
+                    ref_id: str,
+                    score: float,
+                    model_prediction_annotation_id: str | None
+                    ground_truth_annotation_id: str | None,
+                }[],
+                'autotag': {
+                    id: str,
+                    name: str,
+                    status: 'started' | 'completed',
+                    autotag_level: 'Image' | 'Object'
+                }
+            }
+        See https://dashboard.nucleus.scale.com/nucleus/docs/api#export-autotag-items for more details on the return types.
+        """
+        response = self._client.make_request(
+            payload={AUTOTAG_SCORE_THRESHOLD: for_scores_greater_than},
+            route=f"dataset/{self.id}/autotag/{autotag_name}/taggedItems",
+            requests_command=requests.get,
+        )
+        return response
 
-        pandas.Dataframe(dataset.autotag_scores(autotag_name))
+    def autotag_training_items(self, autotag_name):
+        """For a given Autotag of this dataset, export its training items. These are user selected positives during refinement.
 
         :return: dictionary of the form
-            {'ref_ids': List[str],
-             'datset_item_ids': List[str],
-             'score': List[float]}
+            {
+                'autotagPositiveTrainingItems': {
+                    ref_id: str,
+                    model_prediction_annotation_id: str | None,
+                    ground_truth_annotation_id: str | None,
+                }[],
+                'autotag': {
+                    id: str,
+                    name: str,
+                    status: 'started' | 'completed',
+                    autotag_level: 'Image' | 'Object'
+                }
+            }
+        See https://dashboard.nucleus.scale.com/nucleus/docs/api#export-autotag-training-items for more details on the return types.
         """
         response = self._client.make_request(
             payload={},
-            route=f"autotag/{self.id}/{autotag_name}/{for_scores_greater_than}",
+            route=f"dataset/{self.id}/autotag/{autotag_name}/trainingItems",
             requests_command=requests.get,
         )
         return response
@@ -146,7 +189,7 @@ class Dataset:
 
     def annotate(
         self,
-        annotations: List[Annotation],
+        annotations: Sequence[Annotation],
         update: Optional[bool] = DEFAULT_ANNOTATION_UPDATE_MODE,
         batch_size: int = 5000,
         asynchronous: bool = False,
@@ -163,9 +206,9 @@ class Dataset:
             "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
             )
@@ -173,9 +216,7 @@ class Dataset:
                 payload={REQUEST_ID_KEY: request_id, UPDATE_KEY: update},
                 route=f"dataset/{self.id}/annotate?async=1",
             )
-
-            return AsyncJob(response[JOB_ID_KEY], self._client)
-
+            return AsyncJob.from_json(response, self._client)
         return self._client.annotate_dataset(
             self.id, annotations, update=update, batch_size=batch_size
         )
@@ -193,16 +234,16 @@ class Dataset:
 
     def append(
         self,
-        dataset_items: List[DatasetItem],
+        items: Union[Sequence[DatasetItem], Sequence[LidarScene]],
         update: Optional[bool] = False,
         batch_size: Optional[int] = 20,
         asynchronous=False,
     ) -> Union[dict, AsyncJob]:
         """
-        Appends images with metadata (dataset items) to the dataset. Overwrites images on collision if forced.
+        Appends images with metadata (dataset items) or scenes to the dataset. Overwrites images on collision if forced.
 
         Parameters:
-        :param dataset_items: items to upload
+        :param items: items to upload
         :param update: if True overwrites images and metadata on collision
         :param batch_size: batch parameter for long uploads
         :param aynchronous: if True, return a job object representing asynchronous ingestion job.
@@ -214,6 +255,17 @@ class Dataset:
             'ignored_items': int,
         }
         """
+        dataset_items = [
+            item for item in items if isinstance(item, DatasetItem)
+        ]
+        scenes = [item for item in items if isinstance(item, LidarScene)]
+        if dataset_items and scenes:
+            raise Exception(
+                "You must append either DatasetItems or Scenes to the dataset."
+            )
+        if scenes:
+            return self.append_scenes(scenes, update, asynchronous)
+
         check_for_duplicate_reference_ids(dataset_items)
 
         if len(dataset_items) > WARN_FOR_LARGE_UPLOAD and not asynchronous:
@@ -233,7 +285,7 @@ class Dataset:
                 payload={REQUEST_ID_KEY: request_id, UPDATE_KEY: update},
                 route=f"dataset/{self.id}/append?async=1",
             )
-            return AsyncJob(response["job_id"], self._client)
+            return AsyncJob.from_json(response, self._client)
 
         return self._client.populate_dataset(
             self.id,
@@ -242,6 +294,57 @@ class Dataset:
             batch_size=batch_size,
         )
 
+    def append_scenes(
+        self,
+        scenes: List[LidarScene],
+        update: Optional[bool] = False,
+        asynchronous: Optional[bool] = False,
+    ) -> Union[dict, AsyncJob]:
+        """
+        Appends scenes with given frames (containing pointclouds and optional images) to the dataset
+
+        Parameters:
+        :param scenes: scenes to upload
+        :param update: if True, overwrite scene on collision
+        :param asynchronous: if True, return a job object representing asynchronous ingestion job
+        :return:
+        {
+            'dataset_id': str,
+            'new_scenes': int,
+            'ignored_scenes': int,
+            'scenes_errored': int,
+            'errors': List[str],
+        }
+        """
+        for scene in scenes:
+            scene.validate()
+
+        if len(scenes) > WARN_FOR_LARGE_SCENES_UPLOAD and not asynchronous:
+            print(
+                "Tip: for large uploads, get faster performance by importing your data "
+                "into Nucleus directly from a cloud storage provider. See "
+                "https://dashboard.scale.com/nucleus/docs/api?language=python#guide-for-large-ingestions"
+                " for details."
+            )
+
+        if asynchronous:
+            check_all_scene_paths_remote(scenes)
+            request_id = serialize_and_write_to_presigned_url(
+                scenes, self.id, self._client
+            )
+            response = self._client.make_request(
+                payload={REQUEST_ID_KEY: request_id, UPDATE_KEY: update},
+                route=f"{self.id}/upload_scenes?async=1",
+            )
+            return AsyncJob.from_json(response, self._client)
+
+        payload = construct_append_scenes_payload(scenes, update)
+        response = self._client.make_request(
+            payload=payload,
+            route=f"{self.id}/upload_scenes",
+        )
+        return response
+
     def iloc(self, i: int) -> dict:
         """
         Returns Dataset Item Info By Dataset Item Number.
@@ -249,7 +352,7 @@ class Dataset:
         :return:
         {
             "item": DatasetItem,
-            "annotations": List[Union[BoxAnnotation, PolygonAnnotation]],
+            "annotations": List[Union[BoxAnnotation, PolygonAnnotation, CuboidAnnotation, SegmentationAnnotation]],
         }
         """
         response = self._client.dataitem_iloc(self.id, i)
@@ -262,7 +365,7 @@ class Dataset:
         :return:
         {
             "item": DatasetItem,
-            "annotations": List[Union[BoxAnnotation, PolygonAnnotation]],
+            "annotations": List[Union[BoxAnnotation, PolygonAnnotation, CuboidAnnotation, SegmentationAnnotation]],
         }
         """
         response = self._client.dataitem_ref_id(self.id, reference_id)
@@ -275,17 +378,31 @@ class Dataset:
         :return:
         {
             "item": DatasetItem,
-            "annotations": List[Union[BoxAnnotation, PolygonAnnotation]],
+            "annotations": List[Union[BoxAnnotation, PolygonAnnotation, CuboidAnnotation, SegmentationAnnotation]],
         }
         """
         response = self._client.dataitem_loc(self.id, dataset_item_id)
         return format_dataset_item_response(response)
 
+    def ground_truth_loc(self, reference_id: str, annotation_id: str):
+        """
+        Returns info for single ground truth Annotation by its id.
+        :param reference_id: User specified id for the dataset item the ground truth is attached to
+        :param annotation_id: User specified, or auto-generated id for the annotation
+        :return:
+        BoxAnnotation | PolygonAnnotation | CuboidAnnotation
+        """
+        response = self._client.make_request(
+            {},
+            f"dataset/{self.id}/groundTruth/loc/{reference_id}/{annotation_id}",
+            requests.get,
+        )
+        return Annotation.from_json(response)
+
     def create_slice(
         self,
         name: str,
-        dataset_item_ids: List[str] = None,
-        reference_ids: List[str] = None,
+        reference_ids: List[str],
     ):
         """
         Creates a slice from items already present in a dataset.
@@ -293,42 +410,60 @@ class Dataset:
         as a means of identifying items in the dataset.
 
         :param name: The human-readable name of the slice.
-        :param dataset_item_ids: An optional list of dataset item ids for the items in the slice
-        :param reference_ids: An optional list of user-specified identifier for the items in the slice
+        :param reference_ids: A list of user-specified identifier for the items in the slice
 
         :return: new Slice object
         """
-        if bool(dataset_item_ids) == bool(reference_ids):
-            raise Exception(
-                "You must specify exactly one of dataset_item_ids or reference_ids."
-            )
-        payload: Dict[str, Any] = {NAME_KEY: name}
-        if dataset_item_ids:
-            payload[DATASET_ITEM_IDS_KEY] = dataset_item_ids
-        if reference_ids:
-            payload[REFERENCE_IDS_KEY] = reference_ids
-        return self._client.create_slice(self.id, payload)
-
-    def delete_item(self, item_id: str = None, reference_id: str = None):
-        if bool(item_id) == bool(reference_id):
-            raise Exception(
-                "You must specify either a reference_id or an item_id for a DatasetItem."
-            )
+        return self._client.create_slice(
+            self.id, {NAME_KEY: name, REFERENCE_IDS_KEY: reference_ids}
+        )
+
+    def delete_item(self, reference_id: str):
         return self._client.delete_dataset_item(
-            self.id, reference_id=reference_id, item_id=item_id
+            self.id, reference_id=reference_id
         )
 
     def list_autotags(self):
         return self._client.list_autotags(self.id)
 
-    def create_custom_index(self, embeddings_url: str):
-        return self._client.create_custom_index(self.id, embeddings_url)
+    def create_custom_index(self, embeddings_urls: list, embedding_dim: int):
+        return AsyncJob.from_json(
+            self._client.create_custom_index(
+                self.id,
+                embeddings_urls,
+                embedding_dim,
+            ),
+            self._client,
+        )
 
     def delete_custom_index(self):
         return self._client.delete_custom_index(self.id)
 
-    def check_index_status(self, job_id: str):
-        return self._client.check_index_status(job_id)
+    def set_continuous_indexing(self, enable: bool = True):
+        return self._client.set_continuous_indexing(self.id, enable)
+
+    def create_image_index(self):
+        response = self._client.create_image_index(self.id)
+        return AsyncJob.from_json(response, self._client)
+
+    def add_taxonomy(
+        self,
+        taxonomy_name: str,
+        taxonomy_type: str,
+        labels: List[str],
+    ):
+        """
+        Creates a new taxonomy.
+        Returns a response with dataset_id, taxonomy_name and type for the new taxonomy.
+        :param taxonomy_name: name of the taxonomy
+        :param type: type of the taxonomy
+        :param labels: list of possible labels for the taxonomy
+        """
+        return self._client.make_request(
+            construct_taxonomy_payload(taxonomy_name, taxonomy_type, labels),
+            f"dataset/{self.id}/add_taxonomy",
+            requests_command=requests.post,
+        )
 
     def items_and_annotations(
         self,
@@ -349,3 +484,29 @@ class Dataset:
             requests_command=requests.get,
         )
         return convert_export_payload(api_payload[EXPORTED_ROWS])
+
+    def export_embeddings(
+        self,
+    ) -> List[Dict[str, Union[str, List[float]]]]:
+        """Returns a pd.Dataframe-ready format of dataset embeddings.
+
+        Returns:
+            A list, where each item is a dict with two keys representing a row
+            in the dataset.
+            * One value in the dict is the reference id
+            * The other value is a list of the embedding values
+        """
+        api_payload = self._client.make_request(
+            payload=None,
+            route=f"dataset/{self.id}/embeddings",
+            requests_command=requests.get,
+        )
+        return api_payload
+
+    def delete_annotations(
+        self, reference_ids: list = None, keep_history=False
+    ):
+        response = self._client.delete_annotations(
+            self.id, reference_ids, keep_history
+        )
+        return AsyncJob.from_json(response, self._client)

nucleus/dataset_item.py

@@ -2,66 +2,180 @@ from collections import Counter
 import json
 import os.path
 from dataclasses import dataclass
-from typing import Optional, Sequence
-from urllib.parse import urlparse
+from typing import Optional, Sequence, Dict, Any
+from enum import Enum
 
+from .annotation import is_local_path, Point3D
 from .constants import (
-    DATASET_ITEM_ID_KEY,
     IMAGE_URL_KEY,
     METADATA_KEY,
     ORIGINAL_IMAGE_URL_KEY,
+    UPLOAD_TO_SCALE_KEY,
     REFERENCE_ID_KEY,
+    TYPE_KEY,
+    URL_KEY,
+    CAMERA_PARAMS_KEY,
+    POINTCLOUD_URL_KEY,
+    X_KEY,
+    Y_KEY,
+    Z_KEY,
+    W_KEY,
+    POSITION_KEY,
+    HEADING_KEY,
+    FX_KEY,
+    FY_KEY,
+    CX_KEY,
+    CY_KEY,
 )
 
 
 @dataclass
-class DatasetItem:
+class Quaternion:
+    x: float
+    y: float
+    z: float
+    w: float
 
-    image_location: str
+    @classmethod
+    def from_json(cls, payload: Dict[str, float]):
+        return cls(
+            payload[X_KEY], payload[Y_KEY], payload[Z_KEY], payload[W_KEY]
+        )
+
+    def to_payload(self) -> dict:
+        return {
+            X_KEY: self.x,
+            Y_KEY: self.y,
+            Z_KEY: self.z,
+            W_KEY: self.w,
+        }
+
+
+@dataclass
+class CameraParams:
+    position: Point3D
+    heading: Quaternion
+    fx: float
+    fy: float
+    cx: float
+    cy: float
+
+    @classmethod
+    def from_json(cls, payload: Dict[str, Any]):
+        return cls(
+            Point3D.from_json(payload[POSITION_KEY]),
+            Quaternion.from_json(payload[HEADING_KEY]),
+            payload[FX_KEY],
+            payload[FY_KEY],
+            payload[CX_KEY],
+            payload[CY_KEY],
+        )
+
+    def to_payload(self) -> dict:
+        return {
+            POSITION_KEY: self.position.to_payload(),
+            HEADING_KEY: self.heading.to_payload(),
+            FX_KEY: self.fx,
+            FY_KEY: self.fy,
+            CX_KEY: self.cx,
+            CY_KEY: self.cy,
+        }
+
+
+class DatasetItemType(Enum):
+    IMAGE = "image"
+    POINTCLOUD = "pointcloud"
+
+
+@dataclass  # pylint: disable=R0902
+class DatasetItem:  # pylint: disable=R0902
+    image_location: Optional[str] = None
     reference_id: Optional[str] = None
-    item_id: Optional[str] = None
     metadata: Optional[dict] = None
+    pointcloud_location: Optional[str] = None
+    upload_to_scale: Optional[bool] = True
 
     def __post_init__(self):
-        self.local = is_local_path(self.image_location)
+        assert self.reference_id is not None, "reference_id is required."
+        assert bool(self.image_location) != bool(
+            self.pointcloud_location
+        ), "Must specify exactly one of the image_location, pointcloud_location parameters"
+        if self.pointcloud_location and not self.upload_to_scale:
+            raise NotImplementedError(
+                "Skipping upload to Scale is not currently implemented for pointclouds."
+            )
+        self.local = (
+            is_local_path(self.image_location) if self.image_location else None
+        )
+        self.type = (
+            DatasetItemType.IMAGE
+            if self.image_location
+            else DatasetItemType.POINTCLOUD
+        )
+        camera_params = (
+            self.metadata.get(CAMERA_PARAMS_KEY, None)
+            if self.metadata
+            else None
+        )
+        self.camera_params = (
+            CameraParams.from_json(camera_params) if camera_params else None
+        )
 
     @classmethod
-    def from_json(cls, payload: dict):
-        url = payload.get(IMAGE_URL_KEY, "") or payload.get(
-            ORIGINAL_IMAGE_URL_KEY, ""
+    def from_json(cls, payload: dict, is_scene=False):
+        image_url = payload.get(IMAGE_URL_KEY, None) or payload.get(
+            ORIGINAL_IMAGE_URL_KEY, None
         )
+
+        if is_scene:
+            return cls(
+                image_location=image_url,
+                pointcloud_location=payload.get(POINTCLOUD_URL_KEY, None),
+                reference_id=payload.get(REFERENCE_ID_KEY, None),
+                metadata=payload.get(METADATA_KEY, {}),
+            )
+
         return cls(
-            image_location=url,
+            image_location=image_url,
             reference_id=payload.get(REFERENCE_ID_KEY, None),
-            item_id=payload.get(DATASET_ITEM_ID_KEY, None),
             metadata=payload.get(METADATA_KEY, {}),
+            upload_to_scale=payload.get(UPLOAD_TO_SCALE_KEY, None),
         )
 
     def local_file_exists(self):
         return os.path.isfile(self.image_location)
 
-    def to_payload(self) -> dict:
-        payload = {
-            IMAGE_URL_KEY: self.image_location,
+    def to_payload(self, is_scene=False) -> dict:
+        payload: Dict[str, Any] = {
             METADATA_KEY: self.metadata or {},
         }
-        if self.reference_id:
-            payload[REFERENCE_ID_KEY] = self.reference_id
-        if self.item_id:
-            payload[DATASET_ITEM_ID_KEY] = self.item_id
+
+        payload[REFERENCE_ID_KEY] = self.reference_id
+
+        if is_scene:
+            if self.image_location:
+                payload[URL_KEY] = self.image_location
+            elif self.pointcloud_location:
+                payload[URL_KEY] = self.pointcloud_location
+            payload[TYPE_KEY] = self.type.value
+            if self.camera_params:
+                payload[CAMERA_PARAMS_KEY] = self.camera_params.to_payload()
+        else:
+            assert (
+                self.image_location
+            ), "Must specify image_location for DatasetItems not in a LidarScene"
+            payload[IMAGE_URL_KEY] = self.image_location
+            payload[UPLOAD_TO_SCALE_KEY] = self.upload_to_scale
+
         return payload
 
     def to_json(self) -> str:
         return json.dumps(self.to_payload(), allow_nan=False)
 
 
-def is_local_path(path: str) -> bool:
-    return urlparse(path).scheme not in {"https", "http", "s3", "gs"}
-
-
 def check_all_paths_remote(dataset_items: Sequence[DatasetItem]):
     for item in dataset_items:
-        if is_local_path(item.image_location):
+        if item.image_location and is_local_path(item.image_location):
             raise ValueError(
                 f"All paths must be remote, but {item.image_location} is either "
                 "local, or a remote URL type that is not supported."
@@ -79,6 +193,5 @@ def check_for_duplicate_reference_ids(dataset_items: Sequence[DatasetItem]):
             for key, value in Counter(ref_ids).items()
         }
         raise ValueError(
-            "Duplicate reference ids found among dataset_items: %s"
-            % duplicates
+            f"Duplicate reference ids found among dataset_items: {duplicates}"
         )

nucleus/errors.py

@@ -1,3 +1,15 @@
+import pkg_resources
+
+nucleus_client_version = pkg_resources.get_distribution(
+    "scale-nucleus"
+).version
+
+INFRA_FLAKE_MESSAGES = [
+    "downstream duration timeout",
+    "upstream connect error or disconnect/reset before headers. reset reason: local reset",
+]
+
+
 class ModelCreationError(Exception):
     def __init__(self, message="Could not create the model"):
         self.message = message
@@ -28,9 +40,9 @@ class NucleusAPIError(Exception):
     def __init__(
         self, endpoint, command, requests_response=None, aiohttp_response=None
     ):
-
+        message = f"Your client is on version {nucleus_client_version}. If you have not recently done so, please make sure you have updated to the latest version of the client by running pip install --upgrade scale-nucleus\n"
         if requests_response is not None:
-            message = f"Tried to {command.__name__} {endpoint}, but received {requests_response.status_code}: {requests_response.reason}."
+            message += f"Tried to {command.__name__} {endpoint}, but received {requests_response.status_code}: {requests_response.reason}."
             if hasattr(requests_response, "text"):
                 if requests_response.text:
                     message += (
@@ -39,8 +51,14 @@ class NucleusAPIError(Exception):
 
         if aiohttp_response is not None:
             status, reason, data = aiohttp_response
-            message = f"Tried to {command.__name__} {endpoint}, but received {status}: {reason}."
+            message += f"Tried to {command.__name__} {endpoint}, but received {status}: {reason}."
             if data:
                 message += f"\nThe detailed error is:\n{data}"
 
+        if any(
+            infra_flake_message in message
+            for infra_flake_message in INFRA_FLAKE_MESSAGES
+        ):
+            message += "\n This likely indicates temporary downtime of the API, please try again in a minute or two"
+
         super().__init__(message)