Point fix and most tests passing

Author: Diego 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 (
@@ -175,10 +175,25 @@ 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 +202,40 @@ 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"
+            ):
+                raise ValueError(
+                    "Use the Point object, not a dictionary for vertices"
+                )
 
     @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/utils.py

@@ -92,7 +92,7 @@ def convert_export_payload(api_payload):
         return_payload_row = {}
         return_payload_row[ITEM_KEY] = DatasetItem.from_json(row[ITEM_KEY])
         annotations = defaultdict(list)
-        if row[SEGMENTATION_TYPE] is not None:
+        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(

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}",
     }
@@ -139,10 +141,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(

tests/test_annotation.py

@@ -18,6 +18,7 @@
     SegmentationAnnotation,
     DatasetItem,
     Segment,
+    Point,
 )
 from nucleus.constants import ERROR_PAYLOAD
 
@@ -77,7 +78,7 @@ def test_box_gt_upload(dataset):
 
 
 def test_polygon_gt_upload(dataset):
-    annotation = PolygonAnnotation(**TEST_POLYGON_ANNOTATIONS[0])
+    annotation = PolygonAnnotation.from_json(TEST_POLYGON_ANNOTATIONS[0])
     response = dataset.annotate(annotations=[annotation])
 
     assert response["dataset_id"] == dataset.id
@@ -241,7 +242,7 @@ def test_box_gt_upload_ignore(dataset):
 
 
 def test_polygon_gt_upload_update(dataset):
-    annotation = PolygonAnnotation(**TEST_POLYGON_ANNOTATIONS[0])
+    annotation = PolygonAnnotation.from_json(TEST_POLYGON_ANNOTATIONS[0])
     response = dataset.annotate(annotations=[annotation])
 
     assert response["annotations_processed"] == 1
@@ -255,7 +256,7 @@ def test_polygon_gt_upload_update(dataset):
         "reference_id"
     ]
 
-    annotation_update = PolygonAnnotation(**annotation_update_params)
+    annotation_update = PolygonAnnotation.from_json(annotation_update_params)
     response = dataset.annotate(annotations=[annotation_update], update=True)
 
     assert response["annotations_processed"] == 1
@@ -272,7 +273,7 @@ def test_polygon_gt_upload_update(dataset):
 
 
 def test_polygon_gt_upload_ignore(dataset):
-    annotation = PolygonAnnotation(**TEST_POLYGON_ANNOTATIONS[0])
+    annotation = PolygonAnnotation.from_json(TEST_POLYGON_ANNOTATIONS[0])
     response = dataset.annotate(annotations=[annotation])
 
     assert response["annotations_processed"] == 1
@@ -286,7 +287,7 @@ def test_polygon_gt_upload_ignore(dataset):
         "reference_id"
     ]
 
-    annotation_update = PolygonAnnotation(**annotation_update_params)
+    annotation_update = PolygonAnnotation.from_json(annotation_update_params)
     # Default behavior is ignore.
     response = dataset.annotate(annotations=[annotation_update])
 

tests/test_dataset.py

@@ -21,6 +21,8 @@
     ERROR_PAYLOAD,
     IGNORED_ITEMS,
     NEW_ITEMS,
+    POLYGON_TYPE,
+    SEGMENTATION_TYPE,
     UPDATED_ITEMS,
     ITEM_KEY,
     ANNOTATIONS_KEY,
@@ -340,7 +342,13 @@ def test_annotate_async_with_error(dataset: Dataset):
 def test_append_and_export(dataset):
     # Dataset upload
     url = TEST_IMG_URLS[0]
-    annotation = BoxAnnotation(**TEST_BOX_ANNOTATIONS[0])
+    box_annotation = BoxAnnotation(**TEST_BOX_ANNOTATIONS[0])
+    segmentation_annotation = SegmentationAnnotation.from_json(
+        TEST_SEGMENTATION_ANNOTATIONS[0]
+    )
+    polygon_annotation = PolygonAnnotation.from_json(
+        TEST_POLYGON_ANNOTATIONS[0]
+    )
 
     ds_items = [
         DatasetItem(
@@ -352,12 +360,34 @@ def test_append_and_export(dataset):
     response = dataset.append(ds_items)
     assert ERROR_PAYLOAD not in response.json()
 
-    dataset.annotate(annotations=[annotation])
+    dataset.annotate(
+        annotations=[
+            box_annotation,
+            polygon_annotation,
+            segmentation_annotation,
+        ]
+    )
 
-    expected_box_annotation = copy.deepcopy(annotation)
-    expected_box_annotation.annotation_id = None
-    expected_box_annotation.metadata = {}
+    # We don't export everything on the annotations in order to speed up export.
+    def clear_fields(annotation):
+        cleared_annotation = copy.deepcopy(annotation)
+        cleared_annotation.annotation_id = None
+        cleared_annotation.metadata = {}
+        return cleared_annotation
+
+    def sort_labelmap(segmentation_annotation):
+        segmentation_annotation.annotations = sorted(
+            segmentation_annotation.annotations, key=lambda x: x.index
+        )
 
     exported = dataset.items_and_annotations()
     assert exported[0][ITEM_KEY] == ds_items[0]
-    assert exported[0][ANNOTATIONS_KEY][BOX_TYPE][0] == expected_box_annotation
+    assert exported[0][ANNOTATIONS_KEY][BOX_TYPE][0] == clear_fields(
+        box_annotation
+    )
+    assert sort_labelmap(
+        exported[0][ANNOTATIONS_KEY][SEGMENTATION_TYPE]
+    ) == sort_labelmap(clear_fields(segmentation_annotation))
+    assert exported[0][ANNOTATIONS_KEY][POLYGON_TYPE][0] == clear_fields(
+        polygon_annotation
+    )