added CuboidPrediction dataclass

Author: Drew Kaul

nucleus/prediction.py

@@ -5,6 +5,8 @@
     PolygonAnnotation,
     Segment,
     SegmentationAnnotation,
+    CuboidAnnotation,
+    Point3D,
 )
 from .constants import (
     ANNOTATION_ID_KEY,
@@ -23,6 +25,9 @@
     ANNOTATIONS_KEY,
     ITEM_ID_KEY,
     MASK_URL_KEY,
+    POSITION_KEY,
+    DIMENSIONS_KEY,
+    YAW_KEY,
 )
 
 
@@ -146,3 +151,56 @@ def from_json(cls, payload: dict):
             metadata=payload.get(METADATA_KEY, {}),
             class_pdf=payload.get(CLASS_PDF_KEY, None),
         )
+
+
+class CuboidPrediction(CuboidAnnotation):
+    def __init__(
+        self,
+        label: str,
+        position: Point3D,
+        dimensions: Point3D,
+        yaw: float,
+        reference_id: Optional[str] = None,
+        item_id: Optional[str] = None,
+        confidence: Optional[float] = None,
+        annotation_id: Optional[str] = None,
+        metadata: Optional[Dict] = None,
+        class_pdf: Optional[Dict] = None,
+    ):
+        super().__init__(
+            label=label,
+            position=position,
+            dimensions=dimensions,
+            yaw=yaw,
+            reference_id=reference_id,
+            item_id=item_id,
+            annotation_id=annotation_id,
+            metadata=metadata,
+        )
+        self.confidence = confidence
+        self.class_pdf = class_pdf
+
+    def to_payload(self) -> dict:
+        payload = super().to_payload()
+        if self.confidence is not None:
+            payload[CONFIDENCE_KEY] = self.confidence
+        if self.class_pdf is not None:
+            payload[CLASS_PDF_KEY] = self.class_pdf
+
+        return payload
+
+    @classmethod
+    def from_json(cls, payload: dict):
+        geometry = payload.get(GEOMETRY_KEY, {})
+        return cls(
+            label=payload.get(LABEL_KEY, 0),
+            position=Point3D.from_json(geometry.get(POSITION_KEY, {})),
+            dimensions=Point3D.from_json(geometry.get(DIMENSIONS_KEY, {})),
+            yaw=payload.get(YAW_KEY, 0),
+            reference_id=payload.get(REFERENCE_ID_KEY, None),
+            item_id=payload.get(DATASET_ITEM_ID_KEY, None),
+            confidence=payload.get(CONFIDENCE_KEY, None),
+            annotation_id=payload.get(ANNOTATION_ID_KEY, None),
+            metadata=payload.get(METADATA_KEY, {}),
+            class_pdf=payload.get(CLASS_PDF_KEY, None),
+        )