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

vuer_cli/scripts/mcap_playback.py

@@ -0,0 +1,661 @@
+"""MCAP Playback - Stream MCAP data to Vuer clients or server downlink.
+
+This module provides a playback system for MCAP files that generates typed events
+for each message type and can push them to Vuer clients via websocket or directly
+to the server's downlink pipe.
+
+Usage:
+    from vuer_cli.scripts.mcap_playback import McapPlayback, RgbImageEvent
+
+    playback = McapPlayback("/path/to/file.mcap")
+
+    # Register handlers for specific event types
+    @playback.on(RgbImageEvent)
+    async def handle_rgb(event):
+        print(f"RGB image from {event.camera_name} at {event.timestamp_s}")
+
+    # Start playback
+    await playback.play(fps=30, loop=True)
+"""
+
+import asyncio
+from dataclasses import dataclass, field
+from pathlib import Path
+from typing import Any, Callable, Dict, List, Optional, Type, TypeVar
+
+import numpy as np
+
+
+# =============================================================================
+# Event Types
+# =============================================================================
+
+
+@dataclass
+class McapEvent:
+    """Base class for all MCAP events."""
+
+    timestamp_ns: int
+    timestamp_s: float
+    topic: str
+
+    @property
+    def event_type(self) -> str:
+        return self.__class__.__name__
+
+
+@dataclass
+class RgbImageEvent(McapEvent):
+    """RGB image data from a camera."""
+
+    camera_name: str
+    data: bytes  # JPEG or raw image bytes
+    width: int = 0
+    height: int = 0
+    encoding: str = "jpeg"
+
+
+@dataclass
+class DepthImageEvent(McapEvent):
+    """Depth image data from a camera."""
+
+    camera_name: str
+    data: bytes  # PNG or raw depth bytes
+    width: int = 0
+    height: int = 0
+    encoding: str = "16UC1"
+
+
+@dataclass
+class ImuEvent(McapEvent):
+    """IMU sensor data."""
+
+    orientation: List[float] = field(default_factory=lambda: [0, 0, 0, 1])  # xyzw
+    angular_velocity: List[float] = field(default_factory=lambda: [0, 0, 0])
+    linear_acceleration: List[float] = field(default_factory=lambda: [0, 0, 0])
+
+
+@dataclass
+class TransformEvent(McapEvent):
+    """TF transform data."""
+
+    parent_frame: str = ""
+    child_frame: str = ""
+    translation: List[float] = field(default_factory=lambda: [0, 0, 0])
+    rotation: List[float] = field(default_factory=lambda: [0, 0, 0, 1])  # xyzw
+
+
+@dataclass
+class OdometryEvent(McapEvent):
+    """Odometry data."""
+
+    frame_id: str = ""
+    child_frame_id: str = ""
+    position: List[float] = field(default_factory=lambda: [0, 0, 0])
+    orientation: List[float] = field(default_factory=lambda: [0, 0, 0, 1])
+    linear_velocity: List[float] = field(default_factory=lambda: [0, 0, 0])
+    angular_velocity: List[float] = field(default_factory=lambda: [0, 0, 0])
+
+
+@dataclass
+class JointStatesEvent(McapEvent):
+    """Joint states data."""
+
+    frame_id: str = ""
+    names: List[str] = field(default_factory=list)
+    positions: List[float] = field(default_factory=list)
+    velocities: List[float] = field(default_factory=list)
+    efforts: List[float] = field(default_factory=list)
+
+
+@dataclass
+class LidarEvent(McapEvent):
+    """LiDAR point cloud data."""
+
+    data: bytes = b""
+    point_count: int = 0
+
+
+@dataclass
+class CameraInfoEvent(McapEvent):
+    """Camera intrinsics information."""
+
+    camera_name: str = ""
+    width: int = 0
+    height: int = 0
+    K: List[float] = field(default_factory=list)  # 3x3 intrinsic matrix
+    D: List[float] = field(default_factory=list)  # distortion coefficients
+
+
+# Type alias for event handlers
+E = TypeVar("E", bound=McapEvent)
+EventHandler = Callable[[McapEvent], Any]
+
+
+# =============================================================================
+# MCAP Playback Class
+# =============================================================================
+
+
+class McapPlayback:
+    """Playback MCAP files and dispatch events to registered handlers.
+
+    Supports both callback-based handlers and direct streaming to Vuer sessions.
+
+    Example:
+        playback = McapPlayback("/path/to/recording.mcap")
+
+        @playback.on(RgbImageEvent)
+        async def handle_rgb(event):
+            sess.upsert @ ImagePlane(event.data, key=event.camera_name)
+
+        await playback.play(fps=30)
+    """
+
+    def __init__(
+        self,
+        mcap_path: str,
+        start_time_ns: Optional[int] = None,
+        end_time_ns: Optional[int] = None,
+    ):
+        self.mcap_path = Path(mcap_path)
+        self.start_time_ns = start_time_ns
+        self.end_time_ns = end_time_ns
+
+        self._handlers: Dict[Type[McapEvent], List[EventHandler]] = {}
+        self._global_handlers: List[EventHandler] = []
+
+        # Lazy-loaded decoder
+        self._decoder_factory = None
+
+    def on(self, event_type: Type[E]) -> Callable[[EventHandler], EventHandler]:
+        """Decorator to register a handler for a specific event type.
+
+        Example:
+            @playback.on(RgbImageEvent)
+            async def handle_rgb(event):
+                print(event.camera_name)
+        """
+
+        def decorator(handler: EventHandler) -> EventHandler:
+            if event_type not in self._handlers:
+                self._handlers[event_type] = []
+            self._handlers[event_type].append(handler)
+            return handler
+
+        return decorator
+
+    def on_all(self, handler: EventHandler) -> EventHandler:
+        """Register a handler that receives all events."""
+        self._global_handlers.append(handler)
+        return handler
+
+    async def _dispatch(self, event: McapEvent):
+        """Dispatch an event to all registered handlers."""
+        # Type-specific handlers
+        handlers = self._handlers.get(type(event), [])
+        for handler in handlers:
+            result = handler(event)
+            if asyncio.iscoroutine(result):
+                await result
+
+        # Global handlers
+        for handler in self._global_handlers:
+            result = handler(event)
+            if asyncio.iscoroutine(result):
+                await result
+
+    def _get_decoder_factory(self):
+        """Lazy-load the decoder factory."""
+        if self._decoder_factory is None:
+            from mcap_ros2.decoder import DecoderFactory
+
+            self._decoder_factory = DecoderFactory()
+        return self._decoder_factory
+
+    def _sanitize_topic_name(self, topic: str) -> str:
+        """Convert topic name to a clean identifier."""
+        return topic.replace("/", "_").strip("_")
+
+    def _parse_rgb_image(self, topic, message, schema, channel) -> Optional[RgbImageEvent]:
+        """Parse RGB image message into event."""
+        camera_name = self._sanitize_topic_name(topic)
+        timestamp_ns = message.log_time
+        timestamp_s = timestamp_ns / 1e9
+
+        if "image_compressed" in topic:
+            # Compressed JPEG
+            data = message.data[4:]
+            jpeg_start = data.find(b"\xff\xd8")
+            if jpeg_start == -1:
+                return None
+            jpeg_data = data[jpeg_start:]
+            return RgbImageEvent(
+                timestamp_ns=timestamp_ns,
+                timestamp_s=timestamp_s,
+                topic=topic,
+                camera_name=camera_name,
+                data=jpeg_data,
+                encoding="jpeg",
+            )
+        else:
+            # Raw image
+            try:
+                decoder = self._get_decoder_factory().decoder_for(channel.message_encoding, schema)
+                msg = decoder(message.data)
+                return RgbImageEvent(
+                    timestamp_ns=timestamp_ns,
+                    timestamp_s=timestamp_s,
+                    topic=topic,
+                    camera_name=camera_name,
+                    data=bytes(msg.data),
+                    width=msg.width,
+                    height=msg.height,
+                    encoding=msg.encoding,
+                )
+            except Exception:
+                return None
+
+    def _parse_depth_image(self, topic, message, schema, channel) -> Optional[DepthImageEvent]:
+        """Parse depth image message into event."""
+        camera_name = self._sanitize_topic_name(topic)
+        timestamp_ns = message.log_time
+        timestamp_s = timestamp_ns / 1e9
+
+        if "compressedDepth" in topic:
+            # Compressed PNG
+            data = message.data[4:]
+            png_start = data.find(b"\x89PNG")
+            if png_start == -1:
+                return None
+            png_data = data[png_start:]
+            return DepthImageEvent(
+                timestamp_ns=timestamp_ns,
+                timestamp_s=timestamp_s,
+                topic=topic,
+                camera_name=camera_name,
+                data=png_data,
+                encoding="png",
+            )
+        else:
+            # Raw depth
+            try:
+                decoder = self._get_decoder_factory().decoder_for(channel.message_encoding, schema)
+                msg = decoder(message.data)
+                return DepthImageEvent(
+                    timestamp_ns=timestamp_ns,
+                    timestamp_s=timestamp_s,
+                    topic=topic,
+                    camera_name=camera_name,
+                    data=bytes(msg.data),
+                    width=msg.width,
+                    height=msg.height,
+                    encoding=msg.encoding,
+                )
+            except Exception:
+                return None
+
+    def _parse_imu(self, topic, message, schema, channel) -> Optional[ImuEvent]:
+        """Parse IMU message into event."""
+        timestamp_ns = message.log_time
+        timestamp_s = timestamp_ns / 1e9
+
+        if not schema or "Imu" not in schema.name:
+            return None
+
+        try:
+            decoder = self._get_decoder_factory().decoder_for(channel.message_encoding, schema)
+            msg = decoder(message.data)
+            return ImuEvent(
+                timestamp_ns=timestamp_ns,
+                timestamp_s=timestamp_s,
+                topic=topic,
+                orientation=[msg.orientation.x, msg.orientation.y, msg.orientation.z, msg.orientation.w],
+                angular_velocity=[msg.angular_velocity.x, msg.angular_velocity.y, msg.angular_velocity.z],
+                linear_acceleration=[
+                    msg.linear_acceleration.x,
+                    msg.linear_acceleration.y,
+                    msg.linear_acceleration.z,
+                ],
+            )
+        except Exception:
+            return None
+
+    def _parse_transform(self, topic, message, schema, channel) -> List[TransformEvent]:
+        """Parse TF message into events (one per transform)."""
+        timestamp_ns = message.log_time
+        timestamp_s = timestamp_ns / 1e9
+        events = []
+
+        try:
+            decoder = self._get_decoder_factory().decoder_for(channel.message_encoding, schema)
+            msg = decoder(message.data)
+
+            if hasattr(msg, "transforms"):
+                for tf in msg.transforms:
+                    events.append(
+                        TransformEvent(
+                            timestamp_ns=timestamp_ns,
+                            timestamp_s=timestamp_s,
+                            topic=topic,
+                            parent_frame=tf.header.frame_id,
+                            child_frame=tf.child_frame_id,
+                            translation=[
+                                tf.transform.translation.x,
+                                tf.transform.translation.y,
+                                tf.transform.translation.z,
+                            ],
+                            rotation=[
+                                tf.transform.rotation.x,
+                                tf.transform.rotation.y,
+                                tf.transform.rotation.z,
+                                tf.transform.rotation.w,
+                            ],
+                        )
+                    )
+        except Exception:
+            pass
+
+        return events
+
+    def _parse_odometry(self, topic, message, schema, channel) -> Optional[OdometryEvent]:
+        """Parse odometry message into event."""
+        timestamp_ns = message.log_time
+        timestamp_s = timestamp_ns / 1e9
+
+        try:
+            decoder = self._get_decoder_factory().decoder_for(channel.message_encoding, schema)
+            msg = decoder(message.data)
+            return OdometryEvent(
+                timestamp_ns=timestamp_ns,
+                timestamp_s=timestamp_s,
+                topic=topic,
+                frame_id=msg.header.frame_id,
+                child_frame_id=msg.child_frame_id,
+                position=[msg.pose.pose.position.x, msg.pose.pose.position.y, msg.pose.pose.position.z],
+                orientation=[
+                    msg.pose.pose.orientation.x,
+                    msg.pose.pose.orientation.y,
+                    msg.pose.pose.orientation.z,
+                    msg.pose.pose.orientation.w,
+                ],
+                linear_velocity=[
+                    msg.twist.twist.linear.x,
+                    msg.twist.twist.linear.y,
+                    msg.twist.twist.linear.z,
+                ],
+                angular_velocity=[
+                    msg.twist.twist.angular.x,
+                    msg.twist.twist.angular.y,
+                    msg.twist.twist.angular.z,
+                ],
+            )
+        except Exception:
+            return None
+
+    def _parse_joint_states(self, topic, message, schema, channel) -> Optional[JointStatesEvent]:
+        """Parse joint states message into event."""
+        timestamp_ns = message.log_time
+        timestamp_s = timestamp_ns / 1e9
+
+        try:
+            decoder = self._get_decoder_factory().decoder_for(channel.message_encoding, schema)
+            msg = decoder(message.data)
+            return JointStatesEvent(
+                timestamp_ns=timestamp_ns,
+                timestamp_s=timestamp_s,
+                topic=topic,
+                frame_id=msg.header.frame_id,
+                names=list(msg.name),
+                positions=list(msg.position) if msg.position else [],
+                velocities=list(msg.velocity) if msg.velocity else [],
+                efforts=list(msg.effort) if msg.effort else [],
+            )
+        except Exception:
+            return None
+
+    def _parse_camera_info(self, topic, message, schema, channel) -> Optional[CameraInfoEvent]:
+        """Parse camera info message into event."""
+        camera_name = self._sanitize_topic_name(topic)
+        timestamp_ns = message.log_time
+        timestamp_s = timestamp_ns / 1e9
+
+        try:
+            decoder = self._get_decoder_factory().decoder_for(channel.message_encoding, schema)
+            msg = decoder(message.data)
+            return CameraInfoEvent(
+                timestamp_ns=timestamp_ns,
+                timestamp_s=timestamp_s,
+                topic=topic,
+                camera_name=camera_name,
+                width=msg.width,
+                height=msg.height,
+                K=list(msg.k) if hasattr(msg, "k") else [],
+                D=list(msg.d) if hasattr(msg, "d") else [],
+            )
+        except Exception:
+            return None
+
+    def _parse_lidar(self, topic, message) -> Optional[LidarEvent]:
+        """Parse lidar message into event."""
+        timestamp_ns = message.log_time
+        timestamp_s = timestamp_ns / 1e9
+
+        try:
+            data = message.data[4:]
+            return LidarEvent(
+                timestamp_ns=timestamp_ns,
+                timestamp_s=timestamp_s,
+                topic=topic,
+                data=data,
+            )
+        except Exception:
+            return None
+
+    async def iter_events(self):
+        """Iterate through all events in the MCAP file."""
+        from mcap.reader import make_reader
+
+        with open(self.mcap_path, "rb") as f:
+            reader = make_reader(f)
+
+            for schema, channel, message in reader.iter_messages():
+                topic = channel.topic
+
+                # Apply time filter
+                if self.start_time_ns is not None and message.log_time < self.start_time_ns:
+                    continue
+                if self.end_time_ns is not None and message.log_time > self.end_time_ns:
+                    continue
+
+                # Route to parser based on topic
+                events = []
+
+                if "compressedDepth" in topic or "depth/image" in topic:
+                    event = self._parse_depth_image(topic, message, schema, channel)
+                    if event:
+                        events.append(event)
+
+                elif "color/image" in topic:
+                    event = self._parse_rgb_image(topic, message, schema, channel)
+                    if event:
+                        events.append(event)
+
+                elif "imu" in topic.lower():
+                    event = self._parse_imu(topic, message, schema, channel)
+                    if event:
+                        events.append(event)
+
+                elif "/tf" in topic:
+                    events.extend(self._parse_transform(topic, message, schema, channel))
+
+                elif "/odom" in topic:
+                    event = self._parse_odometry(topic, message, schema, channel)
+                    if event:
+                        events.append(event)
+
+                elif "joint_states" in topic.lower():
+                    event = self._parse_joint_states(topic, message, schema, channel)
+                    if event:
+                        events.append(event)
+
+                elif "camera_info" in topic:
+                    event = self._parse_camera_info(topic, message, schema, channel)
+                    if event:
+                        events.append(event)
+
+                elif "pandar_points" in topic or "hesai" in topic:
+                    event = self._parse_lidar(topic, message)
+                    if event:
+                        events.append(event)
+
+                for event in events:
+                    yield event
+
+    async def play(
+        self,
+        fps: float = 30.0,
+        loop: bool = False,
+        speed: float = 1.0,
+    ):
+        """Play back the MCAP file at the specified rate.
+
+        Args:
+            fps: Target frames per second
+            loop: Whether to loop playback
+            speed: Playback speed multiplier (1.0 = realtime)
+        """
+        frame_delay = 1.0 / fps
+
+        while True:
+            last_time_ns = None
+
+            async for event in self.iter_events():
+                # Dispatch to handlers
+                await self._dispatch(event)
+
+                # Timing control
+                if last_time_ns is not None:
+                    delta_ns = event.timestamp_ns - last_time_ns
+                    delta_s = delta_ns / 1e9 / speed
+                    if delta_s > 0 and delta_s < 1.0:
+                        await asyncio.sleep(min(delta_s, frame_delay))
+                    else:
+                        await asyncio.sleep(frame_delay)
+                else:
+                    await asyncio.sleep(frame_delay)
+
+                last_time_ns = event.timestamp_ns
+
+            if not loop:
+                break
+
+    def push_to_session(self, sess):
+        """Create handlers that push events directly to a Vuer session.
+
+        This registers default handlers for common event types that update
+        the Vuer scene.
+
+        Args:
+            sess: VuerSession instance
+        """
+
+        @self.on(RgbImageEvent)
+        async def push_rgb(event: RgbImageEvent):
+            from vuer.schemas import ImagePlane
+
+            sess.upsert @ ImagePlane(
+                src=event.data,
+                key=f"rgb-{event.camera_name}",
+            )
+
+        @self.on(DepthImageEvent)
+        async def push_depth(event: DepthImageEvent):
+            from vuer.schemas import DepthPointCloud
+
+            sess.upsert @ DepthPointCloud(
+                src=event.data,
+                key=f"depth-{event.camera_name}",
+            )
+
+        @self.on(TransformEvent)
+        async def push_transform(event: TransformEvent):
+            # Transform events can update object positions
+            # This is a placeholder - actual implementation depends on scene structure
+            pass
+
+        return self
+
+
+# =============================================================================
+# Vuer Integration
+# =============================================================================
+
+
+class VuerMcapPlayer:
+    """High-level integration of MCAP playback with Vuer server.
+
+    Example:
+        player = VuerMcapPlayer("/path/to/file.mcap")
+
+        @app.spawn(start=True)
+        async def main(sess):
+            await player.stream_to_session(sess, fps=30)
+    """
+
+    def __init__(
+        self,
+        mcap_path: str,
+        start_time_ns: Optional[int] = None,
+        end_time_ns: Optional[int] = None,
+    ):
+        self.playback = McapPlayback(mcap_path, start_time_ns, end_time_ns)
+
+    def on(self, event_type: Type[E]) -> Callable[[EventHandler], EventHandler]:
+        """Register an event handler."""
+        return self.playback.on(event_type)
+
+    async def stream_to_session(
+        self,
+        sess,
+        fps: float = 30.0,
+        loop: bool = False,
+        speed: float = 1.0,
+    ):
+        """Stream MCAP data to a Vuer session.
+
+        Args:
+            sess: VuerSession instance
+            fps: Target frames per second
+            loop: Whether to loop playback
+            speed: Playback speed multiplier
+        """
+        self.playback.push_to_session(sess)
+        await self.playback.play(fps=fps, loop=loop, speed=speed)
+
+    async def push_to_downlink(
+        self,
+        server,
+        fps: float = 30.0,
+        loop: bool = False,
+        speed: float = 1.0,
+    ):
+        """Push MCAP data directly to server's downlink pipe.
+
+        This bypasses the websocket and pushes events directly to all
+        connected clients via the server's internal downlink.
+
+        Args:
+            server: Vuer server instance
+            fps: Target frames per second
+            loop: Whether to loop playback
+            speed: Playback speed multiplier
+        """
+        # Register handlers that use server.broadcast or server.downlink
+        @self.playback.on_all
+        async def broadcast(event: McapEvent):
+            # Convert event to Vuer schema and broadcast
+            # This is a placeholder - actual implementation depends on server API
+            pass
+
+        await self.playback.play(fps=fps, loop=loop, speed=speed)