72 lines
2.7 KiB
Python
72 lines
2.7 KiB
Python
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)
|