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

nucleus/__init__.py

@@ -7,8 +7,7 @@ import asyncio
 import json
 import logging
 import os
-import urllib.request
-from asyncio.tasks import Task
+import time
 from typing import Any, Dict, List, Optional, Union
 
 import aiohttp
@@ -26,6 +25,7 @@ from .annotation import (
     Point,
     Point3D,
     PolygonAnnotation,
+    CategoryAnnotation,
     Segment,
     SegmentationAnnotation,
 )
@@ -41,15 +41,15 @@ from .constants import (
     ERROR_ITEMS,
     ERROR_PAYLOAD,
     ERRORS_KEY,
-    JOB_ID_KEY,
-    JOB_LAST_KNOWN_STATUS_KEY,
-    JOB_TYPE_KEY,
-    JOB_CREATION_TIME_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,
@@ -63,7 +63,7 @@ from .constants import (
     UPDATE_KEY,
 )
 from .dataset import Dataset
-from .dataset_item import DatasetItem, CameraParams, Quaternion
+from .dataset_item import CameraParams, DatasetItem, Quaternion
 from .errors import (
     DatasetItemRetrievalError,
     ModelCreationError,
@@ -87,9 +87,9 @@ from .prediction import (
     PolygonPrediction,
     SegmentationPrediction,
 )
+from .scene import Frame, LidarScene
 from .slice import Slice
 from .upload_response import UploadResponse
-from .scene import Frame, LidarScene
 
 # pylint: disable=E1101
 # TODO: refactor to reduce this file to under 1000 lines.
@@ -105,6 +105,11 @@ logging.getLogger(requests.packages.urllib3.__package__).setLevel(
 )
 
 
+class RetryStrategy:
+    statuses = {503, 504}
+    sleep_times = [1, 3, 9]
+
+
 class NucleusClient:
     """
     Nucleus client.
@@ -511,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,
@@ -564,6 +582,7 @@ class NucleusClient:
                 BoxAnnotation,
                 PolygonAnnotation,
                 CuboidAnnotation,
+                CategoryAnnotation,
                 SegmentationAnnotation,
             ]
         ],
@@ -1130,13 +1149,6 @@ class NucleusClient:
             requests_command=requests.post,
         )
 
-    def check_index_status(self, job_id: str):
-        return self.make_request(
-            {},
-            f"indexing/{job_id}",
-            requests_command=requests.get,
-        )
-
     def delete_custom_index(self, dataset_id: str):
         return self.make_request(
             {},
@@ -1191,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)

nucleus/annotation.py

@@ -8,6 +8,7 @@ from .constants import (
     ANNOTATION_ID_KEY,
     ANNOTATIONS_KEY,
     BOX_TYPE,
+    CATEGORY_TYPE,
     CUBOID_TYPE,
     DIMENSIONS_KEY,
     GEOMETRY_KEY,
@@ -20,6 +21,7 @@ from .constants import (
     POLYGON_TYPE,
     POSITION_KEY,
     REFERENCE_ID_KEY,
+    TAXONOMY_NAME_KEY,
     TYPE_KEY,
     VERTICES_KEY,
     WIDTH_KEY,
@@ -41,6 +43,8 @@ class Annotation:
             return PolygonAnnotation.from_json(payload)
         elif payload.get(TYPE_KEY, None) == CUBOID_TYPE:
             return CuboidAnnotation.from_json(payload)
+        elif payload.get(TYPE_KEY, None) == CATEGORY_TYPE:
+            return CategoryAnnotation.from_json(payload)
         else:
             return SegmentationAnnotation.from_json(payload)
 
@@ -120,6 +124,7 @@ class AnnotationTypes(Enum):
     BOX = BOX_TYPE
     POLYGON = POLYGON_TYPE
     CUBOID = CUBOID_TYPE
+    CATEGORY = CATEGORY_TYPE
 
 
 @dataclass  # pylint: disable=R0902
@@ -291,6 +296,36 @@ class CuboidAnnotation(Annotation):  # pylint: disable=R0902
         return payload
 
 
+@dataclass
+class CategoryAnnotation(Annotation):
+    label: str
+    taxonomy_name: str
+    reference_id: str
+    metadata: Optional[Dict] = None
+
+    def __post_init__(self):
+        self.metadata = self.metadata if self.metadata else {}
+
+    @classmethod
+    def from_json(cls, payload: dict):
+        return cls(
+            label=payload[LABEL_KEY],
+            taxonomy_name=payload[TAXONOMY_NAME_KEY],
+            reference_id=payload[REFERENCE_ID_KEY],
+            metadata=payload.get(METADATA_KEY, {}),
+        )
+
+    def to_payload(self) -> dict:
+        return {
+            LABEL_KEY: self.label,
+            TAXONOMY_NAME_KEY: self.taxonomy_name,
+            TYPE_KEY: CATEGORY_TYPE,
+            GEOMETRY_KEY: {},
+            REFERENCE_ID_KEY: self.reference_id,
+            METADATA_KEY: self.metadata,
+        }
+
+
 def is_local_path(path: str) -> bool:
     return urlparse(path).scheme not in {"https", "http", "s3", "gs"}
 

nucleus/constants.py

@@ -10,7 +10,13 @@ SEGMENTATION_TYPE = "segmentation"
 CUBOID_TYPE = "cuboid"
 CATEGORY_TYPE = "category"
 MULTICATEGORY_TYPE = "multicategory"
-ANNOTATION_TYPES = (BOX_TYPE, POLYGON_TYPE, SEGMENTATION_TYPE, CUBOID_TYPE)
+ANNOTATION_TYPES = (
+    BOX_TYPE,
+    POLYGON_TYPE,
+    SEGMENTATION_TYPE,
+    CUBOID_TYPE,
+    CATEGORY_TYPE,
+)
 ANNOTATION_UPDATE_KEY = "update"
 AUTOTAGS_KEY = "autotags"
 AUTOTAG_SCORE_THRESHOLD = "score_threshold"

nucleus/dataset.py

@@ -217,7 +217,6 @@ class Dataset:
                 route=f"dataset/{self.id}/annotate?async=1",
             )
             return AsyncJob.from_json(response, self._client)
-
         return self._client.annotate_dataset(
             self.id, annotations, update=update, batch_size=batch_size
         )
@@ -428,10 +427,13 @@ class Dataset:
         return self._client.list_autotags(self.id)
 
     def create_custom_index(self, embeddings_urls: list, embedding_dim: int):
-        return self._client.create_custom_index(
-            self.id,
-            embeddings_urls,
-            embedding_dim,
+        return AsyncJob.from_json(
+            self._client.create_custom_index(
+                self.id,
+                embeddings_urls,
+                embedding_dim,
+            ),
+            self._client,
         )
 
     def delete_custom_index(self):
@@ -463,9 +465,6 @@ class Dataset:
             requests_command=requests.post,
         )
 
-    def check_index_status(self, job_id: str):
-        return self._client.check_index_status(job_id)
-
     def items_and_annotations(
         self,
     ) -> List[Dict[str, Union[DatasetItem, Dict[str, List[Annotation]]]]]:

nucleus/errors.py

@@ -4,6 +4,11 @@ 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"):
@@ -35,7 +40,7 @@ class NucleusAPIError(Exception):
     def __init__(
         self, endpoint, command, requests_response=None, aiohttp_response=None
     ):
-        message = f"Your client is on version {nucleus_client_version}. Before reporting this error, please make sure you update to the latest version of the client by running pip install --upgrade scale-nucleus\n"
+        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}."
             if hasattr(requests_response, "text"):
@@ -50,4 +55,10 @@ class NucleusAPIError(Exception):
             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)

nucleus/payload_constructor.py

@@ -5,6 +5,7 @@ from .annotation import (
     BoxAnnotation,
     CuboidAnnotation,
     PolygonAnnotation,
+    CategoryAnnotation,
     SegmentationAnnotation,
 )
 from .prediction import (
@@ -60,6 +61,7 @@ def construct_annotation_payload(
             BoxAnnotation,
             PolygonAnnotation,
             CuboidAnnotation,
+            CategoryAnnotation,
             SegmentationAnnotation,
         ]
     ],

nucleus/utils.py

@@ -14,6 +14,7 @@ from nucleus.annotation import (
     BoxAnnotation,
     CuboidAnnotation,
     PolygonAnnotation,
+    CategoryAnnotation,
     SegmentationAnnotation,
 )
 
@@ -22,6 +23,7 @@ from .constants import (
     ANNOTATIONS_KEY,
     BOX_TYPE,
     CUBOID_TYPE,
+    CATEGORY_TYPE,
     ITEM_KEY,
     POLYGON_TYPE,
     REFERENCE_ID_KEY,
@@ -116,6 +118,11 @@ def convert_export_payload(api_payload):
         for cuboid in row[CUBOID_TYPE]:
             cuboid[REFERENCE_ID_KEY] = row[ITEM_KEY][REFERENCE_ID_KEY]
             annotations[CUBOID_TYPE].append(CuboidAnnotation.from_json(cuboid))
+        for category in row[CATEGORY_TYPE]:
+            category[REFERENCE_ID_KEY] = row[ITEM_KEY][REFERENCE_ID_KEY]
+            annotations[CATEGORY_TYPE].append(
+                CategoryAnnotation.from_json(category)
+            )
         return_payload_row[ANNOTATIONS_KEY] = annotations
         return_payload.append(return_payload_row)
     return return_payload

{scale_nucleus-0.1.22.dist-info → scale_nucleus-0.1.24.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: scale-nucleus
-Version: 0.1.22
+Version: 0.1.24
 Summary: The official Python client library for Nucleus, the Data Platform for AI
 Home-page: https://scale.com/nucleus
 License: MIT

scale_nucleus-0.1.24.dist-info/RECORD

@@ -0,0 +1,21 @@
+nucleus/__init__.py,sha256=105pVyWKhc34vRxhXTFbL9APvyH9Ka6FWOMOCElFsp8,40780
+nucleus/annotation.py,sha256=tjkO_DCJIXQTTMI9gkWXe9W3lveyFsIQjlsM5jfyFyw,10007
+nucleus/autocurate.py,sha256=ogEX3kbuKCciWODOnTjUHU-JSwhQ_34wbNvW4xA79oY,854
+nucleus/constants.py,sha256=86tEkPqITYgd3SB_OWcG5LDcuAUGuc78kBtS5WOqo64,3026
+nucleus/dataset.py,sha256=0amQbRnC3JbcDz_coJNvQsZsmfp41EYiqbXEtVh_m00,18290
+nucleus/dataset_item.py,sha256=lKMMwNH9Iz5jxf1beIJSWrcD1UYNXbMbnPwenVW1He0,5781
+nucleus/errors.py,sha256=quBOj9Dwi8NrC6SIqSI6DLv-fT49e315OSLirSiF4kQ,2338
+nucleus/job.py,sha256=N2Ei3zJflcUyiZBavJOph3eLvckLANMrL7SwYzLUYAA,2301
+nucleus/model.py,sha256=akuWKehw6u5fp-FfBuI2RobkSoceNN-huh9_G3rxWPo,2147
+nucleus/model_run.py,sha256=-m_YzEqv253foD_ZQAIvD66CuDipvtKedzq9Pk0IBs4,7983
+nucleus/payload_constructor.py,sha256=UN9J0NEL6gJqh-EAvwEc51eXJSTaK9ZMH1p0FDgMDsI,3567
+nucleus/prediction.py,sha256=WJu5echvJKBjL67lQ6U9jM_LlbXvA1SPhUHyzdTeVpE,6276
+nucleus/scene.py,sha256=w8mNU5Pt7U-jn9WQCL4Ch7AaZ2RHVPW8nTtIhlqTx0k,7803
+nucleus/slice.py,sha256=zVLF6YyxU0ShJTERGTydcm1XiEx1yaVfJ1coq4H5KrI,5737
+nucleus/upload_response.py,sha256=pwOb3iS6TbpoumC1Mao6Pyli7dXBRDcI0zjNfCMU4_c,2729
+nucleus/url_utils.py,sha256=6iODEEVAa061-ROkqYM_Zhc4RbPHqOSYMczqYGVv4y0,660
+nucleus/utils.py,sha256=WDBx8tw5MEFA1afS9Z0difBi6SQCk56SJX-hfDkBq5k,6194
+scale_nucleus-0.1.24.dist-info/LICENSE,sha256=jaTGyQSQIZeWMo5iyYqgbAYHR9Bdy7nOzgE-Up3m_-g,1075
+scale_nucleus-0.1.24.dist-info/WHEEL,sha256=DRf8A_Psd1SF2kVqTQOOFU1Xzl3-A2qljAxBMTOusUs,83
+scale_nucleus-0.1.24.dist-info/METADATA,sha256=sxWeNc6pC9LBbOll4dfwRyqymKYOljHVHy8LslAoZvM,2656
+scale_nucleus-0.1.24.dist-info/RECORD,,

{scale_nucleus-0.1.22.dist-info → scale_nucleus-0.1.24.dist-info}/WHEEL

@@ -1,4 +1,4 @@
 Wheel-Version: 1.0
-Generator: poetry 1.0.4
+Generator: poetry 1.0.6
 Root-Is-Purelib: true
 Tag: py3-none-any

scale_nucleus-0.1.22.dist-info/RECORD

@@ -1,21 +0,0 @@
-nucleus/__init__.py,sha256=PRe5fr-9tPNkYDhqiGmJrEePcmd8h4JgbDUOFsXn7W8,40120
-nucleus/annotation.py,sha256=csyeeC_C50CCvR7J6BXJpfc85JfLxP3FX5QqABRtwm0,8992
-nucleus/autocurate.py,sha256=ogEX3kbuKCciWODOnTjUHU-JSwhQ_34wbNvW4xA79oY,854
-nucleus/constants.py,sha256=kjpahNP4_gRLWOdd0LXT_v6z3UYOQ3qbF2CX-37wJAc,2988
-nucleus/dataset.py,sha256=yRr6u7ucc-A0zjNmjbZbryjX7skivOvoUCckfLJqf5E,18309
-nucleus/dataset_item.py,sha256=lKMMwNH9Iz5jxf1beIJSWrcD1UYNXbMbnPwenVW1He0,5781
-nucleus/errors.py,sha256=XqBV5t0tEbV5yFM1RwZDjAvHWrkzsmTdprckzLZ6Kng,1924
-nucleus/job.py,sha256=N2Ei3zJflcUyiZBavJOph3eLvckLANMrL7SwYzLUYAA,2301
-nucleus/model.py,sha256=akuWKehw6u5fp-FfBuI2RobkSoceNN-huh9_G3rxWPo,2147
-nucleus/model_run.py,sha256=-m_YzEqv253foD_ZQAIvD66CuDipvtKedzq9Pk0IBs4,7983
-nucleus/payload_constructor.py,sha256=2O4jd774NBTOjDy_fBsN1SOTYKmtDkz4yYVPa_nuq-0,3511
-nucleus/prediction.py,sha256=WJu5echvJKBjL67lQ6U9jM_LlbXvA1SPhUHyzdTeVpE,6276
-nucleus/scene.py,sha256=w8mNU5Pt7U-jn9WQCL4Ch7AaZ2RHVPW8nTtIhlqTx0k,7803
-nucleus/slice.py,sha256=zVLF6YyxU0ShJTERGTydcm1XiEx1yaVfJ1coq4H5KrI,5737
-nucleus/upload_response.py,sha256=pwOb3iS6TbpoumC1Mao6Pyli7dXBRDcI0zjNfCMU4_c,2729
-nucleus/url_utils.py,sha256=6iODEEVAa061-ROkqYM_Zhc4RbPHqOSYMczqYGVv4y0,660
-nucleus/utils.py,sha256=sayZavDsvdPkvXm0QmvoSJ2_FZczje5n60uavHpOpgQ,5918
-scale_nucleus-0.1.22.dist-info/LICENSE,sha256=jaTGyQSQIZeWMo5iyYqgbAYHR9Bdy7nOzgE-Up3m_-g,1075
-scale_nucleus-0.1.22.dist-info/WHEEL,sha256=N0LZrBtofpkS5mJXgVHTCEy52Sam4D6PHQWC8HnMeTs,83
-scale_nucleus-0.1.22.dist-info/METADATA,sha256=Zh_k6vNFP6hkUn_FZgC0gPCL8Iwe2vSpBB2CM77mOSE,2656
-scale_nucleus-0.1.22.dist-info/RECORD,,