This just might work

Author: Diego Ardila

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 Any, Dict, List, Optional, Sequence, Union
+from nucleus.dataset_item import is_local_path
 
 from .constants import (
     ANNOTATION_ID_KEY,
@@ -13,6 +14,7 @@
     INDEX_KEY,
     ITEM_ID_KEY,
     LABEL_KEY,
+    MASK_TYPE,
     MASK_URL_KEY,
     METADATA_KEY,
     POLYGON_TYPE,
@@ -108,6 +110,7 @@ def from_json(cls, payload: dict):
 
     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,
@@ -206,3 +209,14 @@ def to_payload(self) -> dict:
             ANNOTATION_ID_KEY: self.annotation_id,
             METADATA_KEY: self.metadata,
         }
+
+
+def check_all_annotation_paths_remote(
+    annotations: Sequence[Union[Annotation]],
+):
+    for annotation in annotations:
+        if hasattr(annotation, "mask_url"):
+            if is_local_path(getattr(annotation, "mask_url")):
+                raise ValueError(
+                    f"Found an annotation with a local path, which cannot be uploaded asynchronously. Use a remote path instead. {annotation}"
+                )

nucleus/constants.py

@@ -5,6 +5,7 @@
 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"

nucleus/dataset.py

@@ -1,4 +1,3 @@
-import uuid
 from typing import Any, Dict, List, Optional, Union
 
 import requests
@@ -24,7 +23,6 @@
 )
 from .dataset_item import (
     DatasetItem,
-    check_all_paths_remote,
     check_for_duplicate_reference_ids,
 )
 from .payload_constructor import construct_model_run_creation_payload
@@ -198,15 +196,8 @@ def append(
         check_for_duplicate_reference_ids(dataset_items)
 
         if asynchronous:
-            check_all_paths_remote(dataset_items)
-            request_id = uuid.uuid4().hex
-            response = self._client.make_request(
-                payload={},
-                route=f"dataset/{self.id}/signedUrl/{request_id}",
-                requests_command=requests.get,
-            )
-            serialize_and_write_to_presigned_url(
-                dataset_items, response["signed_url"]
+            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},

nucleus/model.py

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

nucleus/model_run.py

@@ -1,10 +1,17 @@
-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,
+    MASK_TYPE,
     POLYGON_TYPE,
-    SEGMENTATION_TYPE,
+    REQUEST_ID_KEY,
+    UPDATE_KEY,
 )
 from .prediction import (
     BoxPrediction,
@@ -84,7 +91,9 @@ def predict(
             Union[BoxPrediction, PolygonPrediction, SegmentationPrediction]
         ],
         update: Optional[bool] = DEFAULT_ANNOTATION_UPDATE_MODE,
-    ) -> dict:
+        asynchronous: bool = False,
+        dataset_id: Optional[str] = None,
+    ) -> 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 +104,24 @@ def predict(
             "predictions_ignored": int,
         }
         """
-        return self._client.predict(self.model_run_id, annotations, update)
+        if asynchronous:
+            check_all_annotation_paths_remote(annotations)
+
+            assert (
+                dataset_id is not None
+            ), "For now, you must pass a dataset id to predict for asynchronous uploads."
+
+            request_id = serialize_and_write_to_presigned_url(
+                annotations, 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"], self._client)
+        else:
+            return self._client.predict(self.model_run_id, annotations, update)
 
     def iloc(self, i: int):
         """
@@ -153,7 +179,7 @@ def _format_prediction_response(
         ] = {
             BOX_TYPE: BoxPrediction,
             POLYGON_TYPE: PolygonPrediction,
-            SEGMENTATION_TYPE: SegmentationPrediction,
+            MASK_TYPE: SegmentationPrediction,
         }
         for type_key in annotation_payload:
             type_class = type_key_to_class[type_key]

nucleus/utils.py

@@ -2,6 +2,7 @@
 
 
 import io
+import uuid
 from typing import IO, Dict, List, Sequence, Union
 
 import requests
@@ -104,9 +105,19 @@ def upload_to_presigned_url(presigned_url: str, file_pointer: IO):
 
 
 def serialize_and_write_to_presigned_url(
-    upload_units: Sequence[Union[DatasetItem, Annotation]], presigned_url
+    upload_units: Sequence[Union["DatasetItem", Annotation]],
+    dataset_id: str,
+    client,
 ):
+    request_id = uuid.uuid4().hex
+    response = client.make_request(
+        payload={},
+        route=f"dataset/{dataset_id}/signedUrl/{request_id}",
+        requests_command=requests.get,
+    )
+
     strio = io.StringIO()
     serialize_and_write(upload_units, strio)
     strio.seek(0)
-    upload_to_presigned_url(presigned_url, strio)
+    upload_to_presigned_url(response["presigned_url"], strio)
+    return request_id