Merge pull request #76 from scaleapi/da/export

Batch export functionality for slice and dataset + point fix

Author: ardila

nucleus/__init__.py

@@ -67,6 +67,7 @@
     PolygonAnnotation,
     Segment,
     SegmentationAnnotation,
+    Point,
 )
 from .constants import (
     ANNOTATION_METADATA_SCHEMA_KEY,

nucleus/annotation.py

@@ -1,7 +1,7 @@
 import json
 from dataclasses import dataclass
 from enum import Enum
-from typing import Any, Dict, List, Optional, Sequence, Union
+from typing import Dict, List, Optional, Sequence, Union
 from nucleus.dataset_item import is_local_path
 
 from .constants import (
@@ -174,11 +174,23 @@ def to_payload(self) -> dict:
         }
 
 
-# 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
@@ -187,28 +199,46 @@ 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),
             metadata=payload.get(METADATA_KEY, {}),
         )
 
     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(

nucleus/constants.py

@@ -10,7 +10,7 @@
 ANNOTATION_TYPES = (BOX_TYPE, POLYGON_TYPE, SEGMENTATION_TYPE)
 ANNOTATION_UPDATE_KEY = "update"
 AUTOTAGS_KEY = "autotags"
-
+EXPORTED_ROWS = "exportedRows"
 CLASS_PDF_KEY = "class_pdf"
 CONFIDENCE_KEY = "confidence"
 DATASET_ID_KEY = "dataset_id"

nucleus/dataset.py

@@ -4,6 +4,7 @@
 
 from nucleus.job import AsyncJob
 from nucleus.utils import (
+    convert_export_payload,
     format_dataset_item_response,
     serialize_and_write_to_presigned_url,
 )
@@ -16,6 +17,7 @@
     DATASET_NAME_KEY,
     DATASET_SLICES_KEY,
     DEFAULT_ANNOTATION_UPDATE_MODE,
+    EXPORTED_ROWS,
     JOB_ID_KEY,
     NAME_KEY,
     REFERENCE_IDS_KEY,
@@ -327,3 +329,23 @@ def delete_custom_index(self):
 
     def check_index_status(self, job_id: str):
         return self._client.check_index_status(job_id)
+
+    def items_and_annotations(
+        self,
+    ) -> List[Dict[str, Union[DatasetItem, Dict[str, List[Annotation]]]]]:
+        """Returns a list of all DatasetItems and Annotations in this slice.
+
+        Returns:
+            A list, where each item is a dict with two keys representing a row
+            in the dataset.
+            * One value in the dict is the DatasetItem, containing a reference to the
+                item that was annotated.
+            * The other value is a dictionary containing all the annotations for this
+                dataset item, sorted by annotation type.
+        """
+        api_payload = self._client.make_request(
+            payload=None,
+            route=f"dataset/{self.id}/exportForTraining",
+            requests_command=requests.get,
+        )
+        return convert_export_payload(api_payload[EXPORTED_ROWS])

nucleus/prediction.py

@@ -1,6 +1,7 @@
-from typing import Dict, Optional, List, Any
+from typing import Dict, Optional, List
 from .annotation import (
     BoxAnnotation,
+    Point,
     PolygonAnnotation,
     Segment,
     SegmentationAnnotation,
@@ -102,7 +103,7 @@ class PolygonPrediction(PolygonAnnotation):
     def __init__(
         self,
         label: str,
-        vertices: List[Any],
+        vertices: List[Point],
         reference_id: Optional[str] = None,
         item_id: Optional[str] = None,
         confidence: Optional[float] = None,
@@ -135,7 +136,9 @@ 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),
             confidence=payload.get(CONFIDENCE_KEY, None),

nucleus/slice.py

@@ -1,10 +1,12 @@
-from typing import Dict, List, Iterable, Set, Tuple, Optional, Union
-from nucleus.dataset_item import DatasetItem
+from typing import Dict, Iterable, List, Set, Tuple, Union
+
+import requests
+
 from nucleus.annotation import Annotation
-from nucleus.utils import format_dataset_item_response
+from nucleus.dataset_item import DatasetItem
 from nucleus.job import AsyncJob
-
-from .constants import DEFAULT_ANNOTATION_UPDATE_MODE
+from nucleus.utils import convert_export_payload, format_dataset_item_response
+from nucleus.constants import EXPORTED_ROWS
 
 
 class Slice:
@@ -109,42 +111,12 @@ def items_and_annotations(
             * The other value is a dictionary containing all the annotations for this
                 dataset item, sorted by annotation type.
         """
-        return list(self.items_and_annotation_generator())
-
-    def annotate(
-        self,
-        annotations: List[Annotation],
-        update: Optional[bool] = DEFAULT_ANNOTATION_UPDATE_MODE,
-        batch_size: int = 5000,
-        strict=True,
-    ):
-        """Update annotations within this slice.
-
-        Args:
-            annotations: List of annotations to upload
-            batch_size: How many annotations to send per request.
-            strict: Whether to first check that the annotations belong to this slice.
-                Set to false to avoid this check and speed up upload.
-        """
-        if strict:
-            (
-                annotations_are_in_slice,
-                item_ids_not_found_in_slice,
-                reference_ids_not_found_in_slice,
-            ) = check_annotations_are_in_slice(annotations, self)
-            if not annotations_are_in_slice:
-                message = "Not all annotations are in this slice.\n"
-                if item_ids_not_found_in_slice:
-                    message += f"Item ids not found in slice: {item_ids_not_found_in_slice} \n"
-                if reference_ids_not_found_in_slice:
-                    message += f"Reference ids not found in slice: {reference_ids_not_found_in_slice}"
-                raise ValueError(message)
-        self._client.annotate_dataset(
-            dataset_id=self.dataset_id,
-            annotations=annotations,
-            update=update,
-            batch_size=batch_size,
+        api_payload = self._client.make_request(
+            payload=None,
+            route=f"slice/{self.slice_id}/exportForTraining",
+            requests_command=requests.get,
         )
+        return convert_export_payload(api_payload[EXPORTED_ROWS])
 
     def send_to_labeling(self, project_id: str):
         response = self._client.make_request(

nucleus/utils.py

@@ -1,16 +1,29 @@
 """Shared stateless utility function library"""
 
-
+from collections import defaultdict
 import io
 import uuid
 from typing import IO, Dict, List, Sequence, Union
 
 import requests
 from requests.models import HTTPError
 
-from nucleus.annotation import Annotation
-
-from .constants import ANNOTATION_TYPES, ANNOTATIONS_KEY, ITEM_KEY
+from nucleus.annotation import (
+    Annotation,
+    BoxAnnotation,
+    PolygonAnnotation,
+    SegmentationAnnotation,
+)
+
+from .constants import (
+    ANNOTATION_TYPES,
+    ANNOTATIONS_KEY,
+    BOX_TYPE,
+    ITEM_KEY,
+    POLYGON_TYPE,
+    REFERENCE_ID_KEY,
+    SEGMENTATION_TYPE,
+)
 from .dataset_item import DatasetItem
 from .prediction import BoxPrediction, PolygonPrediction
 
@@ -73,6 +86,31 @@ def format_dataset_item_response(response: dict) -> dict:
     }
 
 
+def convert_export_payload(api_payload):
+    return_payload = []
+    for row in api_payload:
+        return_payload_row = {}
+        return_payload_row[ITEM_KEY] = DatasetItem.from_json(row[ITEM_KEY])
+        annotations = defaultdict(list)
+        if row.get(SEGMENTATION_TYPE) is not None:
+            segmentation = row[SEGMENTATION_TYPE]
+            segmentation[REFERENCE_ID_KEY] = row[ITEM_KEY][REFERENCE_ID_KEY]
+            annotations[SEGMENTATION_TYPE] = SegmentationAnnotation.from_json(
+                segmentation
+            )
+        for polygon in row[POLYGON_TYPE]:
+            polygon[REFERENCE_ID_KEY] = row[ITEM_KEY][REFERENCE_ID_KEY]
+            annotations[POLYGON_TYPE].append(
+                PolygonAnnotation.from_json(polygon)
+            )
+        for box in row[BOX_TYPE]:
+            box[REFERENCE_ID_KEY] = row[ITEM_KEY][REFERENCE_ID_KEY]
+            annotations[BOX_TYPE].append(BoxAnnotation.from_json(box))
+        return_payload_row[ANNOTATIONS_KEY] = annotations
+        return_payload.append(return_payload_row)
+    return return_payload
+
+
 def serialize_and_write(
     upload_units: Sequence[Union[DatasetItem, Annotation]], file_pointer
 ):

pyproject.toml

@@ -21,7 +21,7 @@ exclude = '''
 
 [tool.poetry]
 name = "scale-nucleus"
-version = "0.1.9"
+version = "0.1.10"
 description = "The official Python client library for Nucleus, the Data Platform for AI"
 license =  "MIT"
 authors = ["Scale AI Nucleus Team <nucleusapi@scaleapi.com>"]

tests/helpers.py

@@ -57,13 +57,15 @@ def reference_id_from_url(url):
 TEST_POLYGON_ANNOTATIONS = [
     {
         "label": f"[Pytest] Polygon Annotation ${i}",
-        "vertices": [
-            {
-                "x": 50 + i * 10 + j,
-                "y": 60 + i * 10 + j,
-            }
-            for j in range(3)
-        ],
+        "geometry": {
+            "vertices": [
+                {
+                    "x": 50 + i * 10 + j,
+                    "y": 60 + i * 10 + j,
+                }
+                for j in range(3)
+            ],
+        },
         "reference_id": reference_id_from_url(TEST_IMG_URLS[i]),
         "annotation_id": f"[Pytest] Polygon Annotation Annotation Id{i}",
     }
@@ -149,10 +151,10 @@ def assert_polygon_annotation_matches_dict(
         annotation_instance.annotation_id == annotation_dict["annotation_id"]
     )
     for instance_pt, dict_pt in zip(
-        annotation_instance.vertices, annotation_dict["vertices"]
+        annotation_instance.vertices, annotation_dict["geometry"]["vertices"]
     ):
-        assert instance_pt["x"] == dict_pt["x"]
-        assert instance_pt["y"] == dict_pt["y"]
+        assert instance_pt.x == dict_pt["x"]
+        assert instance_pt.y == dict_pt["y"]
 
 
 def assert_segmentation_annotation_matches_dict(