Passes local tests

Author: Diego Ardila

nucleus/annotation.py

@@ -174,9 +174,6 @@ def to_payload(self) -> dict:
         }
 
 
-# TODO: Add Generic type for 2D point
-
-
 @dataclass
 class Point:
     x: float
@@ -206,9 +203,15 @@ def __post_init__(self):
             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"
-                )
+                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):

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,
@@ -101,7 +102,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,
@@ -133,7 +134,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),

pyproject.toml

@@ -21,7 +21,7 @@ exclude = '''
 
 [tool.poetry]
 name = "scale-nucleus"
-version = "0.1.8"
+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>"]
@@ -36,7 +36,7 @@ python = "^3.6.2"
 requests = "^2.23.0"
 tqdm = "^4.41.0"
 dataclasses = { version = "^0.7", python = "^3.6.1, <3.7" }
-aiohttp = "^3.7.4"
+aiohttp = "3.7.4.post0"
 
 [tool.poetry.dev-dependencies]
 poetry = "^1.1.5"

tests/helpers.py

@@ -112,7 +112,7 @@ def reference_id_from_url(url):
     {
         **TEST_POLYGON_ANNOTATIONS[i],
         "confidence": 0.10 * i,
-        "class_pdf": TEST_POLYGON_MODEL_PDF,
+        "class_pdf": None,
     }
     for i in range(len(TEST_POLYGON_ANNOTATIONS))
 ]

tests/test_prediction.py

@@ -22,6 +22,7 @@
     DatasetItem,
     Segment,
     ModelRun,
+    Point,
 )
 from nucleus.constants import ERROR_PAYLOAD
 
@@ -89,7 +90,7 @@ def test_box_pred_upload(model_run):
 
 
 def test_polygon_pred_upload(model_run):
-    prediction = PolygonPrediction(**TEST_POLYGON_PREDICTIONS[0])
+    prediction = PolygonPrediction.from_json(TEST_POLYGON_PREDICTIONS[0])
     response = model_run.predict(annotations=[prediction])
 
     assert response["model_run_id"] == model_run.model_run_id
@@ -189,7 +190,7 @@ def test_box_pred_upload_ignore(model_run):
 
 
 def test_polygon_pred_upload_update(model_run):
-    prediction = PolygonPrediction(**TEST_POLYGON_PREDICTIONS[0])
+    prediction = PolygonPrediction.from_json(TEST_POLYGON_PREDICTIONS[0])
     response = model_run.predict(annotations=[prediction])
 
     assert response["predictions_processed"] == 1
@@ -203,7 +204,7 @@ def test_polygon_pred_upload_update(model_run):
         "reference_id"
     ]
 
-    prediction_update = PolygonPrediction(**prediction_update_params)
+    prediction_update = PolygonPrediction.from_json(prediction_update_params)
     response = model_run.predict(annotations=[prediction_update], update=True)
 
     assert response["predictions_processed"] == 1
@@ -217,7 +218,7 @@ def test_polygon_pred_upload_update(model_run):
 
 
 def test_polygon_pred_upload_ignore(model_run):
-    prediction = PolygonPrediction(**TEST_POLYGON_PREDICTIONS[0])
+    prediction = PolygonPrediction.from_json(TEST_POLYGON_PREDICTIONS[0])
     response = model_run.predict(annotations=[prediction])
 
     assert response["predictions_processed"] == 1
@@ -231,7 +232,7 @@ def test_polygon_pred_upload_ignore(model_run):
         "reference_id"
     ]
 
-    prediction_update = PolygonPrediction(**prediction_update_params)
+    prediction_update = PolygonPrediction.from_json(prediction_update_params)
     # Default behavior is ignore.
     response = model_run.predict(annotations=[prediction_update])
 
@@ -249,7 +250,9 @@ def test_mixed_pred_upload(model_run):
     prediction_semseg = SegmentationPrediction.from_json(
         TEST_SEGMENTATION_PREDICTIONS[0]
     )
-    prediction_polygon = PolygonPrediction(**TEST_POLYGON_PREDICTIONS[0])
+    prediction_polygon = PolygonPrediction.from_json(
+        TEST_POLYGON_PREDICTIONS[0]
+    )
     prediction_bbox = BoxPrediction(**TEST_BOX_PREDICTIONS[0])
     response = model_run.predict(
         annotations=[prediction_semseg, prediction_polygon, prediction_bbox]
@@ -276,7 +279,9 @@ def test_mixed_pred_upload_async(model_run: ModelRun):
     prediction_semseg = SegmentationPrediction.from_json(
         TEST_SEGMENTATION_PREDICTIONS[0]
     )
-    prediction_polygon = PolygonPrediction(**TEST_POLYGON_PREDICTIONS[0])
+    prediction_polygon = PolygonPrediction.from_json(
+        TEST_POLYGON_PREDICTIONS[0]
+    )
     prediction_bbox = BoxPrediction(**TEST_BOX_PREDICTIONS[0])
     job: AsyncJob = model_run.predict(
         annotations=[prediction_semseg, prediction_polygon, prediction_bbox],
@@ -310,7 +315,9 @@ def test_mixed_pred_upload_async_with_error(model_run: ModelRun):
     prediction_semseg = SegmentationPrediction.from_json(
         TEST_SEGMENTATION_PREDICTIONS[0]
     )
-    prediction_polygon = PolygonPrediction(**TEST_POLYGON_PREDICTIONS[0])
+    prediction_polygon = PolygonPrediction.from_json(
+        TEST_POLYGON_PREDICTIONS[0]
+    )
     prediction_bbox = BoxPrediction(**TEST_BOX_PREDICTIONS[0])
     prediction_bbox.reference_id = "fake_garbage"