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

vuer_cli/scripts/viz_ptc_proxie.py

@@ -0,0 +1,483 @@
+"""Visualize GLB model with camera frustums and optional point cloud."""
+
+import warnings
+from dataclasses import dataclass
+from pathlib import Path
+from typing import List, Optional
+
+from ..utils import print_error
+
+# Suppress numpy warnings for matrix operations
+warnings.filterwarnings("ignore", category=RuntimeWarning)
+
+
+def create_camera_positions_around_center(radius=1.0, height=0.0):
+  """Create four camera positions around a center point at 90-degree intervals."""
+  import numpy as np
+
+  positions = []
+  for angle in [0, np.pi / 2, np.pi, 3 * np.pi / 2]:
+    x = radius * np.cos(angle)
+    y = radius * np.sin(angle)
+    z = height
+    positions.append([x, y, z])
+  return positions
+
+
+def create_camera_rotations_looking_at_center():
+  """Create rotations for cameras to look at the center point."""
+  import numpy as np
+
+  rotations = [
+    [0, -np.pi / 2, np.pi / 2],
+    [np.pi / 2, 0, 0],
+    [-np.pi / 2, np.pi / 2, 0],
+    [np.pi / 2, -np.pi, np.pi],
+  ]
+  return rotations
+
+
+def load_scene_pointcloud(
+  data_dir: str = "demcap_data",
+  camera_names: Optional[List[str]] = None,
+  max_frames: Optional[int] = None,
+  frame_step: int = 5,
+  pixel_step: int = 5,
+  max_depth: float = 10.0,
+  color_by_camera: bool = False,
+):
+  """Load and merge point clouds from multiple cameras into a single scene point cloud."""
+  import cv2
+  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():
+    return None, None
+
+  if camera_names is None:
+    camera_names = ["nav_front_d455", "nav_right_d455"]
+
+  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]
+  except Exception:
+    return None, None
+
+  camera_color_map = {
+    "nav_front_d455": np.array([0.0, 0.0, 1.0], dtype=np.float32),
+    "nav_right_d455": np.array([1.0, 0.0, 0.0], dtype=np.float32),
+    "nav_left_d455": np.array([0.0, 1.0, 0.0], dtype=np.float32),
+    "nav_rear_d455": np.array([1.0, 1.0, 0.0], dtype=np.float32),
+  }
+
+  all_points = []
+  all_colors = []
+
+  reference_camera = camera_names[0]
+  first_camera_transform = np.eye(4)
+
+  for cam_name in camera_names:
+    try:
+      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]
+
+      if max_frames:
+        pairs = pairs[:max_frames:frame_step]
+      else:
+        pairs = pairs[::frame_step]
+
+      if cam_name == reference_camera and len(pairs) > 0:
+        first_pair = pairs[0]
+        t_world_to_base = robot_pose_interpolator.interpolate_pose(
+          first_pair["timestamp_s"]
+        )
+        if t_world_to_base is not None:
+          t_base_to_camera = static_transform.get_transform()
+          first_camera_transform = t_world_to_base @ t_base_to_camera
+
+      for frame_idx, pair in enumerate(pairs):
+        rgb = cv2.imread(str(pair["rgb_file"]))
+        rgb = cv2.cvtColor(rgb, cv2.COLOR_BGR2RGB)
+        depth = cv2.imread(str(pair["depth_file"]), cv2.IMREAD_UNCHANGED)
+
+        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
+
+        points_cam, colors = generator.backproject_rgbd(
+          rgb, depth, max_depth, pixel_step
+        )
+
+        if len(points_cam) == 0:
+          continue
+
+        if color_by_camera:
+          camera_color = camera_color_map.get(
+            cam_name, np.array([0.5, 0.5, 0.5], dtype=np.float32)
+          )
+          colors = np.tile(camera_color, (len(points_cam), 1))
+
+        points_cam_h = np.hstack([points_cam, np.ones((len(points_cam), 1))])
+        points_world_h = (t_world_to_camera @ points_cam_h.T).T
+        points_ref_h = (np.linalg.inv(first_camera_transform) @ points_world_h.T).T
+
+        flip_xy = np.array([[-1, 0, 0, 0], [0, -1, 0, 0], [0, 0, 1, 0], [0, 0, 0, 1]])
+        points_ref_threejs_h = (flip_xy @ points_ref_h.T).T
+        points_ref = points_ref_threejs_h[:, :3]
+
+        all_points.append(points_ref.astype(np.float32))
+        all_colors.append(colors.astype(np.float32))
+
+    except Exception:
+      continue
+
+  if len(all_points) == 0:
+    return None, None
+
+  combined_points = np.vstack(all_points)
+  combined_colors = np.vstack(all_colors)
+
+  total_points_before = len(combined_points)
+
+  y_values = combined_points[:, 1]
+  y_threshold = np.percentile(y_values, 65)
+  mask = y_values < y_threshold
+  combined_points = combined_points[mask]
+  combined_colors = combined_colors[mask]
+
+  return combined_points, combined_colors
+
+
+def create_proxie_group(
+  group_id: int,
+  initial_position=None,
+  initial_rotation=None,
+  opacity: float = 1.0,
+  glb_url: str = None,
+):
+  """Create a single group containing a GLB model with four camera frustums."""
+  from vuer.schemas import Frustum, Glb, Group
+
+  if initial_position is None:
+    initial_position = [0, 0, 0]
+  if initial_rotation is None:
+    initial_rotation = [0, 0, 0]
+
+  camera_positions = create_camera_positions_around_center(radius=0.1, height=1.625)
+  camera_rotations = create_camera_rotations_looking_at_center()
+
+  frustums = []
+  for i, (pos, rot) in enumerate(zip(camera_positions, camera_rotations, strict=False)):
+    if i == 3:
+      pos[0] += 0.185
+      pos[1] += 0.02
+    elif i == 0:
+      pos[0] += 0.20
+      pos[1] += 0.008
+    elif i == 1:
+      pos[0] += 0.17
+      pos[1] += -0.025
+    elif i == 2:
+      pos[0] += 0.18
+      pos[1] += 0.005
+    frustums.append(
+      Frustum(
+        key=f"frustum-{group_id}-{i}",
+        position=pos,
+        rotation=rot,
+        scale=0.1,
+        fov=60,
+        aspect=16 / 9,
+        near=0.1,
+        far=0.2,
+        showFrustum=True,
+        showImagePlane=True,
+        showFocalPlane=False,
+        colorFrustum="cyan",
+      )
+    )
+
+  initial_position[0] += 0
+  initial_position[1] += 2
+  initial_position[2] += 0
+  initial_rotation[0] += -1.57
+  initial_rotation[1] += 0
+  initial_rotation[2] += 0
+
+  return Group(
+    Glb(
+      src=glb_url,
+      position=[0, 0, 0],
+      rotation=[1.57, 0, 0],
+      scale=0.01,
+      opacity=opacity,
+      key=f"model-{group_id}",
+    ),
+    *frustums,
+    position=initial_position,
+    rotation=initial_rotation,
+    scale=1.0,
+    key=f"group-{group_id}",
+  )
+
+
+async def main_viz(
+  sess,
+  show_pointcloud: bool = True,
+  data_dir: str = "demcap_data",
+  camera_names: Optional[List[str]] = None,
+  max_frames: Optional[int] = None,
+  frame_step: int = 5,
+  pixel_step: int = 5,
+  max_depth: float = 10.0,
+  point_size: float = 0.01,
+  color_by_camera: bool = False,
+  num_groups: int = 3,
+  group_spacing: float = 2.0,
+  base_position: Optional[List[float]] = None,
+  glb_url: str = None,
+):
+  """Main visualization function that creates multiple groups with GLB models and cameras."""
+  from vuer.schemas import DefaultScene, PointCloud
+
+  print("Loading scene...")
+
+  point_cloud_vertices = None
+  point_cloud_colors = None
+  if show_pointcloud:
+    print("Loading point cloud...")
+    point_cloud_vertices, point_cloud_colors = load_scene_pointcloud(
+      data_dir=data_dir,
+      camera_names=camera_names,
+      max_frames=max_frames,
+      frame_step=frame_step,
+      pixel_step=pixel_step,
+      max_depth=max_depth,
+      color_by_camera=color_by_camera,
+    )
+
+  if num_groups > 1:
+    opacity_values = [1.0 - (0.3 * i / (num_groups - 1)) for i in range(num_groups)]
+  else:
+    opacity_values = [1.0]
+
+  group_configs = [
+    {
+      "id": i,
+      "position": [
+        base_position[0],
+        base_position[1],
+        base_position[2] - i * group_spacing,
+      ],
+      "rotation": [0, 0, 0],
+      "opacity": opacity_values[i],
+    }
+    for i in range(num_groups)
+  ]
+
+  groups = [
+    create_proxie_group(
+      config["id"],
+      config["position"],
+      config.get("rotation"),
+      config.get("opacity", 1.0),
+      glb_url=glb_url,
+    )
+    for config in group_configs
+  ]
+
+  scene_elements = []
+  scene_elements.extend(groups)
+
+  if point_cloud_vertices is not None and point_cloud_colors is not None:
+    print(f"Point cloud: {len(point_cloud_vertices):,} points")
+    point_cloud = PointCloud(
+      key="scene_pointcloud",
+      vertices=point_cloud_vertices,
+      colors=point_cloud_colors,
+      size=point_size,
+      position=[0, 2, 0],
+      rotation=[0, 0, 0],
+    )
+    scene_elements.append(point_cloud)
+
+  print("Setting up scene...")
+  sess.set @ DefaultScene(
+    *scene_elements,
+    up=[0, 1, 0],
+    # show_helper=False,
+    # bgChildren=[OrbitControls(key="OrbitControls")],
+  )
+
+  print("Ready! Open browser to view.")
+  await sess.forever()
+
+
+@dataclass
+class VizPtcProxie:
+  """Visualize GLB robot model with camera frustums and optional point cloud.
+
+  DESCRIPTION
+      Creates an interactive 3D visualization combining a GLB robot model with
+      camera frustums and an optional RGB-D point cloud background. Useful for
+      visualizing robot poses within a reconstructed environment.
+
+      Features:
+      - Multiple robot model instances with configurable spacing
+      - Four camera frustums per robot showing sensor coverage
+      - Optional point cloud background from RGB-D data
+      - Opacity gradient for depth perception of robot poses
+
+      The visualization server runs at http://localhost:8012 by default.
+      Open this URL in a browser to view the 3D scene.
+
+  GLB MODEL
+      Place your robot GLB model at: vuer-cli/assets/proxie.glb
+      Or specify a custom URL with --glb-url
+
+  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 with point cloud
+      vuer viz_ptc_proxie --data-dir /path/to/extracted_data
+
+      # Without point cloud (faster startup)
+      vuer viz_ptc_proxie --data-dir /path/to/data --show-pointcloud False
+
+      # Custom number of robot instances
+      vuer viz_ptc_proxie --data-dir /path/to/data --num-groups 5 --group-spacing 1.0
+
+      # Custom GLB model URL
+      vuer viz_ptc_proxie --data-dir /path/to/data --glb-url http://localhost:8012/static/my_robot.glb
+
+      # Specify cameras for point cloud
+      vuer viz_ptc_proxie --data-dir /path/to/data --cameras nav_front_d455,nav_right_d455
+
+      # Faster point cloud with sparser sampling
+      vuer viz_ptc_proxie --data-dir /path/to/data --frame-step 10 --pixel-step 10
+
+  DEPENDENCIES
+      Requires: pip install 'vuer-cli[viz]'
+  """
+
+  # Required: path to extracted MCAP data directory
+  data_dir: str
+
+  # Point cloud options
+  show_pointcloud: bool = True  # Enable point cloud visualization
+  cameras: str = (
+    None  # Comma-separated camera names (default: nav_front_d455,nav_right_d455)
+  )
+  max_frames: int = None  # Max frames per camera (default: all)
+  frame_step: int = 5  # Process every Nth frame
+  pixel_step: int = 10  # Sample every Nth pixel
+  max_depth: float = 10.0  # Maximum depth in meters
+  point_size: float = 0.01  # Size of points
+
+  # Model options
+  assets_folder: str = (
+    None  # Static assets folder for GLB model (default: vuer-cli/assets)
+  )
+  glb_url: str = "proxie.glb"  # GLB model URL
+  num_groups: int = 10  # Number of proxie groups to create
+  group_spacing: float = 0.5  # Spacing between groups along Z-axis
+  base_position: str = "-2.0,-2.2,0"  # Starting position [x,y,z] for first group
+
+  # Server options
+  host: str = "0.0.0.0"  # Server host
+  port: int = 8012  # Server port
+
+  def __call__(self) -> int:
+    """Execute viz_ptc_proxie command."""
+    try:
+      import cv2
+      import numpy as np
+      import pandas as pd
+      from vuer import Vuer
+    except ImportError as e:
+      print_error(
+        f"Missing dependencies for viz_ptc_proxie 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")
+
+      # Parse camera names
+      camera_list = None
+      if self.cameras:
+        camera_list = [c.strip() for c in self.cameras.split(",")]
+
+      # Parse base_position
+      base_pos_list = [float(x.strip()) for x in self.base_position.split(",")]
+
+      # Create Vuer app - use default assets folder if not specified
+      assets_folder = self.assets_folder
+      if not assets_folder:
+        # Default to vuer-cli/assets folder
+        import vuer_cli
+
+        assets_folder = Path(vuer_cli.__file__).parent.parent.parent / "assets"
+        if not assets_folder.exists():
+          assets_folder = None
+
+      if assets_folder:
+        print(f"Using assets folder: {assets_folder}")
+        app = Vuer(host=self.host, port=self.port, workspace=assets_folder)
+      else:
+        app = Vuer(host=self.host, port=self.port)
+
+      print("Starting visualization server...")
+
+      @app.spawn(start=True)
+      async def main_with_args(sess):
+        print("Session started, loading data...")
+        await main_viz(
+          sess,
+          show_pointcloud=self.show_pointcloud,
+          data_dir=self.data_dir,
+          camera_names=camera_list,
+          max_frames=self.max_frames,
+          frame_step=self.frame_step,
+          pixel_step=self.pixel_step,
+          max_depth=self.max_depth,
+          point_size=self.point_size,
+          color_by_camera=False,
+          num_groups=self.num_groups,
+          group_spacing=self.group_spacing,
+          base_position=base_pos_list,
+          glb_url=self.glb_url,
+        )
+
+      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