Merge pull request #30 from scaleapi/jihan-idempotency

Adds annotation_id as optional parameter to annotations and prediction classes. This will be used to uniquely identify each annotation/prediction within its dataset item. Also adds update as an optional boolean parameter when uploading predictions and annotations. Setting update=False, the default behavior, will skip uploading annotations to dataset images that already have annotations with the same annotation_id. Setting update=True will overwrite such annotations.

Author: jihan-yin

README.md

@@ -92,8 +92,8 @@ item = dataset.loc("dataset_item_id")
 Upload groundtruth annotations for the items in your dataset.
 Box2DAnnotation has same format as https://dashboard.scale.com/nucleus/docs/api#add-ground-truth
 ```python
-annotation_1 = BoxAnnotation(reference_id="1", label="label", x=0, y=0, width=10, height=10, metadata={})
-annotation_2 = BoxAnnotation(reference_id="2", label="label", x=0, y=0, width=10, height=10, metadata={})
+annotation_1 = BoxAnnotation(reference_id="1", label="label", x=0, y=0, width=10, height=10, annotation_id="ann_1", metadata={})
+annotation_2 = BoxAnnotation(reference_id="2", label="label", x=0, y=0, width=10, height=10, annotation_id="ann_2", metadata={})
 response = dataset.annotate([annotation_1, annotation_2])
 ```
 
@@ -113,8 +113,8 @@ Returns the associated model_id, human-readable name of the run, status, and use
 Takes a list of Box2DPredictions within the payload, where Box2DPrediction
 is formulated as in https://dashboard.scale.com/nucleus/docs/api#upload-model-outputs
 ```python
-prediction_1 = BoxPrediction(reference_id="1", label="label", x=0, y=0, width=10, height=10, confidence=0.9)
-prediction_2 = BoxPrediction(reference_id="2", label="label", x=0, y=0, width=10, height=10, confidence=0.2)
+prediction_1 = BoxPrediction(reference_id="1", label="label", x=0, y=0, width=10, height=10, annotation_id="pred_1", confidence=0.9)
+prediction_2 = BoxPrediction(reference_id="2", label="label", x=0, y=0, width=10, height=10, annotation_id="pred_2", confidence=0.2)
 
 model_run = model.create_run(name="My Model Run", metadata={"timestamp": "121012401"}, dataset=dataset, predictions=[prediction_1, prediction_2])
 ```

nucleus/__init__.py

@@ -58,7 +58,7 @@
 from typing import List, Union, Dict, Callable, Any, Optional
 
 import tqdm
-import tqdm.notebook
+import tqdm.notebook as tqdm_notebook
 
 import grequests
 import requests
@@ -76,7 +76,7 @@
 from .upload_response import UploadResponse
 from .payload_constructor import (
     construct_append_payload,
-    construct_box_annotation_payload,
+    construct_annotation_payload,
     construct_model_creation_payload,
     construct_box_predictions_payload,
 )
@@ -94,7 +94,9 @@
     DATASET_ITEM_ID_KEY,
     SLICE_ID_KEY,
     ANNOTATIONS_PROCESSED_KEY,
+    ANNOTATIONS_IGNORED_KEY,
     PREDICTIONS_PROCESSED_KEY,
+    PREDICTIONS_IGNORED_KEY,
     STATUS_CODE_KEY,
     SUCCESS_STATUS_CODES,
     DATASET_NAME_KEY,
@@ -129,7 +131,7 @@ def __init__(self, api_key: str, use_notebook: bool = False):
         self.api_key = api_key
         self.tqdm_bar = tqdm.tqdm
         if use_notebook:
-            self.tqdm_bar = tqdm.notebook.tqdm
+            self.tqdm_bar = tqdm_notebook.tqdm
 
     def list_models(self) -> List[str]:
         """
@@ -462,13 +464,15 @@ def exception_handler(request, exception):
     def annotate_dataset(
         self,
         dataset_id: str,
-        annotations: List[DatasetItem],
+        annotations: List[Union[BoxAnnotation, PolygonAnnotation]],
+        update: bool,
         batch_size: int = 100,
     ):
         """
         Uploads ground truth annotations for a given dataset.
         :param dataset_id: id of the dataset
-        :param annotations: List[DatasetItem]
+        :param annotations: List[Union[BoxAnnotation, PolygonAnnotation]]
+        :param update: whether to update or ignore conflicting annotations
         :return: {"dataset_id: str, "annotations_processed": int}
         """
 
@@ -480,12 +484,13 @@ def annotate_dataset(
         agg_response = {
             DATASET_ID_KEY: dataset_id,
             ANNOTATIONS_PROCESSED_KEY: 0,
+            ANNOTATIONS_IGNORED_KEY: 0,
         }
 
         tqdm_batches = self.tqdm_bar(batches)
 
         for batch in tqdm_batches:
-            payload = construct_box_annotation_payload(batch)
+            payload = construct_annotation_payload(batch, update)
             response = self._make_request(
                 payload, f"dataset/{dataset_id}/annotate"
             )
@@ -495,6 +500,9 @@ def annotate_dataset(
                 agg_response[ANNOTATIONS_PROCESSED_KEY] += response[
                     ANNOTATIONS_PROCESSED_KEY
                 ]
+                agg_response[ANNOTATIONS_IGNORED_KEY] += response[
+                    ANNOTATIONS_IGNORED_KEY
+                ]
 
         return agg_response
 
@@ -570,39 +578,40 @@ def create_model_run(self, dataset_id: str, payload: dict) -> ModelRun:
     def predict(
         self,
         model_run_id: str,
-        payload: Dict[str, List[Union[BoxPrediction, PolygonPrediction]]],
+        annotations: List[Union[BoxPrediction, PolygonPrediction]],
+        update: bool,
         batch_size: int = 100,
     ):
         """
         Uploads model outputs as predictions for a model_run. Returns info about the upload.
-        :param payload:
-        {
-            "annotations": List[Box2DPrediction],
-        }
+        :param annotations: List[Union[BoxPrediction, PolygonPrediction]],
+        :param update: bool
         :return:
         {
             "dataset_id": str,
             "model_run_id": str,
-            "annotations_processed: int,
+            "predictions_processed": int,
+            "predictions_ignored": int,
         }
         """
-        predictions: List[Union[BoxPrediction, PolygonPrediction]] = payload[
-            ANNOTATIONS_KEY
-        ]
         batches = [
-            predictions[i : i + batch_size]
-            for i in range(0, len(predictions), batch_size)
+            annotations[i : i + batch_size]
+            for i in range(0, len(annotations), batch_size)
         ]
 
         agg_response = {
             MODEL_RUN_ID_KEY: model_run_id,
             PREDICTIONS_PROCESSED_KEY: 0,
+            PREDICTIONS_IGNORED_KEY: 0,
         }
 
         tqdm_batches = self.tqdm_bar(batches)
 
         for batch in tqdm_batches:
-            batch_payload = {ANNOTATIONS_KEY: batch}
+            batch_payload = construct_box_predictions_payload(
+                annotations,
+                update,
+            )
             response = self._make_request(
                 batch_payload, f"modelRun/{model_run_id}/predict"
             )
@@ -612,6 +621,9 @@ def predict(
                 agg_response[PREDICTIONS_PROCESSED_KEY] += response[
                     PREDICTIONS_PROCESSED_KEY
                 ]
+                agg_response[PREDICTIONS_IGNORED_KEY] += response[
+                    PREDICTIONS_IGNORED_KEY
+                ]
 
         return agg_response
         # return self._make_request(payload, f"modelRun/{model_run_id}/predict")
@@ -699,7 +711,7 @@ def predictions_ref_id(self, model_run_id: str, ref_id: str):
         :param reference_id: reference_id of a dataset item.
         :return:
         {
-            "annotations": List[Box2DPrediction],
+            "annotations": List[BoxPrediction],
         }
         """
         return self._make_request(
@@ -724,7 +736,7 @@ def predictions_iloc(self, model_run_id: str, i: int):
         :param i: absolute number of Dataset Item for a dataset corresponding to the model run.
         :return:
         {
-            "annotations": List[Box2DPrediction],
+            "annotations": List[BoxPrediction],
         }
         """
         return self._make_request(
@@ -753,7 +765,7 @@ def predictions_loc(self, model_run_id: str, dataset_item_id: str):
         :param dataset_item_id: dataset_item_id of a dataset item.
         :return:
         {
-            "annotations": List[Box2DPrediction],
+            "annotations": List[BoxPrediction],
         }
         """
         return self._make_request(

nucleus/annotation.py

@@ -1,6 +1,7 @@
 from enum import Enum
 from typing import Dict, Optional, Any, Union, List
 from .constants import (
+    ANNOTATION_ID_KEY,
     DATASET_ITEM_ID_KEY,
     REFERENCE_ID_KEY,
     METADATA_KEY,
@@ -32,8 +33,9 @@ def __init__(
         y: Union[float, int],
         width: Union[float, int],
         height: Union[float, int],
-        reference_id: str = None,
-        item_id: str = None,
+        reference_id: Optional[str] = None,
+        item_id: Optional[str] = None,
+        annotation_id: Optional[str] = None,
         metadata: Optional[Dict] = None,
     ):
         if bool(reference_id) == bool(item_id):
@@ -47,6 +49,7 @@ def __init__(
         self.height = height
         self.reference_id = reference_id
         self.item_id = item_id
+        self.annotation_id = annotation_id
         self.metadata = metadata if metadata else {}
 
     @classmethod
@@ -60,6 +63,7 @@ def from_json(cls, payload: dict):
             height=geometry.get(HEIGHT_KEY, 0),
             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),
             metadata=payload.get(METADATA_KEY, {}),
         )
 
@@ -74,6 +78,7 @@ def to_payload(self) -> dict:
                 HEIGHT_KEY: self.height,
             },
             REFERENCE_ID_KEY: self.reference_id,
+            ANNOTATION_ID_KEY: self.annotation_id,
             METADATA_KEY: self.metadata,
         }
 
@@ -87,8 +92,9 @@ def __init__(
         self,
         label: str,
         vertices: List[Any],
-        reference_id: str = None,
-        item_id: str = None,
+        reference_id: Optional[str] = None,
+        item_id: Optional[str] = None,
+        annotation_id: Optional[str] = None,
         metadata: Optional[Dict] = None,
     ):
         if bool(reference_id) == bool(item_id):
@@ -99,6 +105,7 @@ def __init__(
         self.vertices = vertices
         self.reference_id = reference_id
         self.item_id = item_id
+        self.annotation_id = annotation_id
         self.metadata = metadata if metadata else {}
 
     @classmethod
@@ -109,6 +116,7 @@ def from_json(cls, payload: dict):
             vertices=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),
             metadata=payload.get(METADATA_KEY, {}),
         )
 
@@ -118,8 +126,9 @@ def to_payload(self) -> dict:
             TYPE_KEY: POLYGON_TYPE,
             GEOMETRY_KEY: {VERTICES_KEY: self.vertices},
             REFERENCE_ID_KEY: self.reference_id,
+            ANNOTATION_ID_KEY: self.annotation_id,
             METADATA_KEY: self.metadata,
         }
 
     def __str__(self):
-        return str(self.to_payload())
+        return str(self.to_payload())

nucleus/constants.py

@@ -13,8 +13,13 @@
 ERROR_PAYLOAD = "error_payload"
 ERROR_CODES = "error_codes"
 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_UPDATE_KEY = "update"
+DEFAULT_ANNOTATION_UPDATE_MODE = False
 STATUS_CODE_KEY = "status_code"
 SUCCESS_STATUS_CODES = [200, 201]
 ERRORS_KEY = "errors"

nucleus/dataset.py

@@ -10,11 +10,11 @@
     REFERENCE_IDS_KEY,
     NAME_KEY,
     ITEM_KEY,
+    DEFAULT_ANNOTATION_UPDATE_MODE,
     ANNOTATIONS_KEY,
 )
 from .payload_constructor import construct_model_run_creation_payload
 
-
 class Dataset:
     """
     Nucleus Dataset. You can append images with metadata to your dataset,
@@ -88,6 +88,7 @@ def create_model_run(
     def annotate(
         self,
         annotations: List[Union[BoxAnnotation, PolygonAnnotation]],
+        update: Optional[bool] = DEFAULT_ANNOTATION_UPDATE_MODE,
         batch_size: int = 20,
     ) -> dict:
         """
@@ -103,7 +104,7 @@ def annotate(
         }
         """
         return self._client.annotate_dataset(
-            self.id, annotations, batch_size=batch_size
+            self.id, annotations, update=update, batch_size=batch_size
         )
 
     def ingest_tasks(self, task_ids: dict):
@@ -165,7 +166,7 @@ def refloc(self, reference_id: str) -> dict:
         :return:
         {
             "item": DatasetItem,
-            "annotations": List[Box2DAnnotation],
+            "annotations": List[Union[BoxAnnotation, PolygonAnnotation]],
         }
         """
         response = self._client.dataitem_ref_id(self.id, reference_id)

nucleus/model_run.py

@@ -1,5 +1,5 @@
 from typing import Optional, List, Dict, Any, Union
-from .constants import ANNOTATIONS_KEY
+from .constants import ANNOTATIONS_KEY, DEFAULT_ANNOTATION_UPDATE_MODE
 from .prediction import BoxPrediction, PolygonPrediction
 from .payload_constructor import construct_box_predictions_payload
 
@@ -61,30 +61,27 @@ def commit(self, payload: Optional[dict] = None) -> dict:
         return self._client.commit_model_run(self.model_run_id, payload)
 
     def predict(
-        self, annotations: List[Union[BoxPrediction, PolygonPrediction]]
+        self,
+        annotations: List[Union[BoxPrediction, PolygonPrediction]],
+        update: Optional[bool] = DEFAULT_ANNOTATION_UPDATE_MODE,
     ) -> dict:
         """
         Uploads model outputs as predictions for a model_run. Returns info about the upload.
         :param annotations: List[Union[BoxPrediction, PolygonPrediction]],
         :return:
         {
-            "dataset_id": str,
             "model_run_id": str,
-            "annotations_processed: int,
+            "predictions_processed": int,
+            "predictions_ignored": int,
         }
         """
-        payload: Dict[str, List[Any]] = construct_box_predictions_payload(
-            annotations
-        )
-        return self._client.predict(self.model_run_id, payload)
+        return self._client.predict(self.model_run_id, annotations, update)
 
     def iloc(self, i: int):
         """
         Returns Model Run Info For Dataset Item by its number.
         :param i: absolute number of Dataset Item for a dataset corresponding to the model run.
-        :return:
-        {
-            "annotations": List[Union[BoxPrediction, PolygonPrediction]],
+        :return: List[Union[BoxPrediction, PolygonPrediction]],
         }
         """
         response = self._client.predictions_iloc(self.model_run_id, i)
@@ -94,10 +91,7 @@ def refloc(self, reference_id: str):
         """
         Returns Model Run Info For Dataset Item by its reference_id.
         :param reference_id: reference_id of a dataset item.
-        :return:
-        {
-            "annotations": List[Union[BoxPrediction, PolygonPrediction]],
-        }
+        :return: List[Union[BoxPrediction, PolygonPrediction]],
         """
         response = self._client.predictions_ref_id(
             self.model_run_id, reference_id