Allow camera_params update through update_items_metadata (#271)

jean-lucas · web-flow · commit 7363a0f9a8fc · 2022-04-05T17:26:00.000-04:00
diff --git a/nucleus/__init__.py b/nucleus/__init__.py
@@ -59,6 +59,7 @@
     Segment,
     SegmentationAnnotation,
 )
+from .camera_params import CameraParams
 from .connection import Connection
 from .constants import (
     ANNOTATION_METADATA_SCHEMA_KEY,
@@ -97,7 +98,7 @@
 from .data_transfer_object.dataset_details import DatasetDetails
 from .data_transfer_object.dataset_info import DatasetInfo
 from .dataset import Dataset
-from .dataset_item import CameraParams, DatasetItem, Quaternion
+from .dataset_item import DatasetItem
 from .deprecation_warning import deprecated
 from .errors import (
     DatasetItemRetrievalError,
@@ -126,6 +127,7 @@
     PolygonPrediction,
     SegmentationPrediction,
 )
+from .quaternion import Quaternion
 from .retry_strategy import RetryStrategy
 from .scene import Frame, LidarScene, VideoScene
 from .slice import Slice
diff --git a/nucleus/camera_params.py b/nucleus/camera_params.py
@@ -0,0 +1,131 @@
+from dataclasses import dataclass
+from enum import Enum
+from typing import Any, Dict
+
+from .annotation import Point3D
+from .constants import (
+    CAMERA_MODEL_KEY,
+    CX_KEY,
+    CY_KEY,
+    FX_KEY,
+    FY_KEY,
+    HEADING_KEY,
+    K1_KEY,
+    K2_KEY,
+    K3_KEY,
+    K4_KEY,
+    P1_KEY,
+    P2_KEY,
+    POSITION_KEY,
+)
+from .quaternion import Quaternion
+
+REQUIRED_CAMERA_PARAMS_KEYS = {
+    POSITION_KEY,
+    HEADING_KEY,
+    FX_KEY,
+    FY_KEY,
+    CX_KEY,
+    CY_KEY,
+}
+
+
+class CameraModels(str, Enum):
+    BROWN_CONRADY = "brown_conrady"
+    FISHEYE = "fisheye"
+
+    def __contains__(self, item):
+        try:
+            self(item)
+        except ValueError:
+            return False
+        return True
+
+
+@dataclass
+class CameraParams:
+    """Camera position/heading used to record the image.
+
+    Args:
+        position (:class:`Point3D`): World-normalized position of the camera
+        heading (:class:`Quaternion`): Vector4 indicating the quaternion of the
+          camera direction; note that the z-axis of the camera frame
+          represents the camera's optical axis.  See `Heading Examples
+          <https://docs.scale.com/reference/data-types-and-the-frame-objects#heading-examples>`_.
+        fx (float): Focal length in x direction (in pixels).
+        fy (float): Focal length in y direction (in pixels).
+        cx (float): Principal point x value.
+        cy (float): Principal point y value.
+    """
+
+    position: Point3D
+    heading: Quaternion
+    fx: float
+    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 (k for k 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]):
+        """Instantiates camera params object from schematized JSON dict payload."""
+        keys = set(payload.keys())
+        if not keys.issuperset(REQUIRED_CAMERA_PARAMS_KEYS):
+            raise ValueError(
+                f"The following fields must be present in the camera_params dictionary: {REQUIRED_CAMERA_PARAMS_KEYS}"
+            )
+
+        return cls(
+            Point3D.from_json(payload[POSITION_KEY]),
+            Quaternion.from_json(payload[HEADING_KEY]),
+            payload[FX_KEY],
+            payload[FY_KEY],
+            payload[CX_KEY],
+            payload[CY_KEY],
+            payload.get(CAMERA_MODEL_KEY, None),
+            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),
+        )
+
+    def to_payload(self) -> dict:
+        """Serializes camera params object to schematized JSON dict."""
+        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
diff --git a/nucleus/dataset.py b/nucleus/dataset.py
@@ -1561,11 +1561,14 @@ def update_item_metadata(self, mapping: Dict[str, dict]):
         The backed will join the specified mapping metadata to the exisiting metadata.
         If there is a key-collision, the value given in the mapping will take precedence.
 
+        This method may also be used to udpate the `camera_params` for a particular set of items.
+        Just specify the key `camera_params` in the metadata for each reference_id along with all the necessary fields.
+
         Args:
             mapping: key-value pair of <reference_id>: <metadata>
 
         Examples:
-            >>> mapping = {"item_ref_1": {"new_key": "foo"}, "item_ref_2": {"some_value": 123}}
+            >>> mapping = {"item_ref_1": {"new_key": "foo"}, "item_ref_2": {"some_value": 123, "camera_params": {...}}}
             >>> dataset.update_item_metadata(mapping)
 
         Returns:
diff --git a/nucleus/dataset_item.py b/nucleus/dataset_item.py
@@ -5,174 +5,22 @@
 from enum import Enum
 from typing import Any, Dict, Optional, Sequence
 
-from .annotation import Point3D, is_local_path
+from .annotation import is_local_path
+from .camera_params import CameraParams
 from .constants import (
     BACKEND_REFERENCE_ID_KEY,
-    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,
     TYPE_KEY,
     UPLOAD_TO_SCALE_KEY,
     URL_KEY,
-    W_KEY,
-    X_KEY,
-    Y_KEY,
-    Z_KEY,
 )
 
 
-class CameraModels(str, 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.
-
-    We use the Hamilton/right-handed quaternion convention, where
-    ::
-
-        i^2 = j^2 = k^2 = ijk = -1
-
-    The quaternion represented by the tuple ``(x, y, z, w)`` is equal to
-    ``w + x*i + y*j + z*k``.
-
-    Parameters:
-        x (float): The x value.
-        y (float): The y value.
-        x (float): The z value.
-        w (float): The w value.
-    """
-
-    x: float
-    y: float
-    z: float
-    w: float
-
-    @classmethod
-    def from_json(cls, payload: Dict[str, float]):
-        """Instantiates quaternion object from schematized JSON dict payload."""
-        return cls(
-            payload[X_KEY], payload[Y_KEY], payload[Z_KEY], payload[W_KEY]
-        )
-
-    def to_payload(self) -> dict:
-        """Serializes quaternion object to schematized JSON dict."""
-        return {
-            X_KEY: self.x,
-            Y_KEY: self.y,
-            Z_KEY: self.z,
-            W_KEY: self.w,
-        }
-
-
-@dataclass
-class CameraParams:
-    """Camera position/heading used to record the image.
-
-    Args:
-        position (:class:`Point3D`): World-normalized position of the camera
-        heading (:class:`Quaternion`): Vector4 indicating the quaternion of the
-          camera direction; note that the z-axis of the camera frame
-          represents the camera's optical axis.  See `Heading Examples
-          <https://docs.scale.com/reference/data-types-and-the-frame-objects#heading-examples>`_.
-        fx (float): Focal length in x direction (in pixels).
-        fy (float): Focal length in y direction (in pixels).
-        cx (float): Principal point x value.
-        cy (float): Principal point y value.
-    """
-
-    position: Point3D
-    heading: Quaternion
-    fx: float
-    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 (k for k 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]):
-        """Instantiates camera params object from schematized JSON dict payload."""
-        return cls(
-            Point3D.from_json(payload[POSITION_KEY]),
-            Quaternion.from_json(payload[HEADING_KEY]),
-            payload[FX_KEY],
-            payload[FY_KEY],
-            payload[CX_KEY],
-            payload[CY_KEY],
-            payload.get(CAMERA_MODEL_KEY, None),
-            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),
-        )
-
-    def to_payload(self) -> dict:
-        """Serializes camera params object to schematized JSON dict."""
-        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):
     IMAGE = "image"
     POINTCLOUD = "pointcloud"
diff --git a/nucleus/metadata_manager.py b/nucleus/metadata_manager.py
@@ -1,5 +1,8 @@
 from enum import Enum
-from typing import TYPE_CHECKING, Dict
+from typing import TYPE_CHECKING, Dict, Optional
+
+from .camera_params import CameraParams
+from .constants import CAMERA_PARAMS_KEY
 
 if TYPE_CHECKING:
     from . import NucleusClient
@@ -31,11 +34,24 @@ def __init__(
 
         self._payload = self._format_mappings()
 
+    def _extract_camera_params(self, metadata: dict) -> Optional[CameraParams]:
+        camera_params = metadata.get(CAMERA_PARAMS_KEY, None)
+        if camera_params is None:
+            return None
+        return CameraParams.from_json(camera_params)
+
     def _format_mappings(self):
-        payload = []
+        payloads = []
         for ref_id, meta in self.raw_mappings.items():
-            payload.append({"reference_id": ref_id, "metadata": meta})
-        return payload
+            payload = {"reference_id": ref_id, "metadata": meta}
+
+            if self.level.value == ExportMetadataType.DATASET_ITEMS.value:
+                camera_params = self._extract_camera_params(meta)
+                if camera_params:
+                    payload[CAMERA_PARAMS_KEY] = camera_params.to_payload()
+
+            payloads.append(payload)
+        return payloads
 
     def update(self):
         payload = {"metadata": self._payload, "level": self.level.value}
diff --git a/nucleus/quaternion.py b/nucleus/quaternion.py
diff --git a/pyproject.toml b/pyproject.toml
