work in progress refactor of annotation upload

Author: Ubuntu

nucleus/__init__.py

@@ -460,93 +460,6 @@ def populate_dataset(
             dataset_items, batch_size=batch_size, update=update
         )
 
-    def annotate_dataset(
-        self,
-        dataset_id: str,
-        annotations: Sequence[
-            Union[
-                BoxAnnotation,
-                PolygonAnnotation,
-                CuboidAnnotation,
-                CategoryAnnotation,
-                MultiCategoryAnnotation,
-                SegmentationAnnotation,
-            ]
-        ],
-        update: bool,
-        batch_size: int = 5000,
-    ) -> Dict[str, object]:
-        # TODO: deprecate in favor of Dataset.annotate invocation
-
-        # Split payload into segmentations and Box/Polygon
-        segmentations = [
-            ann
-            for ann in annotations
-            if isinstance(ann, SegmentationAnnotation)
-        ]
-        other_annotations = [
-            ann
-            for ann in annotations
-            if not isinstance(ann, SegmentationAnnotation)
-        ]
-
-        batches = [
-            other_annotations[i : i + batch_size]
-            for i in range(0, len(other_annotations), batch_size)
-        ]
-
-        semseg_batches = [
-            segmentations[i : i + batch_size]
-            for i in range(0, len(segmentations), batch_size)
-        ]
-
-        agg_response = {
-            DATASET_ID_KEY: dataset_id,
-            ANNOTATIONS_PROCESSED_KEY: 0,
-            ANNOTATIONS_IGNORED_KEY: 0,
-            ERRORS_KEY: [],
-        }
-
-        total_batches = len(batches) + len(semseg_batches)
-
-        tqdm_batches = self.tqdm_bar(batches)
-
-        with self.tqdm_bar(total=total_batches) as pbar:
-            for batch in tqdm_batches:
-                payload = construct_annotation_payload(batch, update)
-                response = self.make_request(
-                    payload, f"dataset/{dataset_id}/annotate"
-                )
-                pbar.update(1)
-                if STATUS_CODE_KEY in response:
-                    agg_response[ERRORS_KEY] = response
-                else:
-                    agg_response[ANNOTATIONS_PROCESSED_KEY] += response[
-                        ANNOTATIONS_PROCESSED_KEY
-                    ]
-                    agg_response[ANNOTATIONS_IGNORED_KEY] += response[
-                        ANNOTATIONS_IGNORED_KEY
-                    ]
-                    agg_response[ERRORS_KEY] += response[ERRORS_KEY]
-
-            for s_batch in semseg_batches:
-                payload = construct_segmentation_payload(s_batch, update)
-                response = self.make_request(
-                    payload, f"dataset/{dataset_id}/annotate_segmentation"
-                )
-                pbar.update(1)
-                if STATUS_CODE_KEY in response:
-                    agg_response[ERRORS_KEY] = response
-                else:
-                    agg_response[ANNOTATIONS_PROCESSED_KEY] += response[
-                        ANNOTATIONS_PROCESSED_KEY
-                    ]
-                    agg_response[ANNOTATIONS_IGNORED_KEY] += response[
-                        ANNOTATIONS_IGNORED_KEY
-                    ]
-
-        return agg_response
-
     @deprecated(msg="Use Dataset.ingest_tasks instead")
     def ingest_tasks(self, dataset_id: str, payload: dict):
         dataset = self.get_dataset(dataset_id)

nucleus/annotation.py

@@ -1,4 +1,5 @@
 import json
+import os
 from dataclasses import dataclass, field
 from enum import Enum
 from typing import Dict, List, Optional, Sequence, Type, Union
@@ -70,6 +71,9 @@ def to_json(self) -> str:
         """Serializes annotation object to schematized JSON string."""
         return json.dumps(self.to_payload(), allow_nan=False)
 
+    def has_local_files(self) -> bool:
+        return False
+
 
 @dataclass  # pylint: disable=R0902
 class BoxAnnotation(Annotation):  # pylint: disable=R0902
@@ -578,6 +582,13 @@ def to_payload(self) -> dict:
 
         return payload
 
+    def has_files(self) -> bool:
+        if is_local_path(self.mask_url):
+            if not os.path.isfile(self.mask_url):
+                raise Exception(f"Mask file {self.mask_url} does not exist.")
+            return True
+        return False
+
 
 class AnnotationTypes(Enum):
     BOX = BOX_TYPE
@@ -737,12 +748,12 @@ def is_local_path(path: str) -> bool:
 
 
 def check_all_mask_paths_remote(
-    annotations: Sequence[Union[Annotation]],
+    annotations: Sequence[Annotation],
 ):
     for annotation in annotations:
         if hasattr(annotation, MASK_URL_KEY):
             if is_local_path(getattr(annotation, MASK_URL_KEY)):
                 raise ValueError(
                     "Found an annotation with a local path, which is not currently"
-                    f"supported. Use a remote path instead. {annotation}"
+                    f"supported for asynchronous upload. Use a remote path instead, or try synchronous upload. {annotation}"
                 )

nucleus/async_utils.py

@@ -5,6 +5,7 @@
 
 import aiohttp
 import nest_asyncio
+from tqdm import tqdm
 
 from nucleus.constants import DEFAULT_NETWORK_TIMEOUT_SEC
 from nucleus.errors import NucleusAPIError
@@ -27,6 +28,16 @@ class FileFormField:
 FileFormData = Sequence[FileFormField]
 
 
+async def gather_with_concurrency(n, *tasks):
+    semaphore = asyncio.Semaphore(n)
+
+    async def sem_task(task):
+        async with semaphore:
+            return await task
+
+    return await asyncio.gather(*(sem_task(task) for task in tasks))
+
+
 class FormDataContextHandler:
     """A context handler for file form data that handles closing all files in a request.
 
@@ -85,6 +96,8 @@ def make_many_form_data_requests_concurrently(
     client: "NucleusClient",
     requests: Sequence[FormDataContextHandler],
     route: str,
+    progressbar: tqdm,
+    concurrency: int = 30,
 ):
     """
     Makes an async post request with form data to a Nucleus endpoint.
@@ -97,14 +110,18 @@ def make_many_form_data_requests_concurrently(
     """
     loop = get_event_loop()
     return loop.run_until_complete(
-        form_data_request_helper(client, requests, route)
+        form_data_request_helper(
+            client, requests, route, progressbar, concurrency
+        )
     )
 
 
 async def form_data_request_helper(
     client: "NucleusClient",
     requests: Sequence[FormDataContextHandler],
     route: str,
+    progressbar: tqdm,
+    concurrency: int = 30,
 ):
     """
     Makes an async post request with files to a Nucleus endpoint.
@@ -123,18 +140,20 @@ async def form_data_request_helper(
                     request=request,
                     route=route,
                     session=session,
+                    progressbar=progressbar,
                 )
             )
             for request in requests
         ]
-        return await asyncio.gather(*tasks)
+        return await gather_with_concurrency(concurrency, *tasks)
 
 
 async def _post_form_data(
     client: "NucleusClient",
     request: FormDataContextHandler,
     route: str,
     session: aiohttp.ClientSession,
+    progressbar: tqdm,
 ):
     """
     Makes an async post request with files to a Nucleus endpoint.
@@ -175,7 +194,7 @@ async def _post_form_data(
 
                 if response.status == 503:
                     raise TimeoutError(
-                        "The request to upload your max is timing out, please lower the batch size."
+                        "The request to upload your max is timing out, please lower local_files_per_upload_request in your api call."
                     )
 
                 if not response.ok:
@@ -188,5 +207,5 @@ async def _post_form_data(
                             data,
                         ),
                     )
-
+                progressbar.update(1)
                 return data