add support for optional camera_parameters

nucleus/constants.py

@@ -23,6 +23,7 @@
 AUTOTAGS_KEY = "autotags"
 AUTOTAG_SCORE_THRESHOLD = "score_threshold"
 EXPORTED_ROWS = "exportedRows"
+CAMERA_MODEL_KEY = "camera_model"
 CAMERA_PARAMS_KEY = "camera_params"
 CLASS_PDF_KEY = "class_pdf"
 CONFIDENCE_KEY = "confidence"
@@ -62,6 +63,10 @@
 ITEM_KEY = "item"
 ITEM_METADATA_SCHEMA_KEY = "item_metadata_schema"
 JOB_ID_KEY = "job_id"
+K1_KEY = "k1"
+K2_KEY = "k2"
+K3_KEY = "k3"
+K4_KEY = "k4"
 KEEP_HISTORY_KEY = "keep_history"
 LENGTH_KEY = "length"
 JOB_STATUS_KEY = "job_status"
@@ -80,6 +85,8 @@
 NUCLEUS_ENDPOINT = "https://api.scale.com/v1/nucleus"
 NUM_SENSORS_KEY = "num_sensors"
 ORIGINAL_IMAGE_URL_KEY = "original_image_url"
+P1_KEY = "p1"
+P2_KEY = "p2"
 POINTCLOUD_KEY = "pointcloud"
 POINTCLOUD_LOCATION_KEY = "pointcloud_location"
 POINTCLOUD_URL_KEY = "pointcloud_url"

nucleus/dataset_item.py

@@ -7,15 +7,22 @@
 
 from .annotation import Point3D, is_local_path
 from .constants import (
+    CAMERA_MODEL_KEY,
     CAMERA_PARAMS_KEY,
     CX_KEY,
     CY_KEY,
     FX_KEY,
     FY_KEY,
     HEADING_KEY,
     IMAGE_URL_KEY,
+    K1_KEY,
+    K2_KEY,
+    K3_KEY,
+    K4_KEY,
     METADATA_KEY,
     ORIGINAL_IMAGE_URL_KEY,
+    P1_KEY,
+    P2_KEY,
     POINTCLOUD_URL_KEY,
     POSITION_KEY,
     REFERENCE_ID_KEY,
@@ -29,6 +36,18 @@
 )
 
 
+class CameraModels(Enum):
+    BROWN_CONRADY = "brown_conrady"
+    FISHEYE = "fisheye"
+
+    def __contains__(self, item):
+        try:
+            self(item)
+        except ValueError:
+            return False
+        return True
+
+
 @dataclass
 class Quaternion:
     """Quaternion objects are used to represent rotation.
@@ -92,6 +111,20 @@ class CameraParams:
     fy: float
     cx: float
     cy: float
+    camera_model: str
+    k1: float
+    k2: float
+    k3: float
+    k4: float
+    p1: float
+    p2: float
+
+    def __post_init__(self):
+        if self.camera_model is not None:
+            if self.camera_model not in CameraModels:
+                raise ValueError(
+                    f'Invalid Camera Model, the supported options are "{CameraModels.BROWN_CONRADY}" and "{CameraModels.FISHEYE}"'
+                )
 
     @classmethod
     def from_json(cls, payload: Dict[str, Any]):
@@ -103,18 +136,40 @@ def from_json(cls, payload: Dict[str, Any]):
             payload[FY_KEY],
             payload[CX_KEY],
             payload[CY_KEY],
+            payload.get(K1_KEY, None),
+            payload.get(K2_KEY, None),
+            payload.get(K3_KEY, None),
+            payload.get(K4_KEY, None),
+            payload.get(P1_KEY, None),
+            payload.get(P2_KEY, None),
+            payload.get(CAMERA_MODEL_KEY, None),
         )
 
     def to_payload(self) -> dict:
         """Serializes camera params object to schematized JSON dict."""
-        return {
+        payload = {
             POSITION_KEY: self.position.to_payload(),
             HEADING_KEY: self.heading.to_payload(),
             FX_KEY: self.fx,
             FY_KEY: self.fy,
             CX_KEY: self.cx,
             CY_KEY: self.cy,
         }
+        if self.k1:
+            payload[K1_KEY] = self.k1
+        if self.k2:
+            payload[K2_KEY] = self.k2
+        if self.k3:
+            payload[K3_KEY] = self.k3
+        if self.k4:
+            payload[K4_KEY] = self.k4
+        if self.p1:
+            payload[P1_KEY] = self.p1
+        if self.p2:
+            payload[P2_KEY] = self.p2
+        if self.camera_model:
+            payload[CAMERA_MODEL_KEY] = self.camera_model
+        return payload
 
 
 class DatasetItemType(Enum):