nucleus side working

Author: jihan-yin

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
@@ -69,14 +69,14 @@
 
 from .dataset import Dataset
 from .dataset_item import DatasetItem
-from .annotation import BoxAnnotation, PolygonAnnotation
+from .annotation import AnnotationItem, BoxAnnotation, PolygonAnnotation
 from .prediction import BoxPrediction, PolygonPrediction
 from .model_run import ModelRun
 from .slice import Slice
 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,
 )
@@ -129,7 +129,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 +462,14 @@ def exception_handler(request, exception):
     def annotate_dataset(
         self,
         dataset_id: str,
-        annotations: List[DatasetItem],
+        annotations: List[AnnotationItem],
+        edit_mode: str,
         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[AnnotationItem]
         :return: {"dataset_id: str, "annotations_processed": int}
         """
 
@@ -485,7 +486,7 @@ def annotate_dataset(
         tqdm_batches = self.tqdm_bar(batches)
 
         for batch in tqdm_batches:
-            payload = construct_box_annotation_payload(batch)
+            payload = construct_annotation_payload(batch, edit_mode)
             response = self._make_request(
                 payload, f"dataset/{dataset_id}/annotate"
             )

nucleus/annotation.py

@@ -1,6 +1,8 @@
 from enum import Enum
 from typing import Dict, Optional, Any, Union, List
 from .constants import (
+    ANNOTATIONS_KEY,
+    ANNOTATION_ID_KEY,
     DATASET_ITEM_ID_KEY,
     REFERENCE_ID_KEY,
     METADATA_KEY,
@@ -32,22 +34,14 @@ def __init__(
         y: Union[float, int],
         width: Union[float, int],
         height: Union[float, int],
-        reference_id: str = None,
-        item_id: str = None,
-        annotation_id: str = None,
+        annotation_id: Optional[str] = None,
         metadata: Optional[Dict] = None,
     ):
-        if bool(reference_id) == bool(item_id):
-            raise Exception(
-                "You must specify either a reference_id or an item_id for an annotation."
-            )
         self.label = label
         self.x = x
         self.y = y
         self.width = width
         self.height = height
-        self.reference_id = reference_id
-        self.item_id = item_id
         self.annotation_id = annotation_id
         self.metadata = metadata if metadata else {}
 
@@ -60,9 +54,7 @@ def from_json(cls, payload: dict):
             y=geometry.get(Y_KEY, 0),
             width=geometry.get(WIDTH_KEY, 0),
             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, None),
+            annotation_id=payload.get(ANNOTATION_ID_KEY, None),
             metadata=payload.get(METADATA_KEY, {}),
         )
 
@@ -76,8 +68,7 @@ def to_payload(self) -> dict:
                 WIDTH_KEY: self.width,
                 HEIGHT_KEY: self.height,
             },
-            REFERENCE_ID_KEY: self.reference_id,
-            ANNOTATION_ID: self.annotation_id,
+            ANNOTATION_ID_KEY: self.annotation_id,
             METADATA_KEY: self.metadata,
         }
 
@@ -91,19 +82,11 @@ def __init__(
         self,
         label: str,
         vertices: List[Any],
-        reference_id: str = None,
-        item_id: str = None,
-        annotation_id: str = None,
+        annotation_id: Optional[str] = None,
         metadata: Optional[Dict] = None,
     ):
-        if bool(reference_id) == bool(item_id):
-            raise Exception(
-                "You must specify either a reference_id or an item_id for an annotation."
-            )
         self.label = label
         self.vertices = vertices
-        self.reference_id = reference_id
-        self.item_id = item_id
         self.annotation_id = annotation_id
         self.metadata = metadata if metadata else {}
 
@@ -113,9 +96,7 @@ def from_json(cls, payload: dict):
         return cls(
             label=payload.get(LABEL_KEY, 0),
             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, None),
+            annotation_id=payload.get(ANNOTATION_ID_KEY, None),
             metadata=payload.get(METADATA_KEY, {}),
         )
 
@@ -124,10 +105,47 @@ def to_payload(self) -> dict:
             LABEL_KEY: self.label,
             TYPE_KEY: POLYGON_TYPE,
             GEOMETRY_KEY: {VERTICES_KEY: self.vertices},
-            REFERENCE_ID_KEY: self.reference_id,
-            ANNOTATION_ID: self.annotation_id,
+            ANNOTATION_ID_KEY: self.annotation_id,
             METADATA_KEY: self.metadata,
         }
 
     def __str__(self):
         return str(self.to_payload())
+
+
+class AnnotationItem:
+    def __init__(
+        self,
+        annotations: List[Union[PolygonAnnotation, BoxAnnotation]],
+        reference_id: Optional[str] = None,
+        item_id: Optional[str] = None,
+    ):
+        if bool(reference_id) == bool(item_id):
+            raise Exception(
+                "You must specify either a reference_id or an item_id for an annotation."
+            )
+        self.annotations = annotations
+        self.reference_id = reference_id
+        self.item_id = item_id
+
+    @classmethod
+    def from_json(cls, payload: dict):
+        return cls(
+            annotations=[
+                BoxAnnotation.from_json(ann)
+                if ann[TYPE_KEY] == "box"
+                else PolygonAnnotation.from_json(ann)
+                for ann in payload.get(ANNOTATIONS_KEY)
+            ],
+            reference_id=payload.get(REFERENCE_ID_KEY, None),
+            item_id=payload.get(DATASET_ITEM_ID_KEY, None),
+        )
+
+    def to_payload(self) -> dict:
+        return {
+            ANNOTATIONS_KEY: [annotation.to_payload() for annotation in self.annotations],
+            REFERENCE_ID_KEY: self.reference_id,
+        }
+
+    def __str__(self):
+        return str(self.to_payload())

nucleus/constants.py

@@ -14,8 +14,13 @@
 ERROR_PAYLOAD = "error_payload"
 ERROR_CODES = "error_codes"
 ANNOTATIONS_KEY = "annotations"
+ANNOTATION_ITEMS_KEY = "annotation_items"
+ANNOTATION_ID_KEY = "annotation_id"
 ANNOTATIONS_PROCESSED_KEY = "annotations_processed"
 PREDICTIONS_PROCESSED_KEY = "predictions_processed"
+EDIT_MODE_KEY = "edit_mode"
+ANNOTATION_EDIT_MODES = ["update", "insert", "rewrite"]
+DEFAULT_ANNOTATION_EDIT_MODE = "insert"
 STATUS_CODE_KEY = "status_code"
 SUCCESS_STATUS_CODES = [200, 201]
 ERRORS_KEY = "errors"

nucleus/dataset.py

@@ -1,6 +1,6 @@
-from typing import List, Dict, Any, Optional, Union
+from typing import List, Dict, Any, Optional
 from .dataset_item import DatasetItem
-from .annotation import BoxAnnotation, PolygonAnnotation
+from .annotation import AnnotationItem, BoxAnnotation, PolygonAnnotation
 from .constants import (
     DATASET_NAME_KEY,
     DATASET_MODEL_RUNS_KEY,
@@ -10,11 +10,15 @@
     REFERENCE_IDS_KEY,
     NAME_KEY,
     ITEM_KEY,
+    ANNOTATION_EDIT_MODES,
     ANNOTATIONS_KEY,
+    DEFAULT_ANNOTATION_EDIT_MODE,
 )
 from .payload_constructor import construct_model_run_creation_payload
 
 
+
+
 class Dataset:
     """
     Nucleus Dataset. You can append images with metadata to your dataset,
@@ -87,7 +91,8 @@ def create_model_run(
 
     def annotate(
         self,
-        annotations: List[Union[BoxAnnotation, PolygonAnnotation]],
+        annotations: List[AnnotationItem],
+        edit_mode: Optional[str] = None,
         batch_size: int = 20,
     ) -> dict:
         """
@@ -102,8 +107,15 @@ def annotate(
             "ignored_items": int,
         }
         """
+        if edit_mode is None:
+            edit_mode = DEFAULT_ANNOTATION_EDIT_MODE
+        if edit_mode not in ANNOTATION_EDIT_MODES:
+            raise Exception(
+                f'Edit mode must be empty or one of {ANNOTATION_EDIT_MODES}!'
+            )
+
         return self._client.annotate_dataset(
-            self.id, annotations, batch_size=batch_size
+            self.id, annotations, edit_mode=edit_mode, batch_size=batch_size
         )
 
     def ingest_tasks(self, task_ids: dict):

nucleus/payload_constructor.py

@@ -7,9 +7,11 @@
     METADATA_KEY,
     REFERENCE_ID_KEY,
     ANNOTATIONS_KEY,
+    ANNOTATION_ITEMS_KEY,
     ITEMS_KEY,
     FORCE_KEY,
     MODEL_ID_KEY,
+    EDIT_MODE_KEY,
 )
 
 
@@ -27,14 +29,15 @@ def construct_append_payload(
     )
 
 
-def construct_box_annotation_payload(
-    box_annotations: List[BoxAnnotation],
+def construct_annotation_payload(
+    box_annotations: List[Union[BoxAnnotation, PolygonAnnotation]],
+    edit_mode: str,
 ) -> dict:
     annotations = []
     for annotation in box_annotations:
         annotations.append(annotation.to_payload())
 
-    return {ANNOTATIONS_KEY: annotations}
+    return {ANNOTATION_ITEMS_KEY: annotations, EDIT_MODE_KEY: edit_mode}
 
 
 def construct_box_predictions_payload(