vuer-cli 0.0.3__py3-none-any.whl → 0.0.5__py3-none-any.whl

vuer_cli/scripts/ptc_utils.py

@@ -0,0 +1,434 @@
+"""Point cloud utilities for RGB-D visualization.
+
+Provides classes and functions for:
+- Loading camera intrinsics and transforms from extracted MCAP data
+- Interpolating robot poses over time
+- Backprojecting RGB-D images to 3D point clouds
+- Coordinate system transformations (OpenCV/ROS to Three.js)
+"""
+
+import json
+from pathlib import Path
+from typing import Dict, Optional, Tuple
+
+import numpy as np
+import pandas as pd
+
+
+def matrix_to_three_js(mat: np.ndarray) -> list:
+    """Convert a 4x4 transformation matrix to Three.js column-major format."""
+    CONVERSION_INDICES = [0, 4, 8, 12, 1, 5, 9, 13, 2, 6, 10, 14, 3, 7, 11, 15]
+    return mat.flatten()[CONVERSION_INDICES].tolist()
+
+
+class CoordinateTransform:
+    """Handle coordinate system transformations between OpenCV/ROS and Three.js.
+
+    OpenCV/ROS camera frame:
+        +X: right, +Y: down, +Z: forward
+
+    Three.js with up=[0, 0, 1]:
+        +X: right, +Y: backward, +Z: up
+
+    The transformation flips X and Z axes to convert between these systems.
+    """
+
+    def __init__(self):
+        """Initialize the coordinate transformation matrices."""
+        self.flip_xz_3x3 = np.array([[-1, 0, 0], [0, 1, 0], [0, 0, -1]], dtype=np.float64)
+        self.flip_xz_4x4 = np.eye(4, dtype=np.float64)
+        self.flip_xz_4x4[:3, :3] = self.flip_xz_3x3
+
+    def transform_points(self, points: np.ndarray) -> np.ndarray:
+        """Transform 3D points from camera frame to Three.js frame."""
+        return points @ self.flip_xz_3x3.T
+
+    def transform_pose(self, pose: np.ndarray, rotation_only: bool = False) -> np.ndarray:
+        """Transform a 4x4 pose matrix from camera frame to Three.js frame."""
+        if rotation_only:
+            result = pose.copy()
+            result[:3, :3] = self.flip_xz_3x3 @ pose[:3, :3]
+            return result
+        return self.flip_xz_4x4 @ pose @ self.flip_xz_4x4
+
+    def to_threejs_matrix(self, pose: np.ndarray, rotation_only: bool = False) -> list:
+        """Transform pose and convert to Three.js column-major format."""
+        transformed = self.transform_pose(pose, rotation_only=rotation_only)
+        return matrix_to_three_js(transformed)
+
+
+def load_camera_intrinsics(data_dir: Path, camera_name: str) -> Dict:
+    """Load camera intrinsics for a specific camera.
+
+    Args:
+        data_dir: Path to extracted data directory
+        camera_name: Name of the camera (e.g., 'nav_front_d455')
+
+    Returns:
+        dict: Camera intrinsics containing:
+            - width, height: int - Image resolution
+            - fx, fy: float - Focal lengths in pixels
+            - cx, cy: float - Principal point in pixels
+            - K: np.ndarray (3, 3) - Camera calibration matrix
+            - D: np.ndarray or None - Distortion coefficients
+    """
+    with open(data_dir / "camera_intrinsics.json") as f:
+        all_intrinsics = json.load(f)
+
+    key = f"{camera_name}_color_camera_info"
+    if key not in all_intrinsics:
+        raise ValueError(f"Camera {key} not found in intrinsics")
+
+    intr = all_intrinsics[key]
+    K = np.array(intr["K"]).reshape(3, 3)
+
+    return {
+        "width": intr["width"],
+        "height": intr["height"],
+        "fx": K[0, 0],
+        "fy": K[1, 1],
+        "cx": K[0, 2],
+        "cy": K[1, 2],
+        "K": K,
+        "D": np.array(intr["D"]) if intr["D"] else None,
+    }
+
+
+def load_robot_poses(data_dir: Path) -> "TransformInterpolator":
+    """Load robot pose transforms (shared by all cameras).
+
+    Args:
+        data_dir: Path to extracted data directory
+
+    Returns:
+        TransformInterpolator: Time-varying robot poses
+    """
+    transforms_path = data_dir / "transforms" / "transforms.csv"
+    odometry_path = data_dir / "odometry.csv"
+
+    robot_poses = None
+
+    if transforms_path.exists() and transforms_path.stat().st_size > 0:
+        tf = pd.read_csv(transforms_path)
+        robot_poses = tf[(tf["parent_frame"] == "localization") & (tf["child_frame"] == "base_link")].copy()
+        if len(robot_poses) == 0:
+            robot_poses = None
+
+    if robot_poses is None and odometry_path.exists():
+        odom = pd.read_csv(odometry_path)
+        robot_poses = pd.DataFrame(
+            {
+                "timestamp_s": odom["timestamp_s"],
+                "parent_frame": odom["frame_id"],
+                "child_frame": odom["child_frame_id"],
+                "translation_x": odom["position_x"],
+                "translation_y": odom["position_y"],
+                "translation_z": odom["position_z"],
+                "rotation_x": odom["orientation_x"],
+                "rotation_y": odom["orientation_y"],
+                "rotation_z": odom["orientation_z"],
+                "rotation_w": odom["orientation_w"],
+            }
+        )
+
+    if robot_poses is None or len(robot_poses) == 0:
+        raise FileNotFoundError(
+            f"Could not find robot pose data. Looked for:\n"
+            f"  - {transforms_path} (with 'localization -> base_link' transforms)\n"
+            f"  - {odometry_path}\n"
+            f"Please ensure you have extracted the MCAP data correctly."
+        )
+
+    return TransformInterpolator(robot_poses)
+
+
+def load_static_transforms(data_dir: Path, camera_name: str) -> "StaticTransformChain":
+    """Load static camera mounting transforms for a specific camera.
+
+    Args:
+        data_dir: Path to extracted data directory
+        camera_name: Name of the camera (e.g., 'nav_front_d455')
+
+    Returns:
+        StaticTransformChain: Chain of static transforms from base_link to camera
+    """
+    transforms_path = data_dir / "transforms" / "transforms.csv"
+
+    if transforms_path.exists() and transforms_path.stat().st_size > 0:
+        tf = pd.read_csv(transforms_path)
+        static_tf = tf[tf["topic"] == "/tf_static"].copy()
+        if len(static_tf) == 0:
+            static_tf = tf.copy()
+    else:
+        static_tf = pd.DataFrame(
+            columns=[
+                "timestamp_s",
+                "topic",
+                "parent_frame",
+                "child_frame",
+                "translation_x",
+                "translation_y",
+                "translation_z",
+                "rotation_x",
+                "rotation_y",
+                "rotation_z",
+                "rotation_w",
+            ]
+        )
+
+    return StaticTransformChain(static_tf, camera_name)
+
+
+class TransformInterpolator:
+    """Interpolate robot poses between transform messages."""
+
+    def __init__(self, transforms_df: pd.DataFrame):
+        self.df = transforms_df.sort_values("timestamp_s").copy()
+        self.times = self.df["timestamp_s"].values
+
+    def interpolate_pose(self, timestamp_s: float) -> Optional[np.ndarray]:
+        """Interpolate pose at given timestamp.
+
+        Returns:
+            4x4 transformation matrix or None if timestamp out of range
+        """
+        if timestamp_s < self.times[0] or timestamp_s > self.times[-1]:
+            return None
+
+        idx_after = np.searchsorted(self.times, timestamp_s)
+        if idx_after == 0:
+            idx_after = 1
+        idx_before = idx_after - 1
+
+        t_before = self.times[idx_before]
+        t_after = self.times[idx_after]
+
+        if t_after == t_before:
+            alpha = 0.0
+        else:
+            alpha = (timestamp_s - t_before) / (t_after - t_before)
+
+        pose_before = self.df.iloc[idx_before]
+        pose_after = self.df.iloc[idx_after]
+
+        trans_before = np.array(
+            [pose_before["translation_x"], pose_before["translation_y"], pose_before["translation_z"]]
+        )
+        trans_after = np.array([pose_after["translation_x"], pose_after["translation_y"], pose_after["translation_z"]])
+        translation = (1 - alpha) * trans_before + alpha * trans_after
+
+        quat_before = np.array(
+            [pose_before["rotation_x"], pose_before["rotation_y"], pose_before["rotation_z"], pose_before["rotation_w"]]
+        )
+        quat_after = np.array(
+            [pose_after["rotation_x"], pose_after["rotation_y"], pose_after["rotation_z"], pose_after["rotation_w"]]
+        )
+
+        quat = self.slerp(quat_before, quat_after, alpha)
+
+        T = np.eye(4)
+        T[:3, :3] = self.quat_to_matrix(quat)
+        T[:3, 3] = translation
+
+        return T
+
+    @staticmethod
+    def slerp(q1: np.ndarray, q2: np.ndarray, t: float) -> np.ndarray:
+        """Spherical linear interpolation between quaternions."""
+        q1 = q1 / np.linalg.norm(q1)
+        q2 = q2 / np.linalg.norm(q2)
+
+        dot = np.dot(q1, q2)
+
+        if dot < 0.0:
+            q2 = -q2
+            dot = -dot
+
+        if dot > 0.9995:
+            result = q1 + t * (q2 - q1)
+            return result / np.linalg.norm(result)
+
+        theta_0 = np.arccos(np.clip(dot, -1.0, 1.0))
+        theta = theta_0 * t
+
+        q3 = q2 - q1 * dot
+        q3 = q3 / np.linalg.norm(q3)
+
+        return q1 * np.cos(theta) + q3 * np.sin(theta)
+
+    @staticmethod
+    def quat_to_matrix(q: np.ndarray) -> np.ndarray:
+        """Convert quaternion (x, y, z, w) to rotation matrix."""
+        x, y, z, w = q
+        return np.array(
+            [
+                [1 - 2 * (y * y + z * z), 2 * (x * y - w * z), 2 * (x * z + w * y)],
+                [2 * (x * y + w * z), 1 - 2 * (x * x + z * z), 2 * (y * z - w * x)],
+                [2 * (x * z - w * y), 2 * (y * z + w * x), 1 - 2 * (x * x + y * y)],
+            ]
+        )
+
+
+class StaticTransformChain:
+    """Compose static transforms from base_link to camera optical frame."""
+
+    def __init__(self, transforms_df: pd.DataFrame, camera_name: str):
+        self.camera_name = camera_name
+        self.transforms = {}
+
+        for _, row in transforms_df.iterrows():
+            parent = row["parent_frame"]
+            child = row["child_frame"]
+
+            T = np.eye(4)
+            T[:3, 3] = [row["translation_x"], row["translation_y"], row["translation_z"]]
+
+            quat = np.array([row["rotation_x"], row["rotation_y"], row["rotation_z"], row["rotation_w"]])
+            T[:3, :3] = TransformInterpolator.quat_to_matrix(quat)
+
+            self.transforms[(parent, child)] = T
+
+        self.chain = self._build_chain(camera_name)
+
+    def _build_chain(self, camera_name: str) -> np.ndarray:
+        """Build complete transform chain from base_link to camera_optical_frame."""
+        target_frame = f"{camera_name}_color_optical_frame"
+
+        if ("base_link", target_frame) in self.transforms:
+            return self.transforms[("base_link", target_frame)]
+
+        chain_frames = [
+            "base_link",
+            f"{camera_name}_bottom_screw_link",
+            f"{camera_name}_link",
+            f"{camera_name}_color_frame",
+            f"{camera_name}_color_optical_frame",
+        ]
+
+        T_total = np.eye(4)
+        missing_transforms = []
+
+        for i in range(len(chain_frames) - 1):
+            parent = chain_frames[i]
+            child = chain_frames[i + 1]
+
+            if (parent, child) in self.transforms:
+                T_total = T_total @ self.transforms[(parent, child)]
+            else:
+                missing_transforms.append(f"{parent} -> {child}")
+
+        if missing_transforms:
+            return np.eye(4)
+
+        return T_total
+
+    def get_transform(self) -> np.ndarray:
+        """Get composed static transform."""
+        return self.chain
+
+
+class PointCloudGenerator:
+    """Generate point cloud from RGB-D images with camera poses."""
+
+    def __init__(self, data_dir: str = "demcap_data", camera_name: str = "nav_front_d455"):
+        self.data_dir = Path(data_dir)
+        self.camera_name = camera_name
+
+        self.intrinsics = load_camera_intrinsics(self.data_dir, camera_name)
+        self.robot_pose_interpolator = load_robot_poses(self.data_dir)
+        self.static_camera_transform = load_static_transforms(self.data_dir, camera_name)
+
+    def _find_rgb_depth_pairs(self) -> list:
+        """Find matching RGB and depth image pairs."""
+        rgb_dir = None
+        images_base = self.data_dir / "images"
+
+        possible_rgb_dirs = [
+            images_base / f"{self.camera_name}_color_image_compressed",
+            images_base / f"{self.camera_name}_color_image_raw",
+            images_base / self.camera_name,
+        ]
+
+        for candidate in possible_rgb_dirs:
+            if candidate.exists():
+                rgb_dir = candidate
+                break
+
+        if rgb_dir is None:
+            raise ValueError(f"No RGB directory found for {self.camera_name}. Tried: {[str(p) for p in possible_rgb_dirs]}")
+
+        rgb_files = sorted(rgb_dir.glob("*.jpg"))
+
+        depth_dir = self.data_dir / "depth"
+        depth_dirs = list(depth_dir.glob(f"{self.camera_name}*"))
+        if not depth_dirs:
+            raise ValueError(f"No depth directory found for {self.camera_name}")
+        depth_dir = depth_dirs[0]
+        depth_files = sorted(depth_dir.glob("*.png"))
+
+        rgb_times = {int(f.stem): f for f in rgb_files}
+        depth_times = {int(f.stem): f for f in depth_files}
+
+        pairs = []
+        depth_timestamps = np.array(list(depth_times.keys()))
+
+        for rgb_ts, rgb_file in rgb_times.items():
+            idx = np.argmin(np.abs(depth_timestamps - rgb_ts))
+            depth_ts = depth_timestamps[idx]
+            depth_file = depth_times[depth_ts]
+
+            time_diff_ms = abs(rgb_ts - depth_ts) / 1e6
+            if time_diff_ms < 50:
+                pairs.append(
+                    {
+                        "rgb_file": rgb_file,
+                        "depth_file": depth_file,
+                        "rgb_timestamp_ns": rgb_ts,
+                        "depth_timestamp_ns": depth_ts,
+                        "timestamp_s": rgb_ts / 1e9,
+                        "time_diff_ms": time_diff_ms,
+                    }
+                )
+
+        return pairs
+
+    def backproject_rgbd(
+        self, rgb: np.ndarray, depth: np.ndarray, max_depth: float = 10.0, pixel_step: int = 1
+    ) -> Tuple[np.ndarray, np.ndarray]:
+        """Backproject RGB-D image to 3D points in camera frame.
+
+        Args:
+            rgb: RGB image (H, W, 3)
+            depth: Depth image (H, W) in mm
+            max_depth: Maximum depth in meters
+            pixel_step: Sample every Nth pixel
+
+        Returns:
+            points: (N, 3) array of 3D points in camera frame
+            colors: (N, 3) array of RGB colors normalized to [0, 1]
+        """
+        depth_ds = depth[::pixel_step, ::pixel_step]
+        rgb_ds = rgb[::pixel_step, ::pixel_step]
+
+        h_ds, w_ds = depth_ds.shape
+
+        u, v = np.meshgrid(np.arange(w_ds), np.arange(h_ds))
+
+        depth_m = depth_ds.astype(np.float32) / 1000.0
+
+        valid = (depth_m > 0) & (depth_m < max_depth)
+
+        u = u[valid]
+        v = v[valid]
+        z = depth_m[valid]
+
+        u_original = u * pixel_step
+        v_original = v * pixel_step
+
+        x = (u_original - self.intrinsics["cx"]) * z / self.intrinsics["fx"]
+        y = (v_original - self.intrinsics["cy"]) * z / self.intrinsics["fy"]
+
+        points = np.stack([x, y, z], axis=-1)
+        colors = rgb_ds[valid] / 255.0
+
+        return points, colors