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

nucleus/__init__.py

@@ -1,73 +1,33 @@
 """
 Nucleus Python Library.
 
-Data formats used:
-
-_____________________________________________________________________________________________________
-
-DatasetItem
-
-image_url       |   str     |   The URL containing the image for the given row of data.\n
-reference_id    |   str     |   An optional user-specified identifier to reference this given image.\n
-metadata        |   dict    |   All of column definitions for this item.
-                |           |   The keys should match the user-specified column names,
-                |           |   and the corresponding values will populate the cell under the column.\n
-_____________________________________________________________________________________________________
-
-
-Box2DGeometry:
-
-x               |   float   |   The distance, in pixels, between the left border of the bounding box
-                |           |   and the left border of the image.\n
-y               |   float   |   The distance, in pixels, between the top border of the bounding box
-                |           |   and the top border of the image.\n
-width	        |   float   |   The width in pixels of the annotation.\n
-height	        |   float   |   The height in pixels of the annotation.\n
-
-Box2DAnnotation:
-
-item_id         |   str     |   The internally-controlled item identifier to associate this annotation with.
-                |           |   The reference_id field should be empty if this field is populated.\n
-reference_id    |   str     |   The user-specified reference identifier to associate this annotation with.\n
-                |           |   The item_id field should be empty if this field is populated.
-label	        |   str     |	The label for this annotation (e.g. car, pedestrian, bicycle).\n
-type	        |   str     |   The type of this annotation. It should always be the box string literal.\n
-geometry        |   dict    |   Representation of the bounding box in the Box2DGeometry format.\n
-metadata        |   dict    |   An arbitrary metadata blob for the annotation.\n
-
-_____________________________________________________________________________________________________
-
-Box2DDetection:
-
-item_id         |   str     |   The internally-controlled item identifier to associate this annotation with.
-                |           |   The reference_id field should be empty if this field is populated.\n
-reference_id    |   str     |   The user-specified reference identifier to associate this annotation with.
-                |           |   The item_id field should be empty if this field is populated.\n
-label	        |   str     |	The label for this annotation (e.g. car, pedestrian, bicycle).\n
-type	        |   str     |   The type of this annotation. It should always be the box string literal.\n
-confidence      |   float   |   The optional confidence level of this annotation.
-                |           |   It should be between 0 and 1 (inclusive).\n
-geometry        |   dict    |   Representation of the bounding box in the Box2DGeometry format.\n
-metadata        |   dict    |   An arbitrary metadata blob for the annotation.\n
+For full documentation see: https://dashboard.scale.com/nucleus/docs/api?language=python
 """
 import asyncio
 import json
 import logging
 import os
+import time
 from typing import Any, Dict, List, Optional, Union
 
 import aiohttp
+import nest_asyncio
 import pkg_resources
 import requests
 import tqdm
 import tqdm.notebook as tqdm_notebook
 
+from nucleus.url_utils import sanitize_string_args
+
 from .annotation import (
     BoxAnnotation,
+    CuboidAnnotation,
+    Point,
+    Point3D,
     PolygonAnnotation,
+    CategoryAnnotation,
     Segment,
     SegmentationAnnotation,
-    Point,
 )
 from .constants import (
     ANNOTATION_METADATA_SCHEMA_KEY,
@@ -75,16 +35,23 @@ from .constants import (
     ANNOTATIONS_PROCESSED_KEY,
     AUTOTAGS_KEY,
     DATASET_ID_KEY,
-    DATASET_ITEM_IDS_KEY,
     DEFAULT_NETWORK_TIMEOUT_SEC,
+    EMBEDDING_DIMENSION_KEY,
     EMBEDDINGS_URL_KEY,
     ERROR_ITEMS,
     ERROR_PAYLOAD,
     ERRORS_KEY,
     IMAGE_KEY,
     IMAGE_URL_KEY,
+    INDEX_CONTINUOUS_ENABLE_KEY,
     ITEM_METADATA_SCHEMA_KEY,
     ITEMS_KEY,
+    JOB_CREATION_TIME_KEY,
+    JOB_ID_KEY,
+    JOB_LAST_KNOWN_STATUS_KEY,
+    JOB_TYPE_KEY,
+    KEEP_HISTORY_KEY,
+    MESSAGE_KEY,
     MODEL_RUN_ID_KEY,
     NAME_KEY,
     NUCLEUS_ENDPOINT,
@@ -96,7 +63,7 @@ from .constants import (
     UPDATE_KEY,
 )
 from .dataset import Dataset
-from .dataset_item import DatasetItem
+from .dataset_item import CameraParams, DatasetItem, Quaternion
 from .errors import (
     DatasetItemRetrievalError,
     ModelCreationError,
@@ -104,6 +71,7 @@ from .errors import (
     NotFoundError,
     NucleusAPIError,
 )
+from .job import AsyncJob
 from .model import Model
 from .model_run import ModelRun
 from .payload_constructor import (
@@ -115,9 +83,11 @@ from .payload_constructor import (
 )
 from .prediction import (
     BoxPrediction,
+    CuboidPrediction,
     PolygonPrediction,
     SegmentationPrediction,
 )
+from .scene import Frame, LidarScene
 from .slice import Slice
 from .upload_response import UploadResponse
 
@@ -135,6 +105,11 @@ logging.getLogger(requests.packages.urllib3.__package__).setLevel(
 )
 
 
+class RetryStrategy:
+    statuses = {503, 504}
+    sleep_times = [1, 3, 9]
+
+
 class NucleusClient:
     """
     Nucleus client.
@@ -176,11 +151,11 @@ class NucleusClient:
 
         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"]
         ]
@@ -192,6 +167,26 @@ class NucleusClient:
         """
         return self.make_request({}, "dataset/", requests.get)
 
+    def list_jobs(
+        self, show_completed=None, date_limit=None
+    ) -> List[AsyncJob]:
+        """
+        Lists jobs for user.
+        :return: jobs
+        """
+        payload = {show_completed: show_completed, date_limit: date_limit}
+        job_objects = self.make_request(payload, "jobs/", requests.get)
+        return [
+            AsyncJob(
+                job_id=job[JOB_ID_KEY],
+                job_last_known_status=job[JOB_LAST_KNOWN_STATUS_KEY],
+                job_type=job[JOB_TYPE_KEY],
+                job_creation_time=job[JOB_CREATION_TIME_KEY],
+                client=self,
+            )
+            for job in job_objects
+        ]
+
     def get_dataset_items(self, dataset_id) -> List[DatasetItem]:
         """
         Gets all the dataset items inside your repo as a json blob.
@@ -207,11 +202,8 @@ class NucleusClient:
             for item in dataset_items:
                 image_url = item.get("original_image_url")
                 metadata = item.get("metadata", None)
-                item_id = item.get("id", None)
                 ref_id = item.get("ref_id", None)
-                dataset_item = DatasetItem(
-                    image_url, ref_id, item_id, metadata
-                )
+                dataset_item = DatasetItem(image_url, ref_id, metadata)
                 constructed_dataset_items.append(dataset_item)
         elif error:
             raise DatasetItemRetrievalError(message=error)
@@ -226,6 +218,19 @@ class NucleusClient:
         """
         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
@@ -298,9 +303,8 @@ class NucleusClient:
         """
         return self.make_request({}, f"dataset/{dataset_id}", requests.delete)
 
-    def delete_dataset_item(
-        self, dataset_id: str, item_id: str = None, reference_id: str = None
-    ) -> dict:
+    @sanitize_string_args
+    def delete_dataset_item(self, dataset_id: str, reference_id) -> dict:
         """
         Deletes a private dataset based on datasetId.
         Returns an empty payload where response status `200` indicates
@@ -308,16 +312,11 @@ class NucleusClient:
         :param payload: { "name": str }
         :return: { "dataset_id": str, "name": str }
         """
-        if item_id:
-            return self.make_request(
-                {}, f"dataset/{dataset_id}/{item_id}", requests.delete
-            )
-        else:  # Assume reference_id is provided
-            return self.make_request(
-                {},
-                f"dataset/{dataset_id}/refloc/{reference_id}",
-                requests.delete,
-            )
+        return self.make_request(
+            {},
+            f"dataset/{dataset_id}/refloc/{reference_id}",
+            requests.delete,
+        )
 
     def populate_dataset(
         self,
@@ -366,28 +365,33 @@ class NucleusClient:
 
         agg_response = UploadResponse(json={DATASET_ID_KEY: dataset_id})
 
-        tqdm_local_batches = self.tqdm_bar(local_batches)
-
-        tqdm_remote_batches = self.tqdm_bar(remote_batches)
-
         async_responses: List[Any] = []
 
-        for batch in tqdm_local_batches:
-            payload = construct_append_payload(batch, update)
-            responses = self._process_append_requests_local(
-                dataset_id, payload, update
+        if local_batches:
+            tqdm_local_batches = self.tqdm_bar(
+                local_batches, desc="Local file batches"
             )
-            async_responses.extend(responses)
-
-        for batch in tqdm_remote_batches:
-            payload = construct_append_payload(batch, update)
-            responses = self._process_append_requests(
-                dataset_id=dataset_id,
-                payload=payload,
-                update=update,
-                batch_size=batch_size,
+
+            for batch in tqdm_local_batches:
+                payload = construct_append_payload(batch, update)
+                responses = self._process_append_requests_local(
+                    dataset_id, payload, update
+                )
+                async_responses.extend(responses)
+
+        if remote_batches:
+            tqdm_remote_batches = self.tqdm_bar(
+                remote_batches, desc="Remote file batches"
             )
-            async_responses.extend(responses)
+            for batch in tqdm_remote_batches:
+                payload = construct_append_payload(batch, update)
+                responses = self._process_append_requests(
+                    dataset_id=dataset_id,
+                    payload=payload,
+                    update=update,
+                    batch_size=batch_size,
+                )
+                async_responses.extend(responses)
 
         for response in async_responses:
             agg_response.update_response(response)
@@ -402,6 +406,8 @@ class NucleusClient:
         local_batch_size: int = 10,
     ):
         def get_files(batch):
+            for item in batch:
+                item[UPDATE_KEY] = update
             request_payload = [
                 (
                     ITEMS_KEY,
@@ -434,14 +440,20 @@ class NucleusClient:
             files_per_request.append(get_files(batch))
             payload_items.append(batch)
 
-        loop = asyncio.get_event_loop()
-        responses = loop.run_until_complete(
-            self.make_many_files_requests_asynchronously(
-                files_per_request,
-                f"dataset/{dataset_id}/append",
-            )
+        future = self.make_many_files_requests_asynchronously(
+            files_per_request,
+            f"dataset/{dataset_id}/append",
         )
 
+        try:
+            loop = asyncio.get_event_loop()
+        except RuntimeError:  # no event loop running:
+            loop = asyncio.new_event_loop()
+            responses = loop.run_until_complete(future)
+        else:
+            nest_asyncio.apply(loop)
+            return loop.run_until_complete(future)
+
         def close_files(request_items):
             for item in request_items:
                 # file buffer in location [1][1]
@@ -504,28 +516,41 @@ class NucleusClient:
                 content_type=file[1][2],
             )
 
-        async with session.post(
-            endpoint,
-            data=form,
-            auth=aiohttp.BasicAuth(self.api_key, ""),
-            timeout=DEFAULT_NETWORK_TIMEOUT_SEC,
-        ) as response:
-            logger.info("API request has response code %s", response.status)
-
-            try:
-                data = await response.json()
-            except aiohttp.client_exceptions.ContentTypeError:
-                # In case of 404, the server returns text
-                data = await response.text()
-
-            if not response.ok:
-                self.handle_bad_response(
-                    endpoint,
-                    session.post,
-                    aiohttp_response=(response.status, response.reason, data),
+        for sleep_time in RetryStrategy.sleep_times + [-1]:
+            async with session.post(
+                endpoint,
+                data=form,
+                auth=aiohttp.BasicAuth(self.api_key, ""),
+                timeout=DEFAULT_NETWORK_TIMEOUT_SEC,
+            ) as response:
+                logger.info(
+                    "API request has response code %s", response.status
                 )
 
-            return data
+                try:
+                    data = await response.json()
+                except aiohttp.client_exceptions.ContentTypeError:
+                    # In case of 404, the server returns text
+                    data = await response.text()
+                if (
+                    response.status in RetryStrategy.statuses
+                    and sleep_time != -1
+                ):
+                    time.sleep(sleep_time)
+                    continue
+
+                if not response.ok:
+                    self.handle_bad_response(
+                        endpoint,
+                        session.post,
+                        aiohttp_response=(
+                            response.status,
+                            response.reason,
+                            data,
+                        ),
+                    )
+
+                return data
 
     def _process_append_requests(
         self,
@@ -553,7 +578,13 @@ class NucleusClient:
         self,
         dataset_id: str,
         annotations: List[
-            Union[BoxAnnotation, PolygonAnnotation, SegmentationAnnotation]
+            Union[
+                BoxAnnotation,
+                PolygonAnnotation,
+                CuboidAnnotation,
+                CategoryAnnotation,
+                SegmentationAnnotation,
+            ]
         ],
         update: bool,
         batch_size: int = 5000,
@@ -561,11 +592,10 @@ class NucleusClient:
         """
         Uploads ground truth annotations for a given dataset.
         :param dataset_id: id of the dataset
-        :param annotations: List[Union[BoxAnnotation, PolygonAnnotation]]
+        :param annotations: List[Union[BoxAnnotation, PolygonAnnotation, CuboidAnnotation, SegmentationAnnotation]]
         :param update: whether to update or ignore conflicting annotations
         :return: {"dataset_id: str, "annotations_processed": int}
         """
-
         # Split payload into segmentations and Box/Polygon
         segmentations = [
             ann
@@ -706,14 +736,19 @@ class NucleusClient:
         self,
         model_run_id: str,
         annotations: List[
-            Union[BoxPrediction, PolygonPrediction, SegmentationPrediction]
+            Union[
+                BoxPrediction,
+                PolygonPrediction,
+                CuboidPrediction,
+                SegmentationPrediction,
+            ]
         ],
         update: bool,
         batch_size: int = 5000,
     ):
         """
         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]],
         :param update: bool
         :return:
         {
@@ -855,6 +890,7 @@ class NucleusClient:
             {}, f"modelRun/{model_run_id}/info", requests.get
         )
 
+    @sanitize_string_args
     def dataitem_ref_id(self, dataset_id: str, reference_id: str):
         """
         :param dataset_id: internally controlled dataset id
@@ -865,6 +901,7 @@ class NucleusClient:
             {}, f"dataset/{dataset_id}/refloc/{reference_id}", requests.get
         )
 
+    @sanitize_string_args
     def predictions_ref_id(self, model_run_id: str, ref_id: str):
         """
         Returns Model Run info For Dataset Item by model_run_id and item reference_id.
@@ -872,7 +909,7 @@ class NucleusClient:
         :param reference_id: reference_id of a dataset item.
         :return:
         {
-            "annotations": List[BoxPrediction],
+            "annotations": List[Union[BoxPrediction, PolygonPrediction, CuboidPrediction, SegmentationPrediction]],
         }
         """
         return self.make_request(
@@ -897,7 +934,7 @@ class NucleusClient:
         :param i: absolute number of Dataset Item for a dataset corresponding to the model run.
         :return:
         {
-            "annotations": List[BoxPrediction],
+            "annotations": List[Union[BoxPrediction, PolygonPrediction, CuboidPrediction, SegmentationPrediction]],
         }
         """
         return self.make_request(
@@ -926,7 +963,7 @@ class NucleusClient:
         :param dataset_item_id: dataset_item_id of a dataset item.
         :return:
         {
-            "annotations": List[BoxPrediction],
+            "annotations": List[Union[BoxPrediction, PolygonPrediction, CuboidPrediction, SegmentationPrediction]],
         }
         """
         return self.make_request(
@@ -940,9 +977,6 @@ class NucleusClient:
         as a means of identifying items in the dataset.
 
         "name" -- The human-readable name of the slice.
-
-        "dataset_item_ids" -- An optional list of dataset item ids for the items in the slice
-
         "reference_ids" -- An optional list of user-specified identifier for the items in the slice
 
         :param
@@ -950,7 +984,6 @@ class NucleusClient:
         payload:
         {
             "name": str,
-            "dataset_item_ids": List[str],
             "reference_ids": List[str],
         }
         :return: new Slice object
@@ -976,14 +1009,12 @@ class NucleusClient:
 
         :param
         slice_id: id of the slice
-        id_type: the type of IDs you want in response (either "reference_id" or "dataset_item_id")
-        to identify the DatasetItems
 
         :return:
         {
             "name": str,
             "dataset_id": str,
-            "dataset_item_ids": List[str],
+            "reference_ids": List[str],
         }
         """
         response = self.make_request(
@@ -1010,11 +1041,32 @@ class NucleusClient:
         )
         return response
 
+    def delete_annotations(
+        self, dataset_id: str, reference_ids: list = None, keep_history=False
+    ) -> dict:
+        """
+        This endpoint deletes annotations.
+
+        :param
+        slice_id: id of the slice
+
+        :return:
+        {}
+        """
+        payload = {KEEP_HISTORY_KEY: keep_history}
+        if reference_ids:
+            payload[REFERENCE_IDS_KEY] = reference_ids
+        response = self.make_request(
+            payload,
+            f"annotation/{dataset_id}",
+            requests_command=requests.delete,
+        )
+        return response
+
     def append_to_slice(
         self,
         slice_id: str,
-        dataset_item_ids: List[str] = None,
-        reference_ids: List[str] = None,
+        reference_ids: List[str],
     ) -> dict:
         """
         Appends to a slice from items already present in a dataset.
@@ -1022,7 +1074,6 @@ class NucleusClient:
         as a means of identifying items in the dataset.
 
         :param
-        dataset_item_ids: List[str],
         reference_ids: List[str],
 
         :return:
@@ -1030,18 +1081,10 @@ class NucleusClient:
             "slice_id": str,
         }
         """
-        if dataset_item_ids and reference_ids:
-            raise Exception(
-                "You cannot specify both dataset_item_ids and reference_ids"
-            )
-
-        ids_to_append: Dict[str, Any] = {}
-        if dataset_item_ids:
-            ids_to_append[DATASET_ITEM_IDS_KEY] = dataset_item_ids
-        if reference_ids:
-            ids_to_append[REFERENCE_IDS_KEY] = reference_ids
 
-        response = self.make_request(ids_to_append, f"slice/{slice_id}/append")
+        response = self.make_request(
+            {REFERENCE_IDS_KEY: reference_ids}, f"slice/{slice_id}/append"
+        )
         return response
 
     def list_autotags(self, dataset_id: str) -> List[str]:
@@ -1057,6 +1100,16 @@ class NucleusClient:
         )
         return response[AUTOTAGS_KEY] if AUTOTAGS_KEY in response else response
 
+    def delete_autotag(self, autotag_id: str) -> dict:
+        """
+        Deletes an autotag based on autotagId.
+        Returns an empty payload where response status `200` indicates
+        the autotag has been successfully deleted.
+        :param autotag_id: id of the autotag to delete.
+        :return: {}
+        """
+        return self.make_request({}, f"autotag/{autotag_id}", requests.delete)
+
     def delete_model(self, model_id: str) -> dict:
         """
         This endpoint deletes the specified model, along with all
@@ -1075,25 +1128,63 @@ class NucleusClient:
         )
         return response
 
-    def create_custom_index(self, dataset_id: str, embeddings_url: str):
+    def create_custom_index(
+        self, dataset_id: str, embeddings_urls: list, embedding_dim: int
+    ):
+        """
+        Creates a custom index for a given dataset, which will then be used
+        for autotag and similarity search.
+
+        :param
+        dataset_id: id of dataset that the custom index is being added to.
+        embeddings_urls: list of urls, each of which being a json mapping reference_id -> embedding vector
+        embedding_dim: the dimension of the embedding vectors, must be consistent for all embedding vectors in the index.
+        """
         return self.make_request(
-            {EMBEDDINGS_URL_KEY: embeddings_url},
+            {
+                EMBEDDINGS_URL_KEY: embeddings_urls,
+                EMBEDDING_DIMENSION_KEY: embedding_dim,
+            },
             f"indexing/{dataset_id}",
             requests_command=requests.post,
         )
 
-    def check_index_status(self, job_id: str):
+    def delete_custom_index(self, dataset_id: str):
         return self.make_request(
             {},
-            f"indexing/{job_id}",
-            requests_command=requests.get,
+            f"indexing/{dataset_id}",
+            requests_command=requests.delete,
         )
 
-    def delete_custom_index(self, dataset_id: str):
+    def set_continuous_indexing(self, dataset_id: str, enable: bool = True):
+        """
+        Sets continuous indexing for a given dataset, which will automatically generate embeddings whenever
+        new images are uploaded. This endpoint is currently only enabled for enterprise customers.
+        Please reach out to nucleus@scale.com if you wish to learn more.
+
+        :param
+        dataset_id: id of dataset that continuous indexing is being toggled for
+        enable: boolean, sets whether we are enabling or disabling continuous indexing. The default behavior is to enable.
+        """
+        return self.make_request(
+            {INDEX_CONTINUOUS_ENABLE_KEY: enable},
+            f"indexing/{dataset_id}/setContinuous",
+            requests_command=requests.post,
+        )
+
+    def create_image_index(self, dataset_id: str):
+        """
+        Starts generating embeddings for images that don't have embeddings in a given dataset. These embeddings will
+        be used for autotag and similarity search. This endpoint is currently only enabled for enterprise customers.
+        Please reach out to nucleus@scale.com if you wish to learn more.
+
+        :param
+        dataset_id: id of dataset for generating embeddings on.
+        """
         return self.make_request(
             {},
-            f"indexing/{dataset_id}",
-            requests_command=requests.delete,
+            f"indexing/{dataset_id}/internal/image",
+            requests_command=requests.post,
         )
 
     def make_request(
@@ -1112,14 +1203,20 @@ class NucleusClient:
 
         logger.info("Posting to %s", endpoint)
 
-        response = requests_command(
-            endpoint,
-            json=payload,
-            headers={"Content-Type": "application/json"},
-            auth=(self.api_key, ""),
-            timeout=DEFAULT_NETWORK_TIMEOUT_SEC,
-        )
-        logger.info("API request has response code %s", response.status_code)
+        for retry_wait_time in RetryStrategy.sleep_times:
+            response = requests_command(
+                endpoint,
+                json=payload,
+                headers={"Content-Type": "application/json"},
+                auth=(self.api_key, ""),
+                timeout=DEFAULT_NETWORK_TIMEOUT_SEC,
+            )
+            logger.info(
+                "API request has response code %s", response.status_code
+            )
+            if response.status_code not in RetryStrategy.statuses:
+                break
+            time.sleep(retry_wait_time)
 
         if not response.ok:
             self.handle_bad_response(endpoint, requests_command, response)