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

nucleus/connection.py

@@ -0,0 +1,87 @@
+import time
+
+import requests
+
+from .constants import DEFAULT_NETWORK_TIMEOUT_SEC
+from .errors import NucleusAPIError
+from .logger import logger
+from .retry_strategy import RetryStrategy
+
+
+class Connection:
+    """Wrapper of HTTP requests to the Nucleus endpoint."""
+
+    def __init__(self, api_key: str, endpoint: str = None):
+        self.api_key = api_key
+        self.endpoint = endpoint
+
+    def __repr__(self):
+        return (
+            f"Connection(api_key='{self.api_key}', endpoint='{self.endpoint}')"
+        )
+
+    def __eq__(self, other):
+        return (
+            self.api_key == other.api_key and self.endpoint == other.endpoint
+        )
+
+    def delete(self, route: str):
+        return self.make_request({}, route, requests_command=requests.delete)
+
+    def get(self, route: str):
+        return self.make_request({}, route, requests_command=requests.get)
+
+    def post(self, payload: dict, route: str):
+        return self.make_request(
+            payload, route, requests_command=requests.post
+        )
+
+    def put(self, payload: dict, route: str):
+        return self.make_request(payload, route, requests_command=requests.put)
+
+    def make_request(
+        self, payload: dict, route: str, requests_command=requests.post
+    ) -> dict:
+        """
+        Makes a request to Nucleus endpoint and logs a warning if not
+        successful.
+
+        :param payload: given payload
+        :param route: route for the request
+        :param requests_command: requests.post, requests.get, requests.delete
+        :return: response JSON
+        """
+        endpoint = f"{self.endpoint}/{route}"
+
+        logger.info("Make request to %s", endpoint)
+
+        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)
+
+        return response.json()
+
+    def handle_bad_response(
+        self,
+        endpoint,
+        requests_command,
+        requests_response=None,
+        aiohttp_response=None,
+    ):
+        raise NucleusAPIError(
+            endpoint, requests_command, requests_response, aiohttp_response
+        )

nucleus/constants.py

@@ -3,6 +3,7 @@ ANNOTATIONS_KEY = "annotations"
 ANNOTATIONS_PROCESSED_KEY = "annotations_processed"
 ANNOTATION_ID_KEY = "annotation_id"
 ANNOTATION_METADATA_SCHEMA_KEY = "annotation_metadata_schema"
+BACKFILL_JOB_KEY = "backfill_job"
 BOX_TYPE = "box"
 POLYGON_TYPE = "polygon"
 MASK_TYPE = "mask"
@@ -16,6 +17,7 @@ ANNOTATION_TYPES = (
     SEGMENTATION_TYPE,
     CUBOID_TYPE,
     CATEGORY_TYPE,
+    MULTICATEGORY_TYPE,
 )
 ANNOTATION_UPDATE_KEY = "update"
 AUTOTAGS_KEY = "autotags"
@@ -27,6 +29,8 @@ CONFIDENCE_KEY = "confidence"
 CX_KEY = "cx"
 CY_KEY = "cy"
 DATASET_ID_KEY = "dataset_id"
+DATASET_IS_SCENE_KEY = "is_scene"
+DATASET_ITEM_ID_KEY = "dataset_item_id"
 DATASET_LENGTH_KEY = "length"
 DATASET_MODEL_RUNS_KEY = "model_run_ids"
 DATASET_NAME_KEY = "name"
@@ -46,6 +50,7 @@ FY_KEY = "fy"
 GEOMETRY_KEY = "geometry"
 HEADING_KEY = "heading"
 HEIGHT_KEY = "height"
+ID_KEY = "id"
 IGNORED_ITEMS = "ignored_items"
 IMAGE_KEY = "image"
 IMAGE_LOCATION_KEY = "image_location"
@@ -54,7 +59,6 @@ INDEX_KEY = "index"
 INDEX_CONTINUOUS_ENABLE_KEY = "enable"
 ITEMS_KEY = "items"
 ITEM_KEY = "item"
-ITEMS_KEY = "items"
 ITEM_METADATA_SCHEMA_KEY = "item_metadata_schema"
 JOB_ID_KEY = "job_id"
 KEEP_HISTORY_KEY = "keep_history"

nucleus/data_transfer_object/dataset_details.py

@@ -0,0 +1,9 @@
+from nucleus.pydantic_base import DictCompatibleModel
+
+
+class DatasetDetails(DictCompatibleModel):
+    """Pydantic model to parse DatasetDetails response from JSON"""
+
+    id: str
+    name: str
+    is_public: bool

nucleus/data_transfer_object/dataset_info.py

@@ -0,0 +1,26 @@
+from typing import Any, Dict, List, Optional
+
+from nucleus.pydantic_base import DictCompatibleModel
+
+
+class DatasetInfo(DictCompatibleModel):
+    """High-level :class:`Dataset` information
+
+    Attributes:
+        dataset_id: Nucleus-generated dataset ID
+        name: User-defined name of dataset
+        length: Number of :class:`DatasetItem` in :class:`Dataset`
+        model_run_ids: (deprecated)
+        slice_ids: List :class:`Slice` IDs associated with the :class:`Dataset`
+        annotation_metadata_schema: Dict defining annotation-level metadata schema.
+        item_metadata_schema: Dict defining item metadata schema.
+    """
+
+    dataset_id: str
+    name: str
+    length: int
+    model_run_ids: List[str]
+    slice_ids: List[str]
+    # TODO: Expand the following into pydantic models to formalize schema
+    annotation_metadata_schema: Optional[Dict[str, Any]] = None
+    item_metadata_schema: Optional[Dict[str, Any]] = None

nucleus/data_transfer_object/dataset_size.py

@@ -0,0 +1,5 @@
+from nucleus.pydantic_base import DictCompatibleModel
+
+
+class DatasetSize(DictCompatibleModel):
+    count: int

nucleus/data_transfer_object/scenes_list.py

@@ -0,0 +1,18 @@
+from typing import List, Optional
+
+from nucleus.pydantic_base import DictCompatibleImmutableModel
+
+
+class ScenesListEntry(DictCompatibleImmutableModel):
+    """/dataset/{dataset_id}/scenes_list nested payload"""
+
+    id: str
+    reference_id: str
+    type: str
+    metadata: Optional[dict] = {}
+
+
+class ScenesList(DictCompatibleImmutableModel):
+    """/dataset/{dataset_id}/scenes_list payload"""
+
+    scenes: List[ScenesListEntry]