All predictions and annotations have better repr

Author: ardila

nucleus/annotation.py

@@ -21,8 +21,16 @@
     ANNOTATIONS_KEY,
 )
 
+from dataclasses import dataclass
 
 class Annotation:
+
+    def _check_ids(self):
+        if not bool(self.reference_id) and not bool(self.item_id):
+            raise Exception(
+                "You must specify either a reference_id or an item_id for an annotation."
+            )
+
     @classmethod
     def from_json(cls, payload: dict):
         if payload.get(TYPE_KEY, None) == BOX_TYPE:
@@ -32,17 +40,11 @@ def from_json(cls, payload: dict):
         else:
             return SegmentationAnnotation.from_json(payload)
 
-
+@dataclass
 class Segment:
-    def __init__(
-        self, label: str, index: int, metadata: Optional[dict] = None
-    ):
-        self.label = label
-        self.index = index
-        self.metadata = metadata
-
-    def __str__(self):
-        return str(self.to_payload())
+    label: str
+    index: int
+    metadata: Optional[dict] = None
 
     @classmethod
     def from_json(cls, payload: dict):
@@ -62,30 +64,18 @@ def to_payload(self) -> dict:
         return payload
 
 
+@dataclass
 class SegmentationAnnotation(Annotation):
-    def __init__(
-        self,
-        mask_url: str,
-        annotations: List[Segment],
-        reference_id: Optional[str] = None,
-        item_id: Optional[str] = None,
-        annotation_id: Optional[str] = None,
-    ):
-        super().__init__()
-        if not mask_url:
+    mask_url: str
+    annotations: List[Segment]
+    reference_id: Optional[str] = None
+    item_id: Optional[str] = None
+    annotation_id: Optional[str] = None
+    def __post_init__(self):
+        if not self.mask_url:
             raise Exception("You must specify a mask_url.")
-        if bool(reference_id) == bool(item_id):
-            raise Exception(
-                "You must specify either a reference_id or an item_id for an annotation."
-            )
-        self.mask_url = mask_url
-        self.annotations = annotations
-        self.reference_id = reference_id
-        self.item_id = item_id
-        self.annotation_id = annotation_id
-
-    def __str__(self):
-        return str(self.to_payload())
+        self._check_ids()
+        
 
     @classmethod
     def from_json(cls, payload: dict):
@@ -120,35 +110,20 @@ class AnnotationTypes(Enum):
     POLYGON = POLYGON_TYPE
 
 
-# TODO: Add base annotation class to reduce repeated code here
+@dataclass
 class BoxAnnotation(Annotation):
-    # pylint: disable=too-many-instance-attributes
-    def __init__(
-        self,
-        label: str,
-        x: Union[float, int],
-        y: Union[float, int],
-        width: Union[float, int],
-        height: Union[float, int],
-        reference_id: Optional[str] = None,
-        item_id: Optional[str] = None,
-        annotation_id: Optional[str] = None,
-        metadata: Optional[Dict] = None,
-    ):
-        super().__init__()
-        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 {}
+    label: str
+    x: Union[float, int]
+    y: Union[float, int]
+    width: Union[float, int]
+    height: Union[float, int]
+    reference_id: Optional[str] = None
+    item_id: Optional[str] = None
+    annotation_id: Optional[str] = None
+    metadata: Optional[Dict] = None
+    def __post_init__(self):
+        self._check_ids()
+        self.metadata = self.metadata if self.metadata else {}
 
     @classmethod
     def from_json(cls, payload: dict):
@@ -180,32 +155,19 @@ def to_payload(self) -> dict:
             METADATA_KEY: self.metadata,
         }
 
-    def __str__(self):
-        return str(self.to_payload())
-
 
 # TODO: Add Generic type for 2D point
+@dataclass
 class PolygonAnnotation(Annotation):
-    def __init__(
-        self,
-        label: str,
-        vertices: List[Any],
-        reference_id: Optional[str] = None,
-        item_id: Optional[str] = None,
-        annotation_id: Optional[str] = None,
-        metadata: Optional[Dict] = None,
-    ):
-        super().__init__()
-        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 {}
+    label: str
+    vertices: List[Any]
+    reference_id: Optional[str] = None
+    item_id: Optional[str] = None
+    annotation_id: Optional[str] = None
+    metadata: Optional[Dict] = None
+    def __post_init__(self):
+        self._check_ids()
+        self.metadata = self.metadata if self.metadata else {}
 
     @classmethod
     def from_json(cls, payload: dict):
@@ -228,6 +190,3 @@ def to_payload(self) -> dict:
             ANNOTATION_ID_KEY: self.annotation_id,
             METADATA_KEY: self.metadata,
         }
-
-    def __str__(self):
-        return str(self.to_payload())

nucleus/dataset_item.py

@@ -7,21 +7,19 @@
     DATASET_ITEM_ID_KEY,
 )
 
+from dataclasses import dataclass
 
+
+@dataclass
 class DatasetItem:
-    def __init__(
-        self,
-        image_location: str,
-        reference_id: str = None,
-        item_id: str = None,
-        metadata: dict = None,
-    ):
-
-        self.image_url = image_location
-        self.local = self._is_local_path(image_location)
-        self.item_id = item_id
-        self.reference_id = reference_id
-        self.metadata = metadata
+
+    image_location: str
+    reference_id: str = None
+    item_id: str = None
+    metadata: dict = None
+
+    def __post_init__(self):
+        self.local = self._is_local_path(self.image_location)
 
     @classmethod
     def from_json(cls, payload: dict):
@@ -35,9 +33,6 @@ def from_json(cls, payload: dict):
             metadata=payload.get(METADATA_KEY, {}),
         )
 
-    def __str__(self):
-        return str(self.to_payload())
-
     def _is_local_path(self, path: str) -> bool:
         path_components = [comp.lower() for comp in path.split("/")]
         return not (

nucleus/model_run.py

@@ -130,7 +130,6 @@ def _format_prediction_response(
         List[Union[BoxPrediction, PolygonPrediction, SegmentationPrediction]],
     ]:
         annotation_payload = response.get(ANNOTATIONS_KEY, None)
-        SegmentationPrediction.from_json({"asdf": "asdf"})
         if not annotation_payload:
             # An error occurred
             return response

nucleus/prediction.py

@@ -93,9 +93,6 @@ def from_json(cls, payload: dict) -> BoxPrediction:
             metadata=payload.get(METADATA_KEY, {}),
         )
 
-    def __str__(self):
-        return str(self.to_payload())
-
 
 class PolygonPrediction(PolygonAnnotation):
     def __init__(
@@ -132,6 +129,3 @@ def from_json(cls, payload: dict) -> PolygonPrediction:
             annotation_id=payload.get(ANNOTATION_ID_KEY, None),
             metadata=payload.get(METADATA_KEY, {}),
         )
-
-    def __str__(self):
-        return str(self.to_payload())

pyproject.toml

@@ -41,6 +41,7 @@ tqdm = "^4.60.0"
 poetry = "^1.1.5"
 pytest = "^6.2.3"
 pylint = "^2.7.4"
+boto3 = "^1.17.51"
 
 [build-system]
 requires = ["poetry-core>=1.0.0"]

tests/test_annotation.py

@@ -17,10 +17,29 @@
     PolygonAnnotation,
     SegmentationAnnotation,
     DatasetItem,
+    Segment,
 )
 from nucleus.constants import ERROR_PAYLOAD
 
 
+def test_repr(test_object: any):
+    assert eval(str(test_object)) == test_object
+
+
+def test_reprs():
+    [
+        test_repr(SegmentationAnnotation.from_json(_))
+        for _ in TEST_SEGMENTATION_ANNOTATIONS
+    ]
+
+    [test_repr(BoxAnnotation.from_json(_)) for _ in TEST_BOX_ANNOTATIONS]
+
+    [
+        test_repr(PolygonAnnotation.from_json(_))
+        for _ in TEST_POLYGON_ANNOTATIONS
+    ]
+
+
 @pytest.fixture()
 def dataset(CLIENT):
     ds = CLIENT.create_dataset(TEST_DATASET_NAME)

tests/test_prediction.py

@@ -20,10 +20,29 @@
     PolygonPrediction,
     SegmentationPrediction,
     DatasetItem,
+    Segment
 )
 from nucleus.constants import ERROR_PAYLOAD
 
 
+def test_repr(test_object: any):
+    assert eval(str(test_object)) == test_object
+
+
+def test_reprs():
+    [
+        test_repr(SegmentationPrediction.from_json(_))
+        for _ in TEST_SEGMENTATION_PREDICTIONS
+    ]
+
+    [test_repr(BoxPrediction.from_json(_)) for _ in TEST_BOX_PREDICTIONS]
+
+    [
+        test_repr(PolygonPrediction.from_json(_))
+        for _ in TEST_POLYGON_PREDICTIONS
+    ]
+
+
 @pytest.fixture()
 def model_run(CLIENT):
     ds = CLIENT.create_dataset(TEST_DATASET_NAME)