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

nucleus/annotation.py

@@ -1,25 +1,29 @@
+import json
 from dataclasses import dataclass
 from enum import Enum
-from typing import Dict, Optional, Any, Union, List
+from typing import Dict, List, Optional, Sequence, Union
+from nucleus.dataset_item import is_local_path
+
 from .constants import (
     ANNOTATION_ID_KEY,
+    ANNOTATIONS_KEY,
+    BOX_TYPE,
     DATASET_ITEM_ID_KEY,
-    REFERENCE_ID_KEY,
-    METADATA_KEY,
-    X_KEY,
-    Y_KEY,
-    WIDTH_KEY,
-    HEIGHT_KEY,
     GEOMETRY_KEY,
-    BOX_TYPE,
-    POLYGON_TYPE,
+    HEIGHT_KEY,
+    INDEX_KEY,
+    ITEM_ID_KEY,
     LABEL_KEY,
+    MASK_TYPE,
+    MASK_URL_KEY,
+    METADATA_KEY,
+    POLYGON_TYPE,
+    REFERENCE_ID_KEY,
     TYPE_KEY,
     VERTICES_KEY,
-    ITEM_ID_KEY,
-    MASK_URL_KEY,
-    INDEX_KEY,
-    ANNOTATIONS_KEY,
+    WIDTH_KEY,
+    X_KEY,
+    Y_KEY,
 )
 
 
@@ -42,6 +46,15 @@ class Annotation:
         else:
             return SegmentationAnnotation.from_json(payload)
 
+    def to_payload(self):
+        raise NotImplementedError(
+            "For serialization, use a specific subclass (i.e. SegmentationAnnotation), "
+            "not the base annotation class."
+        )
+
+    def to_json(self) -> str:
+        return json.dumps(self.to_payload(), allow_nan=False)
+
 
 @dataclass
 class Segment:
@@ -97,6 +110,7 @@ class SegmentationAnnotation(Annotation):
 
     def to_payload(self) -> dict:
         payload = {
+            TYPE_KEY: MASK_TYPE,
             MASK_URL_KEY: self.mask_url,
             ANNOTATIONS_KEY: [ann.to_payload() for ann in self.annotations],
             ANNOTATION_ID_KEY: self.annotation_id,
@@ -160,11 +174,23 @@ class BoxAnnotation(Annotation):  # pylint: disable=R0902
         }
 
 
-# TODO: Add Generic type for 2D point
+@dataclass
+class Point:
+    x: float
+    y: float
+
+    @classmethod
+    def from_json(cls, payload: Dict[str, float]):
+        return cls(payload[X_KEY], payload[Y_KEY])
+
+    def to_payload(self) -> dict:
+        return {X_KEY: self.x, Y_KEY: self.y}
+
+
 @dataclass
 class PolygonAnnotation(Annotation):
     label: str
-    vertices: List[Any]
+    vertices: List[Point]
     reference_id: Optional[str] = None
     item_id: Optional[str] = None
     annotation_id: Optional[str] = None
@@ -173,13 +199,28 @@ class PolygonAnnotation(Annotation):
     def __post_init__(self):
         self._check_ids()
         self.metadata = self.metadata if self.metadata else {}
+        if len(self.vertices) > 0:
+            if not hasattr(self.vertices[0], X_KEY) or not hasattr(
+                self.vertices[0], "to_payload"
+            ):
+                try:
+                    self.vertices = [
+                        Point(x=vertex[X_KEY], y=vertex[Y_KEY])
+                        for vertex in self.vertices
+                    ]
+                except KeyError as ke:
+                    raise ValueError(
+                        "Use a point object to pass in vertices. For example, vertices=[nucleus.Point(x=1, y=2)]"
+                    ) from ke
 
     @classmethod
     def from_json(cls, payload: dict):
         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),
             annotation_id=payload.get(ANNOTATION_ID_KEY, None),
@@ -187,11 +228,25 @@ class PolygonAnnotation(Annotation):
         )
 
     def to_payload(self) -> dict:
-        return {
+        payload = {
             LABEL_KEY: self.label,
             TYPE_KEY: POLYGON_TYPE,
-            GEOMETRY_KEY: {VERTICES_KEY: self.vertices},
+            GEOMETRY_KEY: {
+                VERTICES_KEY: [_.to_payload() for _ in self.vertices]
+            },
             REFERENCE_ID_KEY: self.reference_id,
             ANNOTATION_ID_KEY: self.annotation_id,
             METADATA_KEY: self.metadata,
         }
+        return payload
+
+
+def check_all_annotation_paths_remote(
+    annotations: Sequence[Union[Annotation]],
+):
+    for annotation in annotations:
+        if hasattr(annotation, MASK_URL_KEY):
+            if is_local_path(getattr(annotation, MASK_URL_KEY)):
+                raise ValueError(
+                    f"Found an annotation with a local path, which cannot be uploaded asynchronously. Use a remote path instead. {annotation}"
+                )

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,23 +1,40 @@
-from typing import List, Dict, Any, Optional
+from typing import Any, Dict, List, Optional, Union
 
-from nucleus.utils import format_dataset_item_response
-from .dataset_item import DatasetItem
-from .annotation import (
-    Annotation,
+import requests
+
+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, check_all_annotation_paths_remote
 from .constants import (
-    DATASET_NAME_KEY,
+    DATASET_ITEM_IDS_KEY,
+    DATASET_LENGTH_KEY,
     DATASET_MODEL_RUNS_KEY,
+    DATASET_NAME_KEY,
     DATASET_SLICES_KEY,
-    DATASET_LENGTH_KEY,
-    DATASET_ITEM_IDS_KEY,
-    REFERENCE_IDS_KEY,
-    NAME_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 .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,
@@ -25,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
 
@@ -58,6 +79,25 @@ 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.
+
+        If you have pandas installed, you can create a pandas dataframe using
+
+        pandas.Dataframe(dataset.autotag_scores(autotag_name))
+
+        :return: dictionary of the form
+            {'ref_ids': List[str],
+             'datset_item_ids': List[str],
+             'score': List[float]}
+        """
+        response = self._client.make_request(
+            payload={},
+            route=f"autotag/{self.id}/{autotag_name}/{for_scores_greater_than}",
+            requests_command=requests.get,
+        )
+        return response
+
     def info(self) -> dict:
         """
         Returns information about existing dataset
@@ -109,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
@@ -122,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
         )
@@ -140,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,
@@ -158,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,
         )
 
@@ -221,9 +298,9 @@ class Dataset:
 
         :return: new Slice object
         """
-        if dataset_item_ids and reference_ids:
+        if bool(dataset_item_ids) == bool(reference_ids):
             raise Exception(
-                "You cannot specify both dataset_item_ids and reference_ids"
+                "You must specify exactly one of dataset_item_ids or reference_ids."
             )
         payload: Dict[str, Any] = {NAME_KEY: name}
         if dataset_item_ids:
@@ -252,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,12 +1,16 @@
-from dataclasses import dataclass
+from collections import Counter
+import json
 import os.path
-from typing import Optional
+from dataclasses import dataclass
+from typing import Optional, Sequence
+from urllib.parse import urlparse
+
 from .constants import (
+    DATASET_ITEM_ID_KEY,
     IMAGE_URL_KEY,
     METADATA_KEY,
-    REFERENCE_ID_KEY,
     ORIGINAL_IMAGE_URL_KEY,
-    DATASET_ITEM_ID_KEY,
+    REFERENCE_ID_KEY,
 )
 
 
@@ -19,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):
@@ -34,20 +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 not (
-            "https:" in path_components
-            or "http:" in path_components
-            or "s3:" in path_components
-        )
-
     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,3 +50,35 @@ class DatasetItem:
         if self.item_id:
             payload[DATASET_ITEM_ID_KEY] = self.item_id
         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):
+            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

@@ -22,3 +22,25 @@ class DatasetItemRetrievalError(Exception):
     def __init__(self, message="Could not retrieve dataset items"):
         self.message = message
         super().__init__(self.message)
+
+
+class NucleusAPIError(Exception):
+    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