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

vuer_cli/scripts/minimap.py

@@ -0,0 +1,365 @@
+"""Minimap visualization with GLB model and depth point clouds."""
+
+import asyncio
+import warnings
+from dataclasses import dataclass
+from pathlib import Path
+from textwrap import dedent
+from typing import List, Optional
+
+from vuer.schemas import (
+  DepthPointCloudProvider,
+  InPlaceDepthPointCloud,
+)
+
+from ..utils import print_error
+
+warnings.filterwarnings("ignore", category=RuntimeWarning)
+
+
+def collect_depth_frames(
+  data_dir: str,
+  camera_names: Optional[List[str]] = None,
+  max_frames: Optional[int] = None,
+  frame_step: int = 5,
+):
+  """Collect depth frame paths and their world transforms."""
+  import numpy as np
+
+  from .ptc_utils import PointCloudGenerator, load_robot_poses, load_static_transforms
+
+  data_path = Path(data_dir)
+  if not data_path.exists():
+    print(f"  Data path does not exist: {data_path}")
+    return []
+
+  if camera_names is None:
+    camera_names = ["nav_front_d455", "nav_right_d455"]
+
+  print(f"  Looking for cameras: {camera_names}")
+
+  try:
+    robot_pose_interpolator = load_robot_poses(data_path)
+    pose_time_min = robot_pose_interpolator.times[0]
+    pose_time_max = robot_pose_interpolator.times[-1]
+    print(f"  Loaded {len(robot_pose_interpolator.times)} robot poses")
+  except Exception as e:
+    print(f"  Failed to load robot poses: {e}")
+    return []
+
+  frames = []
+  reference_camera = camera_names[0]
+  first_camera_transform = np.eye(4)
+  first_transform_set = False
+
+  for cam_name in camera_names:
+    try:
+      print(f"  Processing camera: {cam_name}")
+      static_transform = load_static_transforms(data_path, cam_name)
+      generator = PointCloudGenerator(data_dir=str(data_path), camera_name=cam_name)
+
+      pairs = generator._find_rgb_depth_pairs()
+      pairs = [p for p in pairs if pose_time_min <= p["timestamp_s"] <= pose_time_max]
+      print(f"    Found {len(pairs)} RGB-D pairs in time range")
+
+      if max_frames:
+        pairs = pairs[:max_frames:frame_step]
+      else:
+        pairs = pairs[::frame_step]
+
+      print(f"    Using {len(pairs)} pairs after sampling")
+
+      for pair in pairs:
+        t_world_to_base = robot_pose_interpolator.interpolate_pose(pair["timestamp_s"])
+        if t_world_to_base is None:
+          continue
+
+        t_base_to_camera = static_transform.get_transform()
+        t_world_to_camera = t_world_to_base @ t_base_to_camera
+
+        if cam_name == reference_camera and not first_transform_set:
+          first_camera_transform = t_world_to_camera.copy()
+          first_transform_set = True
+
+        # Transform to reference frame
+        t_ref = np.linalg.inv(first_camera_transform) @ t_world_to_camera
+
+        # Flip XY for three.js coordinate system
+        flip_xy = np.array([[-1, 0, 0, 0], [0, -1, 0, 0], [0, 0, 1, 0], [0, 0, 0, 1]])
+        t_final = flip_xy @ t_ref
+
+        # Extract camera position and rotation
+        position = t_final[:3, 3].tolist()
+
+        from scipy.spatial.transform import Rotation
+
+        rot_matrix = t_final[:3, :3]
+        rotation = Rotation.from_matrix(rot_matrix).as_euler("xyz").tolist()
+
+        # Compute baselink transform (robot base, not camera)
+        t_base_ref = np.linalg.inv(first_camera_transform) @ t_world_to_base
+        t_base_final = flip_xy @ t_base_ref
+        baselink_position = t_base_final[:3, 3].tolist()
+        baselink_rot_matrix = t_base_final[:3, :3]
+        baselink_rotation = (
+          Rotation.from_matrix(baselink_rot_matrix).as_euler("xyz").tolist()
+        )
+
+        frames.append(
+          {
+            "depth_file": pair["depth_file"],
+            "camera_name": cam_name,
+            "timestamp_s": pair["timestamp_s"],
+            "position": position,
+            "rotation": rotation,
+            "baselink_position": baselink_position,
+            "baselink_rotation": baselink_rotation,
+            "intrinsics": generator.intrinsics,
+          }
+        )
+    except Exception as e:
+      print(f"    Error processing {cam_name}: {e}")
+      continue
+
+  return frames
+
+
+async def main_viz(
+  sess,
+  app,
+  scene_store,
+  depth_frames: List[dict],
+  cmap: str = "turbo",
+  color_mode: str = "depth",
+  fov: float = 58.0,
+  max_depth: float = 10.0,
+  glb_url: str = None,
+  fps: float = 10.0,
+  loop: bool = True,
+):
+  """Main visualization with frame-by-frame depth point cloud and robot model playback."""
+  import math
+
+  from vuer.schemas import Glb
+
+  print("Setting up scene...")
+
+  # Link all depth files upfront
+  for i, frame in enumerate(depth_frames):
+    depth_filename = f"depth/{i}.png"
+    app.workspace.link(Path(frame["depth_file"]), f"/{depth_filename}")
+
+  # Construct proper GLB URL using workspace prefix
+  glb_full_url = str(app.localhost_prefix / glb_url)
+
+  # clear the default scene (removes lighting / grid)
+  # scene_store.set @ Scene(up=[0, 1, 0])
+  # Set initial scene from the scene store
+  # sess.upsert @ scene_store.scene
+  await asyncio.sleep(0.01)
+
+  frame_delay = 1.0 / fps
+  frame_idx = 0
+
+  print(f"Starting playback: {len(depth_frames)} frames at {fps} FPS (loop={loop})")
+
+  while True:
+    frame = depth_frames[frame_idx]
+    depth_filename = f"depth/{frame_idx}.png"
+
+    # Compute FOV from intrinsics if available
+    frame_fov = fov
+    if frame["intrinsics"]:
+      fx = frame["intrinsics"].get("fx", None)
+      height = frame["intrinsics"].get("height", None)
+      if fx and height:
+        frame_fov = 2 * math.atan(height / (2 * fx)) * 180 / math.pi
+
+    depth_pc = InPlaceDepthPointCloud(
+      key="depth-pc-current",
+      depth=app.localhost_prefix / depth_filename,
+      cmap=cmap,
+      colorMode=color_mode,
+      fov=frame_fov,
+      position=frame["position"],
+      rotation=frame["rotation"],
+      maxDepth=max_depth,
+    )
+
+    scene_store.upsert @ DepthPointCloudProvider(
+      depth_pc,
+      key="depth-provider",
+      frustumCulling=True,
+    )
+
+    # Update robot model position based on baselink transform
+    baselink_pos = frame["baselink_position"]
+    baselink_rot = frame["baselink_rotation"]
+
+    scene_store.upsert @ Glb(
+      src=glb_full_url,
+      position=[baselink_pos[0], baselink_pos[1] + 2, baselink_pos[2]],
+      rotation=[baselink_rot[0] - 1.57, baselink_rot[1], baselink_rot[2]],
+      scale=[0.01, 0.01, 0.01],
+      key="proxie-model",
+    )
+
+    await asyncio.sleep(frame_delay)
+
+    frame_idx += 1
+    if frame_idx >= len(depth_frames):
+      if loop:
+        frame_idx = 0
+      else:
+        break
+
+  print("Playback complete.")
+
+
+@dataclass
+class Minimap:
+  """Visualize robot trajectory with depth point clouds in a minimap view.
+
+  DESCRIPTION
+      Creates an interactive 3D visualization combining a GLB robot model with
+      depth-based point clouds using DepthPointCloud components. The robot model
+      moves frame-by-frame following the baselink trajectory.
+      Uses Workspace path linking for efficient depth image streaming.
+
+      Features:
+      - High-performance depth point cloud rendering with frustum culling
+      - Robot model that follows the trajectory frame-by-frame
+      - Colormap support (turbo, viridis, inferno, jet)
+      - Multiple color modes (depth, camZ, camDist, localY, worldY)
+
+      The visualization server runs at http://localhost:8012 by default.
+
+  INPUT FORMAT
+      Requires extracted MCAP data (use 'vuer demcap' first):
+          <data_dir>/
+          ├── images/<camera>_color_image_*/    # RGB images
+          ├── depth/<camera>_*/                  # Depth images
+          ├── transforms/transforms.csv          # Robot poses
+          └── camera_intrinsics.json             # Camera calibration
+
+  EXAMPLES
+      # Basic visualization
+      vuer minimap --data-dir /path/to/extracted_data
+
+      # Custom colormap and color mode
+      vuer minimap --data-dir /path/to/data --cmap viridis --color-mode worldY
+
+      # Sparse sampling for faster loading
+      vuer minimap --data-dir /path/to/data --frame-step 10
+
+  DEPENDENCIES
+      Requires: pip install 'vuer-cli[viz]'
+  """
+
+  data_dir: str  # Path to extracted MCAP data directory
+
+  # Depth point cloud options
+  cameras: str = None  # Comma-separated camera names
+  max_frames: int = None  # Max frames per camera
+  frame_step: int = 5  # Process every Nth frame
+  cmap: str = "turbo"  # Colormap (turbo, viridis, inferno, jet)
+  color_mode: str = "worldY"  # Color mode (depth, camZ, camDist, localY, worldY)
+  fov: float = 58.0  # Default camera FOV (RealSense D435)
+  max_depth: float = 10.0  # Maximum depth in meters
+
+  # Model options
+  glb_url: str = "proxie.glb"  # GLB model URL
+
+  # Playback options
+  fps: float = 10.0  # Frames per second for playback
+  loop: bool = True  # Loop playback
+
+  # Server options
+  host: str = "0.0.0.0"
+  port: int = 8012
+
+  def __call__(self) -> int:
+    """Execute minimap command."""
+    try:
+      from vuer import Vuer
+      from vuer.rtc.scene_store import SceneStore
+    except ImportError as e:
+      print_error(
+        f"Missing dependencies for minimap command: {e}\n\n"
+        "Install with:\n"
+        "  pip install 'vuer-cli[viz]'\n"
+        "  # or: uv add 'vuer-cli[viz]'"
+      )
+      return 1
+
+    try:
+      data_path = Path(self.data_dir)
+      if not data_path.exists():
+        raise FileNotFoundError(f"Data directory {self.data_dir} does not exist")
+
+      camera_list = None
+      if self.cameras:
+        camera_list = [c.strip() for c in self.cameras.split(",")]
+
+      # Collect depth frames
+      print("Collecting depth frames...")
+      depth_frames = collect_depth_frames(
+        data_dir=self.data_dir,
+        camera_names=camera_list,
+        max_frames=self.max_frames,
+        frame_step=self.frame_step,
+      )
+
+      if not depth_frames:
+        print_error("No depth frames found in the data directory")
+        return 1
+
+      print(f"Found {len(depth_frames)} depth frames")
+
+      # Use current working directory as workspace
+      app = Vuer(host=self.host, port=self.port, workspace=Path.cwd())
+
+      print(
+        dedent(f"""
+        Starting minimap visualization server...
+          URL: http://localhost:{self.port}
+          Depth frames: {len(depth_frames)}
+          Colormap: {self.cmap}
+          Color mode: {self.color_mode}
+          FPS: {self.fps}
+          Loop: {self.loop}
+      """).strip()
+      )
+
+      scene_store = SceneStore()
+
+      @app.spawn(start=True)
+      async def main_with_args(sess):
+        print("Session started, loading data...")
+        async with scene_store.subscribe(sess):
+          await main_viz(
+            sess,
+            app=app,
+            scene_store=scene_store,
+            depth_frames=depth_frames,
+            cmap=self.cmap,
+            color_mode=self.color_mode,
+            fov=self.fov,
+            max_depth=self.max_depth,
+            glb_url=self.glb_url,
+            fps=self.fps,
+            loop=self.loop,
+          )
+          await sess.forever()
+
+      return 0
+
+    except FileNotFoundError as e:
+      print_error(str(e))
+      return 1
+    except Exception as e:
+      print_error(f"Unexpected error: {e}")
+      import traceback
+
+      traceback.print_exc()
+      return 1