init commit

nimbus/components/data/camera.py

@@ -0,0 +1,71 @@
+from dataclasses import dataclass
+from typing import Any, Dict, List, Optional
+
+import numpy as np
+
+
+@dataclass
+class C2W:
+    """
+    Represents a camera-to-world transformation matrix.
+
+    Attributes:
+        matrix (List[float]): A list of 16 floats representing the 4x4 transformation matrix in row-major order.
+    """
+
+    matrix: List[float]
+
+
+@dataclass
+class Camera:
+    """
+    Represents a single camera pose in the trajectory.
+
+    Attributes:
+        trajectory (List[C2W]): List of C2W transformations for this camera pose.
+        intrinsic (Optional[List[float]]): 3x3 camera intrinsic matrix: [[fx, 0, cx], [0, fy, cy], [0, 0, 1]].
+        extrinsic (Optional[List[float]]): 4x4 tobase_extrinsic matrix representing the camera mounting offset
+            relative to the robot base (height + pitch).
+        length (Optional[int]): Length of the trajectory in number of frames.
+        depths (Optional[list[np.ndarray]]): List of depth images captured by this camera.
+        rgbs (Optional[list[np.ndarray]]): List of RGB images captured by this camera.
+        uv_tracks (Optional[Dict[str, Any]]): UV tracking data in the format
+            {mesh_name: {"per_frame": list, "width": W, "height": H}}.
+        uv_mesh_names (Optional[List[str]]): List of mesh names being tracked in the UV tracking data.
+    """
+
+    trajectory: List[C2W]
+    intrinsic: List[float] = None
+    extrinsic: List[float] = None
+    length: int = None
+    depths: list[np.ndarray] = None
+    rgbs: list[np.ndarray] = None
+    uv_tracks: Optional[Dict[str, Any]] = None
+    uv_mesh_names: Optional[List[str]] = None
+
+    def __len__(self):
+        if self.length is not None:
+            return self.length
+        self._check_length()
+        self.length = len(self.trajectory)
+        return len(self.trajectory)
+
+    def _check_length(self):
+        if self.depths is not None and len(self.depths) != len(self.trajectory):
+            raise ValueError("Length of depths does not match length of trajectory")
+        if self.rgbs is not None and len(self.rgbs) != len(self.trajectory):
+            raise ValueError("Length of rgbs does not match length of trajectory")
+        if self.uv_tracks is not None:
+            for mesh_name, track_data in self.uv_tracks.items():
+                if len(track_data["per_frame"]) != len(self.trajectory):
+                    raise ValueError(f"Length of uv_tracks for mesh {mesh_name} does not match length of trajectory")
+
+    def append_rgb(self, rgb_image: np.ndarray):
+        if self.rgbs is None:
+            self.rgbs = []
+        self.rgbs.append(rgb_image)
+
+    def append_depth(self, depth_image: np.ndarray):
+        if self.depths is None:
+            self.depths = []
+        self.depths.append(depth_image)