diff --git a/nucleus/constants.py b/nucleus/constants.py index b594ee97..076a4a3c 100644 --- a/nucleus/constants.py +++ b/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" diff --git a/nucleus/dataset_item.py b/nucleus/dataset_item.py index 3d0901d5..83cf3b44 100644 --- a/nucleus/dataset_item.py +++ b/nucleus/dataset_item.py @@ -7,6 +7,7 @@ from .annotation import Point3D, is_local_path from .constants import ( + CAMERA_MODEL_KEY, CAMERA_PARAMS_KEY, CX_KEY, CY_KEY, @@ -14,8 +15,14 @@ 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,11 +136,18 @@ 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, @@ -115,6 +155,21 @@ def to_payload(self) -> dict: 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):