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

nucleus/constants.py

@@ -1,63 +1,67 @@
-NUCLEUS_ENDPOINT = "https://api.scale.com/v1/nucleus"
-DEFAULT_NETWORK_TIMEOUT_SEC = 120
-ITEMS_KEY = "items"
-ITEM_KEY = "item"
-REFERENCE_ID_KEY = "reference_id"
-REFERENCE_IDS_KEY = "reference_ids"
-DATASET_ID_KEY = "dataset_id"
-IMAGE_KEY = "image"
-IMAGE_URL_KEY = "image_url"
-NEW_ITEMS = "new_items"
-UPDATED_ITEMS = "updated_items"
-IGNORED_ITEMS = "ignored_items"
-ERROR_ITEMS = "upload_errors"
-ERROR_PAYLOAD = "error_payload"
-ERROR_CODES = "error_codes"
+ANNOTATIONS_IGNORED_KEY = "annotations_ignored"
 ANNOTATIONS_KEY = "annotations"
-ANNOTATION_ID_KEY = "annotation_id"
 ANNOTATIONS_PROCESSED_KEY = "annotations_processed"
-ANNOTATIONS_IGNORED_KEY = "annotations_ignored"
-PREDICTIONS_PROCESSED_KEY = "predictions_processed"
-PREDICTIONS_IGNORED_KEY = "predictions_ignored"
+ANNOTATION_ID_KEY = "annotation_id"
+ANNOTATION_METADATA_SCHEMA_KEY = "annotation_metadata_schema"
+BOX_TYPE = "box"
+POLYGON_TYPE = "polygon"
+MASK_TYPE = "mask"
+SEGMENTATION_TYPE = "segmentation"
+ANNOTATION_TYPES = (BOX_TYPE, POLYGON_TYPE, SEGMENTATION_TYPE)
 ANNOTATION_UPDATE_KEY = "update"
-DEFAULT_ANNOTATION_UPDATE_MODE = False
-STATUS_CODE_KEY = "status_code"
-STATUS_KEY = "status"
-SUCCESS_STATUS_CODES = [200, 201, 202]
-ERRORS_KEY = "errors"
-MODEL_RUN_ID_KEY = "model_run_id"
-MODEL_ID_KEY = "model_id"
-DATASET_ITEM_ID_KEY = "dataset_item_id"
-ITEM_ID_KEY = "item_id"
+AUTOTAGS_KEY = "autotags"
+EXPORTED_ROWS = "exportedRows"
+CLASS_PDF_KEY = "class_pdf"
+CONFIDENCE_KEY = "confidence"
+DATASET_ID_KEY = "dataset_id"
 DATASET_ITEM_IDS_KEY = "dataset_item_ids"
-SLICE_ID_KEY = "slice_id"
-DATASET_NAME_KEY = "name"
+DATASET_ITEM_ID_KEY = "dataset_item_id"
+DATASET_LENGTH_KEY = "length"
 DATASET_MODEL_RUNS_KEY = "model_run_ids"
+DATASET_NAME_KEY = "name"
 DATASET_SLICES_KEY = "slice_ids"
-DATASET_LENGTH_KEY = "length"
-FORCE_KEY = "update"
+DEFAULT_ANNOTATION_UPDATE_MODE = False
+DEFAULT_NETWORK_TIMEOUT_SEC = 120
+EMBEDDINGS_URL_KEY = "embeddings_url"
+ERRORS_KEY = "errors"
+ERROR_CODES = "error_codes"
+ERROR_ITEMS = "upload_errors"
+ERROR_PAYLOAD = "error_payload"
+GEOMETRY_KEY = "geometry"
+HEIGHT_KEY = "height"
+IGNORED_ITEMS = "ignored_items"
+IMAGE_KEY = "image"
+IMAGE_URL_KEY = "image_url"
+INDEX_KEY = "index"
+ITEMS_KEY = "items"
+ITEM_ID_KEY = "item_id"
+ITEM_KEY = "item"
+ITEM_METADATA_SCHEMA_KEY = "item_metadata_schema"
+JOB_ID_KEY = "job_id"
+LABEL_KEY = "label"
+MASK_URL_KEY = "mask_url"
+MESSAGE_KEY = "message"
 METADATA_KEY = "metadata"
+MODEL_ID_KEY = "model_id"
+MODEL_RUN_ID_KEY = "model_run_id"
 NAME_KEY = "name"
-LABEL_KEY = "label"
-CONFIDENCE_KEY = "confidence"
+NEW_ITEMS = "new_items"
+NUCLEUS_ENDPOINT = "https://api.scale.com/v1/nucleus"
 ORIGINAL_IMAGE_URL_KEY = "original_image_url"
-X_KEY = "x"
-Y_KEY = "y"
-WIDTH_KEY = "width"
-HEIGHT_KEY = "height"
+PREDICTIONS_IGNORED_KEY = "predictions_ignored"
+PREDICTIONS_PROCESSED_KEY = "predictions_processed"
+REFERENCE_IDS_KEY = "reference_ids"
+REFERENCE_ID_KEY = "reference_id"
+REQUEST_ID_KEY = "requestId"
+SEGMENTATIONS_KEY = "segmentations"
+SLICE_ID_KEY = "slice_id"
+STATUS_CODE_KEY = "status_code"
+STATUS_KEY = "status"
+SUCCESS_STATUS_CODES = [200, 201, 202]
 TYPE_KEY = "type"
+UPDATED_ITEMS = "updated_items"
+UPDATE_KEY = "update"
 VERTICES_KEY = "vertices"
-BOX_TYPE = "box"
-POLYGON_TYPE = "polygon"
-SEGMENTATION_TYPE = "segmentation"
-ANNOTATION_TYPES = (BOX_TYPE, POLYGON_TYPE, SEGMENTATION_TYPE)
-GEOMETRY_KEY = "geometry"
-AUTOTAGS_KEY = "autotags"
-ANNOTATION_METADATA_SCHEMA_KEY = "annotation_metadata_schema"
-ITEM_METADATA_SCHEMA_KEY = "item_metadata_schema"
-MASK_URL_KEY = "mask_url"
-INDEX_KEY = "index"
-SEGMENTATIONS_KEY = "segmentations"
-EMBEDDINGS_URL_KEY = "embeddings_url"
-JOB_ID_KEY = "job_id"
-MESSAGE_KEY = "message"
+WIDTH_KEY = "width"
+X_KEY = "x"
+Y_KEY = "y"

nucleus/dataset.py

@@ -1,10 +1,15 @@
-from typing import Any, Dict, List, Optional
+from typing import Any, Dict, List, Optional, Union
 
 import requests
 
-from nucleus.utils import format_dataset_item_response
+from nucleus.job import AsyncJob
+from nucleus.utils import (
+    convert_export_payload,
+    format_dataset_item_response,
+    serialize_and_write_to_presigned_url,
+)
 
-from .annotation import Annotation
+from .annotation import Annotation, check_all_annotation_paths_remote
 from .constants import (
     DATASET_ITEM_IDS_KEY,
     DATASET_LENGTH_KEY,
@@ -12,13 +17,24 @@ from .constants import (
     DATASET_NAME_KEY,
     DATASET_SLICES_KEY,
     DEFAULT_ANNOTATION_UPDATE_MODE,
+    EXPORTED_ROWS,
+    JOB_ID_KEY,
     NAME_KEY,
     REFERENCE_IDS_KEY,
+    REQUEST_ID_KEY,
+    UPDATE_KEY,
+)
+from .dataset_item import (
+    DatasetItem,
+    check_all_paths_remote,
+    check_for_duplicate_reference_ids,
 )
-from .dataset_item import DatasetItem
 from .payload_constructor import construct_model_run_creation_payload
 
 
+WARN_FOR_LARGE_UPLOAD = 50000
+
+
 class Dataset:
     """
     Nucleus Dataset. You can append images with metadata to your dataset,
@@ -26,7 +42,11 @@ class Dataset:
     compare model performance on you data.
     """
 
-    def __init__(self, dataset_id: str, client):
+    def __init__(
+        self,
+        dataset_id: str,
+        client: "NucleusClient",  # type:ignore # noqa: F821
+    ):
         self.id = dataset_id
         self._client = client
 
@@ -129,7 +149,8 @@ class Dataset:
         annotations: List[Annotation],
         update: Optional[bool] = DEFAULT_ANNOTATION_UPDATE_MODE,
         batch_size: int = 5000,
-    ) -> dict:
+        asynchronous: bool = False,
+    ) -> Union[Dict[str, Any], AsyncJob]:
         """
         Uploads ground truth annotations for a given dataset.
         :param annotations: ground truth annotations for a given dataset to upload
@@ -142,6 +163,19 @@ class Dataset:
             "ignored_items": int,
         }
         """
+        if asynchronous:
+            check_all_annotation_paths_remote(annotations)
+
+            request_id = serialize_and_write_to_presigned_url(
+                annotations, self.id, self._client
+            )
+            response = self._client.make_request(
+                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 self._client.annotate_dataset(
             self.id, annotations, update=update, batch_size=batch_size
         )
@@ -160,16 +194,18 @@ class Dataset:
     def append(
         self,
         dataset_items: List[DatasetItem],
-        force: Optional[bool] = False,
+        update: Optional[bool] = False,
         batch_size: Optional[int] = 20,
-    ) -> dict:
+        asynchronous=False,
+    ) -> Union[dict, AsyncJob]:
         """
         Appends images with metadata (dataset items) to the dataset. Overwrites images on collision if forced.
 
         Parameters:
         :param dataset_items: items to upload
-        :param force: if True overwrites images on collision
+        :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.
         :return:
         {
             'dataset_id': str,
@@ -178,10 +214,31 @@ class Dataset:
             'ignored_items': int,
         }
         """
+        check_for_duplicate_reference_ids(dataset_items)
+
+        if len(dataset_items) > WARN_FOR_LARGE_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_paths_remote(dataset_items)
+            request_id = serialize_and_write_to_presigned_url(
+                dataset_items, self.id, self._client
+            )
+            response = self._client.make_request(
+                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 self._client.populate_dataset(
             self.id,
             dataset_items,
-            force=force,
+            update=update,
             batch_size=batch_size,
         )
 
@@ -272,3 +329,23 @@ class Dataset:
 
     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]]]]]:
+        """Returns a list of all DatasetItems and Annotations in this slice.
+
+        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 DatasetItem, containing a reference to the
+                item that was annotated.
+            * The other value is a dictionary containing all the annotations for this
+                dataset item, sorted by annotation type.
+        """
+        api_payload = self._client.make_request(
+            payload=None,
+            route=f"dataset/{self.id}/exportForTraining",
+            requests_command=requests.get,
+        )
+        return convert_export_payload(api_payload[EXPORTED_ROWS])

nucleus/dataset_item.py

@@ -1,7 +1,9 @@
+from collections import Counter
 import json
 import os.path
 from dataclasses import dataclass
-from typing import Optional
+from typing import Optional, Sequence
+from urllib.parse import urlparse
 
 from .constants import (
     DATASET_ITEM_ID_KEY,
@@ -21,8 +23,7 @@ class DatasetItem:
     metadata: Optional[dict] = None
 
     def __post_init__(self):
-        self.image_url = self.image_location
-        self.local = self._is_local_path(self.image_location)
+        self.local = is_local_path(self.image_location)
 
     @classmethod
     def from_json(cls, payload: dict):
@@ -36,16 +37,12 @@ class DatasetItem:
             metadata=payload.get(METADATA_KEY, {}),
         )
 
-    def _is_local_path(self, path: str) -> bool:
-        path_components = [comp.lower() for comp in path.split("/")]
-        return path_components[0] not in {"https:", "http:", "s3:", "gs:"}
-
     def local_file_exists(self):
-        return os.path.isfile(self.image_url)
+        return os.path.isfile(self.image_location)
 
     def to_payload(self) -> dict:
         payload = {
-            IMAGE_URL_KEY: self.image_url,
+            IMAGE_URL_KEY: self.image_location,
             METADATA_KEY: self.metadata or {},
         }
         if self.reference_id:
@@ -55,4 +52,33 @@ class DatasetItem:
         return payload
 
     def to_json(self) -> str:
-        return json.dumps(self.to_payload())
+        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):
+            raise ValueError(
+                f"All paths must be remote, but {item.image_location} is either "
+                "local, or a remote URL type that is not supported."
+            )
+
+
+def check_for_duplicate_reference_ids(dataset_items: Sequence[DatasetItem]):
+    ref_ids = []
+    for dataset_item in dataset_items:
+        if dataset_item.reference_id is not None:
+            ref_ids.append(dataset_item.reference_id)
+    if len(ref_ids) != len(set(ref_ids)):
+        duplicates = {
+            f"{key}": f"Count: {value}"
+            for key, value in Counter(ref_ids).items()
+        }
+        raise ValueError(
+            "Duplicate reference ids found among dataset_items: %s"
+            % duplicates
+        )

nucleus/errors.py

@@ -25,9 +25,22 @@ class DatasetItemRetrievalError(Exception):
 
 
 class NucleusAPIError(Exception):
-    def __init__(self, endpoint, command, response):
-        message = f"Tried to {command.__name__} {endpoint}, but received {response.status_code}: {response.reason}."
-        if hasattr(response, "text"):
-            if response.text:
-                message += f"\nThe detailed error is:\n{response.text}"
+    def __init__(
+        self, endpoint, command, requests_response=None, aiohttp_response=None
+    ):
+
+        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"):
+                if requests_response.text:
+                    message += (
+                        f"\nThe detailed error is:\n{requests_response.text}"
+                    )
+
+        if aiohttp_response is not None:
+            status, reason, data = aiohttp_response
+            message = f"Tried to {command.__name__} {endpoint}, but received {status}: {reason}."
+            if data:
+                message += f"\nThe detailed error is:\n{data}"
+
         super().__init__(message)

nucleus/job.py

@@ -0,0 +1,56 @@
+from dataclasses import dataclass
+import time
+from typing import Dict, List
+
+import requests
+
+JOB_POLLING_INTERVAL = 5
+
+
+@dataclass
+class AsyncJob:
+    id: str
+    client: "NucleusClient"  # type: ignore # noqa: F821
+
+    def status(self) -> Dict[str, str]:
+        return self.client.make_request(
+            payload={},
+            route=f"job/{self.id}",
+            requests_command=requests.get,
+        )
+
+    def errors(self) -> List[str]:
+        return self.client.make_request(
+            payload={},
+            route=f"job/{self.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)
+
+
+class JobError(Exception):
+    def __init__(self, job_status: Dict[str, str], job: AsyncJob):
+        final_status_message = job_status["message"]
+        final_status = job_status["status"]
+        message = (
+            f"The job reported a final status of {final_status} "
+            "This could, however, mean a partial success with some successes and some failures. "
+            f"The final status message was: {final_status_message} \n"
+            f"For more detailed error messages you can call {str(job)}.errors()"
+        )
+        super().__init__(message)

nucleus/model.py

@@ -45,6 +45,7 @@ class Model:
             Union[BoxPrediction, PolygonPrediction, SegmentationPrediction]
         ],
         metadata: Optional[Dict] = None,
+        asynchronous: bool = False,
     ) -> ModelRun:
         payload: dict = {
             NAME_KEY: name,
@@ -56,6 +57,6 @@ class Model:
             dataset.id, payload
         )
 
-        model_run.predict(predictions)
+        model_run.predict(predictions, asynchronous=asynchronous)
 
         return model_run

nucleus/model_run.py

@@ -1,10 +1,18 @@
-from typing import Dict, Optional, List, Union, Type
+from typing import Dict, List, Optional, Type, Union
+
+from nucleus.annotation import check_all_annotation_paths_remote
+from nucleus.job import AsyncJob
+from nucleus.utils import serialize_and_write_to_presigned_url
+
 from .constants import (
     ANNOTATIONS_KEY,
-    DEFAULT_ANNOTATION_UPDATE_MODE,
     BOX_TYPE,
+    DEFAULT_ANNOTATION_UPDATE_MODE,
+    JOB_ID_KEY,
     POLYGON_TYPE,
+    REQUEST_ID_KEY,
     SEGMENTATION_TYPE,
+    UPDATE_KEY,
 )
 from .prediction import (
     BoxPrediction,
@@ -19,12 +27,13 @@ class ModelRun:
     Having an open model run is a prerequisite for uploading predictions to your dataset.
     """
 
-    def __init__(self, model_run_id: str, client):
+    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
 
     def __repr__(self):
-        return f"ModelRun(model_run_id='{self.model_run_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:
@@ -84,7 +93,8 @@ class ModelRun:
             Union[BoxPrediction, PolygonPrediction, SegmentationPrediction]
         ],
         update: Optional[bool] = DEFAULT_ANNOTATION_UPDATE_MODE,
-    ) -> dict:
+        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]],
@@ -95,7 +105,20 @@ class ModelRun:
             "predictions_ignored": int,
         }
         """
-        return self._client.predict(self.model_run_id, annotations, update)
+        if asynchronous:
+            check_all_annotation_paths_remote(annotations)
+
+            request_id = serialize_and_write_to_presigned_url(
+                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)
+        else:
+            return self._client.predict(self.model_run_id, annotations, update)
 
     def iloc(self, i: int):
         """

nucleus/payload_constructor.py

@@ -17,7 +17,7 @@ from .constants import (
     REFERENCE_ID_KEY,
     ANNOTATIONS_KEY,
     ITEMS_KEY,
-    FORCE_KEY,
+    UPDATE_KEY,
     MODEL_ID_KEY,
     ANNOTATION_METADATA_SCHEMA_KEY,
     SEGMENTATIONS_KEY,
@@ -34,7 +34,7 @@ def construct_append_payload(
     return (
         {ITEMS_KEY: items}
         if not force
-        else {ITEMS_KEY: items, FORCE_KEY: True}
+        else {ITEMS_KEY: items, UPDATE_KEY: True}
     )
 
 

nucleus/prediction.py

@@ -1,6 +1,7 @@
-from typing import Dict, Optional, List, Any
+from typing import Dict, Optional, List
 from .annotation import (
     BoxAnnotation,
+    Point,
     PolygonAnnotation,
     Segment,
     SegmentationAnnotation,
@@ -16,6 +17,7 @@ from .constants import (
     Y_KEY,
     WIDTH_KEY,
     HEIGHT_KEY,
+    CLASS_PDF_KEY,
     CONFIDENCE_KEY,
     VERTICES_KEY,
     ANNOTATIONS_KEY,
@@ -54,6 +56,7 @@ class BoxPrediction(BoxAnnotation):
         confidence: Optional[float] = None,
         annotation_id: Optional[str] = None,
         metadata: Optional[Dict] = None,
+        class_pdf: Optional[Dict] = None,
     ):
         super().__init__(
             label=label,
@@ -67,11 +70,14 @@ class BoxPrediction(BoxAnnotation):
             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
 
@@ -89,6 +95,7 @@ class BoxPrediction(BoxAnnotation):
             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),
         )
 
 
@@ -96,12 +103,13 @@ class PolygonPrediction(PolygonAnnotation):
     def __init__(
         self,
         label: str,
-        vertices: List[Any],
+        vertices: List[Point],
         reference_id: Optional[str] = None,
         item_id: Optional[str] = None,
         confidence: Optional[float] = None,
         annotation_id: Optional[str] = None,
         metadata: Optional[Dict] = None,
+        class_pdf: Optional[Dict] = None,
     ):
         super().__init__(
             label=label,
@@ -112,11 +120,14 @@ class PolygonPrediction(PolygonAnnotation):
             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
 
@@ -125,10 +136,13 @@ class PolygonPrediction(PolygonAnnotation):
         geometry = payload.get(GEOMETRY_KEY, {})
         return cls(
             label=payload.get(LABEL_KEY, 0),
-            vertices=geometry.get(VERTICES_KEY, []),
+            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),
             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),
         )