Merge pull request #77 from scaleapi/cuboid_upload

added CuboidAnnotation dataclass

Author: sasha-scale

nucleus/annotation.py

@@ -8,7 +8,9 @@
     ANNOTATION_ID_KEY,
     ANNOTATIONS_KEY,
     BOX_TYPE,
+    CUBOID_TYPE,
     DATASET_ITEM_ID_KEY,
+    DIMENSIONS_KEY,
     GEOMETRY_KEY,
     HEIGHT_KEY,
     INDEX_KEY,
@@ -18,12 +20,15 @@
     MASK_URL_KEY,
     METADATA_KEY,
     POLYGON_TYPE,
+    POSITION_KEY,
     REFERENCE_ID_KEY,
     TYPE_KEY,
     VERTICES_KEY,
     WIDTH_KEY,
     X_KEY,
+    YAW_KEY,
     Y_KEY,
+    Z_KEY,
 )
 
 
@@ -43,6 +48,8 @@ def from_json(cls, payload: dict):
             return BoxAnnotation.from_json(payload)
         elif payload.get(TYPE_KEY, None) == POLYGON_TYPE:
             return PolygonAnnotation.from_json(payload)
+        elif payload.get(TYPE_KEY, None) == CUBOID_TYPE:
+            return CuboidAnnotation.from_json(payload)
         else:
             return SegmentationAnnotation.from_json(payload)
 
@@ -125,6 +132,7 @@ def to_payload(self) -> dict:
 class AnnotationTypes(Enum):
     BOX = BOX_TYPE
     POLYGON = POLYGON_TYPE
+    CUBOID = CUBOID_TYPE
 
 
 @dataclass  # pylint: disable=R0902
@@ -187,6 +195,20 @@ def to_payload(self) -> dict:
         return {X_KEY: self.x, Y_KEY: self.y}
 
 
+@dataclass
+class Point3D:
+    x: float
+    y: float
+    z: float
+
+    @classmethod
+    def from_json(cls, payload: Dict[str, float]):
+        return cls(payload[X_KEY], payload[Y_KEY], payload[Z_KEY])
+
+    def to_payload(self) -> dict:
+        return {X_KEY: self.x, Y_KEY: self.y, Z_KEY: self.z}
+
+
 @dataclass
 class PolygonAnnotation(Annotation):
     label: str
@@ -241,6 +263,50 @@ def to_payload(self) -> dict:
         return payload
 
 
+@dataclass  # pylint: disable=R0902
+class CuboidAnnotation(Annotation):  # pylint: disable=R0902
+    label: str
+    position: Point3D
+    dimensions: Point3D
+    yaw: float
+    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):
+        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),
+            annotation_id=payload.get(ANNOTATION_ID_KEY, None),
+            metadata=payload.get(METADATA_KEY, {}),
+        )
+
+    def to_payload(self) -> dict:
+        return {
+            LABEL_KEY: self.label,
+            TYPE_KEY: CUBOID_TYPE,
+            GEOMETRY_KEY: {
+                POSITION_KEY: self.position,
+                DIMENSIONS_KEY: self.dimensions,
+                YAW_KEY: self.yaw,
+            },
+            REFERENCE_ID_KEY: self.reference_id,
+            ANNOTATION_ID_KEY: self.annotation_id,
+            METADATA_KEY: self.metadata,
+        }
+
+
 def check_all_annotation_paths_remote(
     annotations: Sequence[Union[Annotation]],
 ):

nucleus/constants.py

@@ -7,7 +7,8 @@
 POLYGON_TYPE = "polygon"
 MASK_TYPE = "mask"
 SEGMENTATION_TYPE = "segmentation"
-ANNOTATION_TYPES = (BOX_TYPE, POLYGON_TYPE, SEGMENTATION_TYPE)
+CUBOID_TYPE = "cuboid"
+ANNOTATION_TYPES = (BOX_TYPE, POLYGON_TYPE, SEGMENTATION_TYPE, CUBOID_TYPE)
 ANNOTATION_UPDATE_KEY = "update"
 AUTOTAGS_KEY = "autotags"
 EXPORTED_ROWS = "exportedRows"
@@ -22,6 +23,7 @@
 DATASET_SLICES_KEY = "slice_ids"
 DEFAULT_ANNOTATION_UPDATE_MODE = False
 DEFAULT_NETWORK_TIMEOUT_SEC = 120
+DIMENSIONS_KEY = "dimensions"
 EMBEDDINGS_URL_KEY = "embeddings_url"
 ERRORS_KEY = "errors"
 ERROR_CODES = "error_codes"
@@ -48,6 +50,7 @@
 NEW_ITEMS = "new_items"
 NUCLEUS_ENDPOINT = "https://api.scale.com/v1/nucleus"
 ORIGINAL_IMAGE_URL_KEY = "original_image_url"
+POSITION_KEY = "position"
 PREDICTIONS_IGNORED_KEY = "predictions_ignored"
 PREDICTIONS_PROCESSED_KEY = "predictions_processed"
 REFERENCE_IDS_KEY = "reference_ids"
@@ -63,5 +66,7 @@
 UPDATE_KEY = "update"
 VERTICES_KEY = "vertices"
 WIDTH_KEY = "width"
+YAW_KEY = "yaw"
 X_KEY = "x"
 Y_KEY = "y"
+Z_KEY = "z"

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),
+        )