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

nucleus/__init__.py

@@ -50,90 +50,83 @@ confidence      |   float   |   The optional confidence level of this annotation
 geometry        |   dict    |   Representation of the bounding box in the Box2DGeometry format.\n
 metadata        |   dict    |   An arbitrary metadata blob for the annotation.\n
 """
-__version__ = "0.1.0"
-
+import asyncio
 import json
 import logging
-import warnings
 import os
-from typing import List, Union, Dict, Callable, Any, Optional
+from typing import Any, Dict, List, Optional, Union
 
+import aiohttp
+import pkg_resources
+import requests
 import tqdm
 import tqdm.notebook as tqdm_notebook
 
-import grequests
-import requests
-from requests.adapters import HTTPAdapter
-
-# pylint: disable=E1101
-# TODO: refactor to reduce this file to under 1000 lines.
-# pylint: disable=C0302
-from requests.packages.urllib3.util.retry import Retry
-
-from .constants import REFERENCE_IDS_KEY, DATASET_ITEM_IDS_KEY
-from .dataset import Dataset
-from .dataset_item import DatasetItem
 from .annotation import (
     BoxAnnotation,
     PolygonAnnotation,
-    SegmentationAnnotation,
     Segment,
-)
-from .prediction import (
-    BoxPrediction,
-    PolygonPrediction,
-    SegmentationPrediction,
-)
-from .model_run import ModelRun
-from .slice import Slice
-from .upload_response import UploadResponse
-from .payload_constructor import (
-    construct_append_payload,
-    construct_annotation_payload,
-    construct_model_creation_payload,
-    construct_box_predictions_payload,
-    construct_segmentation_payload,
+    SegmentationAnnotation,
+    Point,
 )
 from .constants import (
-    NUCLEUS_ENDPOINT,
+    ANNOTATION_METADATA_SCHEMA_KEY,
+    ANNOTATIONS_IGNORED_KEY,
+    ANNOTATIONS_PROCESSED_KEY,
+    AUTOTAGS_KEY,
+    DATASET_ID_KEY,
+    DATASET_ITEM_IDS_KEY,
     DEFAULT_NETWORK_TIMEOUT_SEC,
-    ERRORS_KEY,
+    EMBEDDINGS_URL_KEY,
     ERROR_ITEMS,
     ERROR_PAYLOAD,
-    ITEMS_KEY,
-    ITEM_KEY,
+    ERRORS_KEY,
     IMAGE_KEY,
     IMAGE_URL_KEY,
-    DATASET_ID_KEY,
+    ITEM_METADATA_SCHEMA_KEY,
+    ITEMS_KEY,
     MODEL_RUN_ID_KEY,
-    DATASET_ITEM_ID_KEY,
-    SLICE_ID_KEY,
-    ANNOTATIONS_PROCESSED_KEY,
-    ANNOTATIONS_IGNORED_KEY,
-    PREDICTIONS_PROCESSED_KEY,
+    NAME_KEY,
+    NUCLEUS_ENDPOINT,
     PREDICTIONS_IGNORED_KEY,
+    PREDICTIONS_PROCESSED_KEY,
+    REFERENCE_IDS_KEY,
+    SLICE_ID_KEY,
     STATUS_CODE_KEY,
-    SUCCESS_STATUS_CODES,
-    DATASET_NAME_KEY,
-    DATASET_MODEL_RUNS_KEY,
-    DATASET_SLICES_KEY,
-    DATASET_LENGTH_KEY,
-    NAME_KEY,
-    ANNOTATIONS_KEY,
-    AUTOTAGS_KEY,
-    ANNOTATION_METADATA_SCHEMA_KEY,
-    ITEM_METADATA_SCHEMA_KEY,
-    FORCE_KEY,
-    EMBEDDINGS_URL_KEY,
+    UPDATE_KEY,
 )
-from .model import Model
+from .dataset import Dataset
+from .dataset_item import DatasetItem
 from .errors import (
+    DatasetItemRetrievalError,
     ModelCreationError,
     ModelRunCreationError,
-    DatasetItemRetrievalError,
     NotFoundError,
     NucleusAPIError,
 )
+from .model import Model
+from .model_run import ModelRun
+from .payload_constructor import (
+    construct_annotation_payload,
+    construct_append_payload,
+    construct_box_predictions_payload,
+    construct_model_creation_payload,
+    construct_segmentation_payload,
+)
+from .prediction import (
+    BoxPrediction,
+    PolygonPrediction,
+    SegmentationPrediction,
+)
+from .slice import Slice
+from .upload_response import UploadResponse
+
+# pylint: disable=E1101
+# TODO: refactor to reduce this file to under 1000 lines.
+# pylint: disable=C0302
+
+
+__version__ = pkg_resources.get_distribution("scale-nucleus").version
 
 logger = logging.getLogger(__name__)
 logging.basicConfig()
@@ -151,11 +144,16 @@ class NucleusClient:
         self,
         api_key: str,
         use_notebook: bool = False,
-        endpoint=NUCLEUS_ENDPOINT,
+        endpoint: str = None,
     ):
         self.api_key = api_key
         self.tqdm_bar = tqdm.tqdm
-        self.endpoint = endpoint
+        if endpoint is None:
+            self.endpoint = os.environ.get(
+                "NUCLEUS_ENDPOINT", NUCLEUS_ENDPOINT
+            )
+        else:
+            self.endpoint = endpoint
         self._use_notebook = use_notebook
         if use_notebook:
             self.tqdm_bar = tqdm_notebook.tqdm
@@ -228,13 +226,15 @@ class NucleusClient:
         """
         return Dataset(dataset_id, self)
 
-    def get_model_run(self, model_run_id: str) -> ModelRun:
+    def get_model_run(self, model_run_id: str, dataset_id: str) -> ModelRun:
         """
         Fetches a model_run for given id
         :param model_run_id: internally controlled model_run_id
+        :param dataset_id: the dataset id which may determine the prediction schema
+            for this model run if present on the dataset.
         :return: model_run
         """
-        return ModelRun(model_run_id, self)
+        return ModelRun(model_run_id, dataset_id, self)
 
     def delete_model_run(self, model_run_id: str):
         """
@@ -324,13 +324,13 @@ class NucleusClient:
         dataset_id: str,
         dataset_items: List[DatasetItem],
         batch_size: int = 100,
-        force: bool = False,
+        update: bool = False,
     ):
         """
         Appends images to a dataset with given dataset_id.
-        Overwrites images on collision if forced.
+        Overwrites images on collision if updated.
         :param dataset_id: id of a dataset
-        :param payload: { "items": List[DatasetItem], "force": bool }
+        :param payload: { "items": List[DatasetItem], "update": bool }
         :param local: flag if images are stored locally
         :param batch_size: size of the batch for long payload
         :return:
@@ -373,21 +373,24 @@ class NucleusClient:
         async_responses: List[Any] = []
 
         for batch in tqdm_local_batches:
-            payload = construct_append_payload(batch, force)
+            payload = construct_append_payload(batch, update)
             responses = self._process_append_requests_local(
-                dataset_id, payload, force
+                dataset_id, payload, update
             )
             async_responses.extend(responses)
 
         for batch in tqdm_remote_batches:
-            payload = construct_append_payload(batch, force)
+            payload = construct_append_payload(batch, update)
             responses = self._process_append_requests(
-                dataset_id, payload, force, batch_size, batch_size
+                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.json())
+            agg_response.update_response(response)
 
         return agg_response
 
@@ -395,25 +398,19 @@ class NucleusClient:
         self,
         dataset_id: str,
         payload: dict,
-        update: bool,
+        update: bool,  # TODO: understand how to pass this in.
         local_batch_size: int = 10,
-        size: int = 10,
     ):
-        def error(batch_items: dict) -> UploadResponse:
-            return UploadResponse(
-                {
-                    DATASET_ID_KEY: dataset_id,
-                    ERROR_ITEMS: len(batch_items),
-                    ERROR_PAYLOAD: batch_items,
-                }
-            )
-
-        def exception_handler(request, exception):
-            logger.error(exception)
-
-        def preprocess_payload(batch):
+        def get_files(batch):
             request_payload = [
-                (ITEMS_KEY, (None, json.dumps(batch), "application/json"))
+                (
+                    ITEMS_KEY,
+                    (
+                        None,
+                        json.dumps(batch, allow_nan=False),
+                        "application/json",
+                    ),
+                )
             ]
             for item in batch:
                 image = open(  # pylint: disable=R1732
@@ -430,27 +427,19 @@ class NucleusClient:
 
         items = payload[ITEMS_KEY]
         responses: List[Any] = []
-        request_payloads = []
+        files_per_request = []
         payload_items = []
         for i in range(0, len(items), local_batch_size):
             batch = items[i : i + local_batch_size]
-            batch_payload = preprocess_payload(batch)
-            request_payloads.append(batch_payload)
+            files_per_request.append(get_files(batch))
             payload_items.append(batch)
 
-        async_requests = [
-            self._make_grequest(
-                payload,
+        loop = asyncio.get_event_loop()
+        responses = loop.run_until_complete(
+            self.make_many_files_requests_asynchronously(
+                files_per_request,
                 f"dataset/{dataset_id}/append",
-                local=True,
             )
-            for payload in request_payloads
-        ]
-
-        async_responses = grequests.map(
-            async_requests,
-            exception_handler=exception_handler,
-            size=size,
         )
 
         def close_files(request_items):
@@ -460,69 +449,106 @@ class NucleusClient:
                     item[1][1].close()
 
         # don't forget to close all open files
-        for p in request_payloads:
+        for p in files_per_request:
             close_files(p)
 
-        # response object will be None if an error occurred
-        async_responses = [
-            response
-            if (response and response.status_code == 200)
-            else error(request_items)
-            for response, request_items in zip(async_responses, payload_items)
-        ]
-        responses.extend(async_responses)
-
         return responses
 
+    async def make_many_files_requests_asynchronously(
+        self, files_per_request, route
+    ):
+        """
+        Makes an async post request with files to a Nucleus endpoint.
+
+        :param files_per_request: A list of lists of tuples (name, (filename, file_pointer, content_type))
+           name will become the name by which the multer can build an array.
+        :param route: route for the request
+        :return: awaitable list(response)
+        """
+        async with aiohttp.ClientSession() as session:
+            tasks = [
+                asyncio.ensure_future(
+                    self._make_files_request(
+                        files=files, route=route, session=session
+                    )
+                )
+                for files in files_per_request
+            ]
+            return await asyncio.gather(*tasks)
+
+    async def _make_files_request(
+        self,
+        files,
+        route: str,
+        session: aiohttp.ClientSession,
+    ):
+        """
+        Makes an async post request with files to a Nucleus endpoint.
+
+        :param files: A list of tuples (name, (filename, file_pointer, file_type))
+        :param route: route for the request
+        :param session: Session to use for post.
+        :return: response
+        """
+        endpoint = f"{self.endpoint}/{route}"
+
+        logger.info("Posting to %s", endpoint)
+
+        form = aiohttp.FormData()
+
+        for file in files:
+            form.add_field(
+                name=file[0],
+                filename=file[1][0],
+                value=file[1][1],
+                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),
+                )
+
+            return data
+
     def _process_append_requests(
         self,
         dataset_id: str,
         payload: dict,
         update: bool,
         batch_size: int = 20,
-        size: int = 10,
     ):
-        def default_error(payload: dict) -> UploadResponse:
-            return UploadResponse(
-                {
-                    DATASET_ID_KEY: dataset_id,
-                    ERROR_ITEMS: len(payload[ITEMS_KEY]),
-                    ERROR_PAYLOAD: payload[ITEMS_KEY],
-                }
-            )
-
-        def exception_handler(request, exception):
-            logger.error(exception)
-
         items = payload[ITEMS_KEY]
         payloads = [
             # batch_size images per request
-            {ITEMS_KEY: items[i : i + batch_size], FORCE_KEY: update}
+            {ITEMS_KEY: items[i : i + batch_size], UPDATE_KEY: update}
             for i in range(0, len(items), batch_size)
         ]
 
-        async_requests = [
-            self._make_grequest(
+        return [
+            self.make_request(
                 payload,
                 f"dataset/{dataset_id}/append",
-                local=False,
             )
             for payload in payloads
         ]
 
-        async_responses = grequests.map(
-            async_requests, exception_handler=exception_handler, size=size
-        )
-
-        async_responses = [
-            response
-            if (response and response.status_code == 200)
-            else default_error(payload)
-            for response, payload in zip(async_responses, payloads)
-        ]
-
-        return async_responses
-
     def annotate_dataset(
         self,
         dataset_id: str,
@@ -672,7 +698,9 @@ class NucleusClient:
         if response.get(STATUS_CODE_KEY, None):
             raise ModelRunCreationError(response.get("error"))
 
-        return ModelRun(response[MODEL_RUN_ID_KEY], self)
+        return ModelRun(
+            response[MODEL_RUN_ID_KEY], dataset_id=dataset_id, client=self
+        )
 
     def predict(
         self,
@@ -1068,61 +1096,20 @@ class NucleusClient:
             requests_command=requests.delete,
         )
 
-    def _make_grequest(
-        self,
-        payload: dict,
-        route: str,
-        session=None,
-        requests_command: Callable = grequests.post,
-        local=True,
-    ):
-        """
-        makes a grequest to Nucleus endpoint
-        :param payload: file dict for multipart-formdata
-        :param route: route for the request
-        :param session: requests.session
-        :param requests_command: grequests.post, grequests.get, grequests.delete
-        :return: An async grequest object
-        """
-        adapter = HTTPAdapter(max_retries=Retry(total=3))
-        sess = requests.Session()
-        sess.mount("https://", adapter)
-        sess.mount("http://", adapter)
-
-        endpoint = f"{self.endpoint}/{route}"
-        logger.info("Posting to %s", endpoint)
-
-        if local:
-            post = requests_command(
-                endpoint,
-                session=sess,
-                files=payload,
-                auth=(self.api_key, ""),
-                timeout=DEFAULT_NETWORK_TIMEOUT_SEC,
-            )
-        else:
-            post = requests_command(
-                endpoint,
-                session=sess,
-                json=payload,
-                headers={"Content-Type": "application/json"},
-                auth=(self.api_key, ""),
-                timeout=DEFAULT_NETWORK_TIMEOUT_SEC,
-            )
-        return post
-
-    def _make_request_raw(
-        self, payload: dict, endpoint: str, requests_command=requests.post
-    ):
+    def make_request(
+        self, payload: dict, route: str, requests_command=requests.post
+    ) -> dict:
         """
-        Makes a request to Nucleus endpoint. This method returns the raw
-        requests.Response object which is useful for unit testing.
+        Makes a request to Nucleus endpoint and logs a warning if not
+        successful.
 
         :param payload: given payload
-        :param endpoint: endpoint + route for the request
+        :param route: route for the request
         :param requests_command: requests.post, requests.get, requests.delete
-        :return: response
+        :return: response JSON
         """
+        endpoint = f"{self.endpoint}/{route}"
+
         logger.info("Posting to %s", endpoint)
 
         response = requests_command(
@@ -1134,28 +1121,18 @@ class NucleusClient:
         )
         logger.info("API request has response code %s", response.status_code)
 
-        return response
-
-    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}"
-
-        response = self._make_request_raw(payload, endpoint, requests_command)
-
         if not response.ok:
             self.handle_bad_response(endpoint, requests_command, response)
 
         return response.json()
 
-    def handle_bad_response(self, endpoint, requests_command, response):
-        raise NucleusAPIError(endpoint, requests_command, response)
+    def handle_bad_response(
+        self,
+        endpoint,
+        requests_command,
+        requests_response=None,
+        aiohttp_response=None,
+    ):
+        raise NucleusAPIError(
+            endpoint, requests_command, requests_response, aiohttp_response
+        )

nucleus/annotation.py

@@ -1,7 +1,8 @@
 import json
 from dataclasses import dataclass
 from enum import Enum
-from typing import Any, Dict, List, Optional, Union
+from typing import Dict, List, Optional, Sequence, Union
+from nucleus.dataset_item import is_local_path
 
 from .constants import (
     ANNOTATION_ID_KEY,
@@ -13,6 +14,7 @@ from .constants import (
     INDEX_KEY,
     ITEM_ID_KEY,
     LABEL_KEY,
+    MASK_TYPE,
     MASK_URL_KEY,
     METADATA_KEY,
     POLYGON_TYPE,
@@ -51,7 +53,7 @@ class Annotation:
         )
 
     def to_json(self) -> str:
-        return json.dumps(self.to_payload())
+        return json.dumps(self.to_payload(), allow_nan=False)
 
 
 @dataclass
@@ -108,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,
@@ -171,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
@@ -184,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),
@@ -198,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}"
+                )