diff --git a/integrations/omnidreams/omnidreams/interactive_drive/README.md b/integrations/omnidreams/omnidreams/interactive_drive/README.md index 83c8914df..93a63dc4e 100644 --- a/integrations/omnidreams/omnidreams/interactive_drive/README.md +++ b/integrations/omnidreams/omnidreams/interactive_drive/README.md @@ -467,6 +467,28 @@ Controls (apply in all three modes): - `X` exit scene (return to the scene selector; HUD mode only) - `Esc` quit +Manifests can enable rollout recording with either backend: + +```yaml +recording_enabled: true +recording_dir: recordings +recording_hotkey: F9 +recording_auto_start: false +``` + +Press the configured hotkey once to start and again to stop. Set +`recording_auto_start: true` to start recording as soon as each rollout begins. +Relative `recording_dir` values are written under the flashdreams repository +root; absolute paths are used as-is. Each saved recording writes +`first_frame.png`, `prompt.txt`, `metadata.json`, `hdmap.mp4`, and +`inferred.mp4` under a timestamped subdirectory. `first_frame.png` is copied +from the first retained inferred frame so it matches the inferred video. To +bound host memory during long auto-start rollouts, each stream keeps the most +recent 600 frames; `metadata.json` records any dropped frame counts. +Raster-only recordings contain the HD-map/raster stream, save `first_frame.png` +from that stream, and omit `inferred.mp4` because there is no world-model +output. + The browser control hint is static today, so it does not confirm every keydown visually. If the world-model backend is still producing a chunk, input can be accepted before the visual response arrives. diff --git a/integrations/omnidreams/omnidreams/interactive_drive/app.py b/integrations/omnidreams/omnidreams/interactive_drive/app.py index 4e3643064..f4261fd70 100644 --- a/integrations/omnidreams/omnidreams/interactive_drive/app.py +++ b/integrations/omnidreams/omnidreams/interactive_drive/app.py @@ -16,6 +16,7 @@ KeyboardState, ) from omnidreams.interactive_drive.presenter import SlangPyPresenter +from omnidreams.interactive_drive.recording import InteractiveDriveRecorder from omnidreams.interactive_drive.runtime.loop import ( LoopConfig, PresenterBackend, @@ -82,7 +83,10 @@ def __init__( """ self._config = config self._backend = backend - self._keyboard = KeyboardState() + self._keyboard = KeyboardState( + recording_enabled=config.recording.enabled, + recording_hotkey=config.recording.hotkey, + ) if config.backend == "omnidreams": self._keyboard.set_view_mode("model_rgb") if presenter is None: @@ -438,60 +442,77 @@ def run_scene(self) -> None: # world model..."); subsequent rollouts come from a manual reset or # OOB respawn, so switch the indicator to "Resetting..." for those. loading_status = self._loading_status_message - while not self._presenter.should_close: - simulation = EgoVehicleKinematics( - initial_state=state_from_initial_pose( - initial_rig_to_world=self._scene.initial_rig_to_world, - initial_yaw_rad=self._scene.initial_yaw_rad, - # Start each rollout at a fixed 10 m/s so the ego is - # already rolling on initial load (and after a manual - # reset / OOB respawn), instead of launching at the - # clip's full recorded speed. - initial_speed_mps=10.0, - ), - vehicle_config=self._config.vehicle, - ground_snapper=self._ground_snapper, - initial_timestamp_us=self._scene.initial_timestamp_us, - map_bounds=self._map_bounds, - oob_margin_m=self._config.oob_margin_m, - oob_warning_zone_m=self._config.oob_warning_zone_m, - ) - # Publish the freshly-built initial state up front so read-side - # speed readouts (the HUD speed digit, the browser ``/state`` - # endpoint) reflect a reset / respawn immediately. Without this - # the last telemetry from the previous rollout would linger on - # screen through the "Resetting..." window until the new rollout - # requested its first chunk -- the "reset doesn't reset the - # displayed speed" symptom. - self._keyboard.update_telemetry(simulation.current_state) - input_backend = KeyboardInputBackend(self._keyboard) - reset_requested = run_main_loop( - presenter=self._presenter, - runtime_controls=self._keyboard, - initial_presented_frame=loading_frame, - input_backend=input_backend, - simulation=simulation, - pipeline=self._pipeline, - config=LoopConfig( - initial_chunk_size=self._config.chunk.initial_chunk_frames, - chunk_size=self._config.chunk.chunk_frames, - frame_interval_s=self._config.chunk.frame_interval_s, - oob_warn_proximity=self._config.oob_warn_proximity, - oob_respawn_proximity=self._config.oob_respawn_proximity, - oob_respawn_debounce_chunks=( - self._config.oob_respawn_debounce_chunks + recorder = self._build_recorder() + try: + while not self._presenter.should_close: + simulation = EgoVehicleKinematics( + initial_state=state_from_initial_pose( + initial_rig_to_world=self._scene.initial_rig_to_world, + initial_yaw_rad=self._scene.initial_yaw_rad, + # Start each rollout at a fixed 10 m/s so the ego is + # already rolling on initial load (and after a manual + # reset / OOB respawn), instead of launching at the + # clip's full recorded speed. + initial_speed_mps=10.0, ), - ), - loading_status=loading_status, - ) - if not reset_requested: - break - self._pipeline.reset() - loading_status = self._resetting_status_message - # Paint the reset indicator at once, before the next rollout's - # setup, so a reset shows on screen the instant it's requested - # rather than after the rebuild completes. - self._present_loading_once(loading_status) + vehicle_config=self._config.vehicle, + ground_snapper=self._ground_snapper, + initial_timestamp_us=self._scene.initial_timestamp_us, + map_bounds=self._map_bounds, + oob_margin_m=self._config.oob_margin_m, + oob_warning_zone_m=self._config.oob_warning_zone_m, + ) + # Publish the freshly-built initial state up front so read-side + # speed readouts (the HUD speed digit, the browser ``/state`` + # endpoint) reflect a reset / respawn immediately. Without this + # the last telemetry from the previous rollout would linger on + # screen through the "Resetting..." window until the new rollout + # requested its first chunk -- the "reset doesn't reset the + # displayed speed" symptom. + self._keyboard.update_telemetry(simulation.current_state) + input_backend = KeyboardInputBackend(self._keyboard) + reset_requested = run_main_loop( + presenter=self._presenter, + runtime_controls=self._keyboard, + initial_presented_frame=loading_frame, + input_backend=input_backend, + simulation=simulation, + pipeline=self._pipeline, + config=LoopConfig( + initial_chunk_size=self._config.chunk.initial_chunk_frames, + chunk_size=self._config.chunk.chunk_frames, + frame_interval_s=self._config.chunk.frame_interval_s, + oob_warn_proximity=self._config.oob_warn_proximity, + oob_respawn_proximity=self._config.oob_respawn_proximity, + oob_respawn_debounce_chunks=( + self._config.oob_respawn_debounce_chunks + ), + ), + loading_status=loading_status, + recorder=recorder, + ) + if not reset_requested: + break + self._pipeline.reset() + loading_status = self._resetting_status_message + # Paint the reset indicator at once, before the next rollout's + # setup, so a reset shows on screen the instant it's requested + # rather than after the rebuild completes. + self._present_loading_once(loading_status) + finally: + if recorder is not None: + recorder.close(reason="scene-end") + + def _build_recorder(self) -> InteractiveDriveRecorder | None: + if not self._config.recording.enabled: + return None + if self._scene is None: + return None + return InteractiveDriveRecorder( + self._config.recording, + scene=self._scene, + fps=self._config.chunk.fps, + ) def _present_loading_once(self, loading_status: Callable[[], str]) -> None: """Render a single loading-overlay frame immediately (used on reset).""" diff --git a/integrations/omnidreams/omnidreams/interactive_drive/cli.py b/integrations/omnidreams/omnidreams/interactive_drive/cli.py index 6d2e49588..9a25a9d68 100644 --- a/integrations/omnidreams/omnidreams/interactive_drive/cli.py +++ b/integrations/omnidreams/omnidreams/interactive_drive/cli.py @@ -4,6 +4,7 @@ from __future__ import annotations import argparse +import warnings from dataclasses import replace from pathlib import Path @@ -21,8 +22,13 @@ WorldModelProfileConfig, ) from omnidreams.interactive_drive.log import configure_logging +from omnidreams.interactive_drive.recording import RecordingConfig from omnidreams.interactive_drive.synthetic_scene import build_synthetic_scene_to_temp -from omnidreams.interactive_drive.world_model.manifest import load_world_model_manifest +from omnidreams.interactive_drive.world_model.manifest import ( + RecordingManifest, + load_recording_manifest, + load_world_model_manifest, +) from omnidreams.scenes import local_scene_archive_path # Package root (from this file's location) so packaged-asset defaults below @@ -31,6 +37,26 @@ # ``$FLASHDREAMS_CACHE_DIR/omnidreams-scenes/`` (shared with the webrtc server). _PACKAGE_ROOT = Path(__file__).resolve().parent _CONFIGS_ROOT = _PACKAGE_ROOT / "configs" +_DEFAULT_RECORDING_DIR_NAME = "recordings" + + +def _find_flashdreams_root(start: Path) -> Path: + for candidate in (start, *start.parents): + if (candidate / "pyproject.toml").is_file() and ( + candidate / "integrations" / "omnidreams" + ).is_dir(): + return candidate + fallback = Path.cwd().resolve() + warnings.warn( + "Could not locate the flashdreams repository root from " + f"{start}; resolving relative recording_dir values from {fallback}.", + RuntimeWarning, + stacklevel=2, + ) + return fallback + + +_FLASHDREAMS_ROOT = _find_flashdreams_root(_PACKAGE_ROOT) # Default scene UUID staged by ``omnidreams-prepare`` (clear-weather base # archive in nvidia/omni-dreams-scenes). @@ -69,6 +95,21 @@ def resolve_manifest_path(path: str | Path) -> Path: return cwd_path +def _resolve_recording_output_dir( + raw_dir: Path | None, *, enabled: bool +) -> Path | None: + if not enabled: + return None + path = ( + Path(_DEFAULT_RECORDING_DIR_NAME) + if raw_dir is None + else Path(raw_dir).expanduser() + ) + if path.is_absolute(): + return path + return (_FLASHDREAMS_ROOT / path).resolve() + + def build_parser() -> argparse.ArgumentParser: parser = argparse.ArgumentParser( description="Single-process flashdreams driving demo" @@ -148,8 +189,10 @@ def build_parser() -> argparse.ArgumentParser: type=Path, default=None, help=( - "Omnidreams pipeline manifest (YAML). Accepts a path or a bundled " - "config filename such as example_world_model_perf.yaml." + "Interactive-drive manifest (YAML). For the omnidreams backend this " + "also configures the world-model pipeline; for raster, recording_* " + "fields are honored when present. Accepts a path or a bundled config " + "filename such as example_world_model_perf.yaml." ), ) parser.add_argument( @@ -369,6 +412,18 @@ def _oob_kwargs(args: argparse.Namespace) -> dict[str, float | int]: return overrides +def _recording_config_from_manifest(manifest: RecordingManifest) -> RecordingConfig: + return RecordingConfig( + enabled=manifest.enabled, + output_dir=_resolve_recording_output_dir( + manifest.dir, + enabled=manifest.enabled, + ), + hotkey=manifest.hotkey, + auto_start=manifest.auto_start, + ) + + def main() -> None: """Stand-alone entry point for ``python -m omnidreams.interactive_drive.cli``. @@ -448,6 +503,13 @@ def prepare_config_and_backend( backend: RenderBackend if config.backend == "raster": + if config.manifest_path is not None: + config = replace( + config, + recording=_recording_config_from_manifest( + load_recording_manifest(config.manifest_path) + ), + ) backend = RasterRenderBackend( chunk=config.chunk, raster=config.raster, bev=config.bev ) @@ -468,6 +530,17 @@ def prepare_config_and_backend( height=manifest.resolution_wh[1], ), ) + config = replace( + config, + recording=_recording_config_from_manifest( + RecordingManifest( + enabled=manifest.recording_enabled, + dir=manifest.recording_dir, + hotkey=manifest.recording_hotkey, + auto_start=manifest.recording_auto_start, + ) + ), + ) backend = WorldModelRenderBackend( manifest=manifest, chunk=config.chunk, diff --git a/integrations/omnidreams/omnidreams/interactive_drive/config.py b/integrations/omnidreams/omnidreams/interactive_drive/config.py index a26f069b7..3e18fae19 100644 --- a/integrations/omnidreams/omnidreams/interactive_drive/config.py +++ b/integrations/omnidreams/omnidreams/interactive_drive/config.py @@ -7,6 +7,8 @@ from pathlib import Path from typing import Literal +from omnidreams.interactive_drive.recording import RecordingConfig + BackendName = Literal["raster", "omnidreams"] ViewMode = Literal["rgb", "model_rgb"] ComputeDeviceName = Literal["automatic", "cuda", "vulkan"] @@ -110,6 +112,7 @@ class AppConfig: vehicle: VehicleConfig = VehicleConfig() world_model_profile: WorldModelProfileConfig = WorldModelProfileConfig() world_model_offload_text_encoder: bool = False + recording: RecordingConfig = RecordingConfig() bev: BevConfig = BevConfig() # OOB thresholds plumbed to LoopConfig (overridable via CLI --oob-*). # Match alpasim's driver-side proximity: warn > 0.6, respawn >= 2.0 diff --git a/integrations/omnidreams/omnidreams/interactive_drive/configs/example_world_model_perf.yaml b/integrations/omnidreams/omnidreams/interactive_drive/configs/example_world_model_perf.yaml index 335ffa8b7..a9d106fce 100644 --- a/integrations/omnidreams/omnidreams/interactive_drive/configs/example_world_model_perf.yaml +++ b/integrations/omnidreams/omnidreams/interactive_drive/configs/example_world_model_perf.yaml @@ -58,3 +58,9 @@ native_dit_attention_backend: cudnn # auto | cudnn | sparge | sage3 | sage3_fp8 # or uncomment and set an absolute or manifest-relative path below. native_vae_encoder: disabled # disabled | fp8 # native_vae_fp8_state_path: /path/to/lightvae_fp8_state.pt + + +recording_enabled: true +recording_dir: recordings +recording_hotkey: F9 +recording_auto_start: true diff --git a/integrations/omnidreams/omnidreams/interactive_drive/input/keyboard.py b/integrations/omnidreams/omnidreams/interactive_drive/input/keyboard.py index 18c610875..ccc8b612e 100644 --- a/integrations/omnidreams/omnidreams/interactive_drive/input/keyboard.py +++ b/integrations/omnidreams/omnidreams/interactive_drive/input/keyboard.py @@ -5,6 +5,7 @@ import time from omnidreams.interactive_drive.input.backend import InputBackend, SampledInput +from omnidreams.interactive_drive.recording import normalize_recording_hotkey from omnidreams.interactive_drive.types import ( ControlSnapshot, DriverCommand, @@ -24,12 +25,17 @@ class KeyboardState: referencing the per-scene simulation object. """ - def __init__(self) -> None: + def __init__( + self, *, recording_enabled: bool = False, recording_hotkey: str = "f9" + ) -> None: self._lock = threading.Lock() self._pressed: set[str] = set() self._view_mode = "rgb" self._drive_command: DriverCommand | None = None self._reset_pending = False + self._recording_enabled = bool(recording_enabled) + self._recording_hotkey = normalize_recording_hotkey(recording_hotkey) + self._recording_toggle_pending = False # Rising-edge "exit the current scene and return to the scene # selector" request, set by a wheel/controller's bound exit button # (the HUD's ``x`` key calls the presenter directly). The presenter @@ -57,6 +63,27 @@ def request_reset(self) -> None: with self._lock: self._reset_pending = True + def request_recording_toggle(self) -> None: + with self._lock: + if self._recording_enabled: + self._recording_toggle_pending = True + + def consume_recording_toggle_request(self) -> bool: + with self._lock: + pending = self._recording_toggle_pending + self._recording_toggle_pending = False + return pending + + @property + def recording_enabled(self) -> bool: + with self._lock: + return self._recording_enabled + + @property + def recording_hotkey(self) -> str: + with self._lock: + return self._recording_hotkey + def request_exit_scene(self) -> None: """Request a return to the scene selector from a bound device button.""" with self._lock: diff --git a/integrations/omnidreams/omnidreams/interactive_drive/presenter.py b/integrations/omnidreams/omnidreams/interactive_drive/presenter.py index 623dd756f..d2e26b072 100644 --- a/integrations/omnidreams/omnidreams/interactive_drive/presenter.py +++ b/integrations/omnidreams/omnidreams/interactive_drive/presenter.py @@ -13,6 +13,7 @@ ) from omnidreams.interactive_drive.input.keyboard import KeyboardState from omnidreams.interactive_drive.loading_overlay import render_loading_overlay +from omnidreams.interactive_drive.recording import slangpy_key_name_candidates from omnidreams.interactive_drive.types import PresentedFrame @@ -326,6 +327,8 @@ def _on_keyboard_event(self, event) -> None: self._keyboard.set_view_mode("rgb") elif is_press and self._matches_key(event.key, "r"): self._keyboard.request_reset() + elif is_press and self._matches_recording_hotkey(event.key): + self._keyboard.request_recording_toggle() def _build_key_codes(self) -> dict[str, object | None]: return { @@ -354,6 +357,14 @@ def _matches_key(self, event_key: object, name: str) -> bool: key_code = self._key_codes.get(name) return key_code is not None and event_key == key_code + def _matches_recording_hotkey(self, event_key: object) -> bool: + if not self._keyboard.recording_enabled: + return False + for name in slangpy_key_name_candidates(self._keyboard.recording_hotkey): + if event_key == self._lookup_key_code(name): + return True + return False + class _CudaRGBInterop: def __init__(self, *, spy: Any, device: Any, width: int, height: int) -> None: diff --git a/integrations/omnidreams/omnidreams/interactive_drive/recording.py b/integrations/omnidreams/omnidreams/interactive_drive/recording.py new file mode 100644 index 000000000..49166242d --- /dev/null +++ b/integrations/omnidreams/omnidreams/interactive_drive/recording.py @@ -0,0 +1,301 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + +from __future__ import annotations + +import json +import re +from dataclasses import dataclass +from datetime import datetime, timezone +from pathlib import Path +from typing import Any + +import numpy as np +from loguru import logger +from omnidreams.interactive_drive.types import PresentedFrame, SceneBundle +from PIL import Image + +RECORDING_HOTKEY_COLLISIONS = frozenset({"r", "x", "1", "2"}) + + +@dataclass(frozen=True) +class RecordingConfig: + enabled: bool = False + output_dir: Path | None = None + hotkey: str = "f9" + auto_start: bool = False + max_buffer_frames: int = 600 + + +@dataclass +class _RecordingSession: + output_dir: Path + start_time_utc: str + hdmap_frames: list[np.ndarray] + inferred_frames: list[np.ndarray] + dropped_hdmap_frames: int = 0 + dropped_inferred_frames: int = 0 + frame_drop_warning_emitted: bool = False + + +def normalize_recording_hotkey(raw: object) -> str: + """Normalize manifest/browser key names into one small vocabulary.""" + value = str(raw).strip() + if not value: + raise ValueError("recording_hotkey must be a non-empty key name") + aliases = { + " ": "space", + "spacebar": "space", + "arrowup": "up", + "arrow_up": "up", + "arrowdown": "down", + "arrow_down": "down", + "arrowleft": "left", + "arrow_left": "left", + "arrowright": "right", + "arrow_right": "right", + } + lowered = value.lower() + return aliases.get(lowered, lowered) + + +def recording_hotkey_matches(configured: str, key: object) -> bool: + """Return whether ``key`` names the configured recording hotkey.""" + try: + return normalize_recording_hotkey(key) == normalize_recording_hotkey(configured) + except ValueError: + return False + + +def recording_hotkey_collides_with_controls(hotkey: str) -> bool: + return normalize_recording_hotkey(hotkey) in RECORDING_HOTKEY_COLLISIONS + + +def slangpy_key_name_candidates(hotkey: str) -> tuple[str, ...]: + """Map a normalized hotkey to likely SlangPy ``KeyCode`` attribute names.""" + key = normalize_recording_hotkey(hotkey) + if len(key) == 1 and key.isdigit(): + return (f"key{key}", f"digit{key}", f"num_{key}") + return (key,) + + +class InteractiveDriveRecorder: + """Collect a rollout recording and write its artifact bundle on stop.""" + + def __init__( + self, + config: RecordingConfig, + *, + scene: SceneBundle, + fps: int, + ) -> None: + if not config.enabled: + raise ValueError("InteractiveDriveRecorder requires recording.enabled") + if config.output_dir is None: + raise ValueError("InteractiveDriveRecorder requires an output directory") + if config.max_buffer_frames <= 0: + raise ValueError("RecordingConfig.max_buffer_frames must be positive") + self._config = config + self._scene = scene + self._fps = int(fps) + self._active_session: _RecordingSession | None = None + self._session_count = 0 + self._closed_paths: list[Path] = [] + + @property + def is_recording(self) -> bool: + return self._active_session is not None + + @property + def closed_paths(self) -> tuple[Path, ...]: + return tuple(self._closed_paths) + + @property + def auto_start(self) -> bool: + return self._config.auto_start + + def toggle(self) -> None: + if self.is_recording: + self.stop(reason="hotkey") + return + self.start() + + def start(self) -> None: + if self.is_recording: + return + self._session_count += 1 + start_time = datetime.now(timezone.utc) + output_dir = self._next_output_dir(start_time) + try: + output_dir.mkdir(parents=True, exist_ok=False) + (output_dir / "prompt.txt").write_text( + self._scene.prompt, + encoding="utf-8", + ) + except Exception as exc: + logger.warning( + "[recording] failed to start dir={} error={!r}", + output_dir, + exc, + ) + return + self._active_session = _RecordingSession( + output_dir=output_dir, + start_time_utc=start_time.isoformat(), + hdmap_frames=[], + inferred_frames=[], + ) + print(f"[recording] started dir={output_dir}", flush=True) + + def stop(self, *, reason: str) -> Path | None: + session = self._active_session + if session is None: + return None + self._active_session = None + metadata = { + "scene_id": self._scene.scene_id, + "scene_path": str(self._scene.scene_path), + "fps": self._fps, + "start_time_utc": session.start_time_utc, + "stop_time_utc": datetime.now(timezone.utc).isoformat(), + "reason": reason, + "hdmap_frames": len(session.hdmap_frames), + "inferred_frames": len(session.inferred_frames), + "dropped_hdmap_frames": session.dropped_hdmap_frames, + "dropped_inferred_frames": session.dropped_inferred_frames, + "max_buffer_frames": self._config.max_buffer_frames, + } + try: + (session.output_dir / "metadata.json").write_text( + json.dumps(metadata, indent=2, sort_keys=True) + "\n", + encoding="utf-8", + ) + first_frame = ( + session.inferred_frames[0] if session.inferred_frames else None + ) + if first_frame is None and session.hdmap_frames: + first_frame = session.hdmap_frames[0] + if first_frame is not None: + Image.fromarray(first_frame).save( + session.output_dir / "first_frame.png", + ) + _write_video( + session.hdmap_frames, + session.output_dir / "hdmap.mp4", + self._fps, + ) + _write_video( + session.inferred_frames, + session.output_dir / "inferred.mp4", + self._fps, + ) + except Exception as exc: + logger.warning( + "[recording] failed to save dir={} reason={} error={!r}", + session.output_dir, + reason, + exc, + ) + return None + self._closed_paths.append(session.output_dir) + print( + "[recording] saved " + f"dir={session.output_dir} " + f"hdmap_frames={len(session.hdmap_frames)} " + f"inferred_frames={len(session.inferred_frames)}", + flush=True, + ) + return session.output_dir + + def record_frame(self, frame: PresentedFrame) -> None: + session = self._active_session + if session is None: + return + self._append_frame( + session, + buffer=session.hdmap_frames, + frame=_as_rgb_host_uint8(frame.rgb_host_uint8), + stream="hdmap", + ) + if frame.model_rgb_host_uint8 is not None: + inferred_frame = _as_rgb_host_uint8(frame.model_rgb_host_uint8) + self._append_frame( + session, + buffer=session.inferred_frames, + frame=inferred_frame, + stream="inferred", + ) + + def close(self, *, reason: str) -> None: + if self.is_recording: + self.stop(reason=reason) + + def _next_output_dir(self, start_time: datetime) -> Path: + root = self._config.output_dir + if root is None: + raise ValueError("Recording output_dir is required") + timestamp = start_time.strftime("%Y%m%d-%H%M%S") + scene_slug = _slugify(self._scene.scene_id or self._scene.scene_path.stem) + base = root / f"{timestamp}-{scene_slug}-{self._session_count:03d}" + candidate = base + suffix = 0 + while candidate.exists(): + suffix += 1 + candidate = root / f"{base.name}-{suffix:02d}" + return candidate + + def _append_frame( + self, + session: _RecordingSession, + *, + buffer: list[np.ndarray], + frame: np.ndarray, + stream: str, + ) -> None: + if len(buffer) >= self._config.max_buffer_frames: + buffer.pop(0) + if stream == "hdmap": + session.dropped_hdmap_frames += 1 + elif stream == "inferred": + session.dropped_inferred_frames += 1 + if not session.frame_drop_warning_emitted: + logger.warning( + "[recording] frame buffer reached max_buffer_frames={}; " + "dropping oldest frames for dir={}", + self._config.max_buffer_frames, + session.output_dir, + ) + session.frame_drop_warning_emitted = True + buffer.append(frame) + + +def _as_rgb_host_uint8(value: Any) -> np.ndarray: + if hasattr(value, "to_numpy"): + array = value.to_numpy() + else: + array = np.asarray(value) + if array.ndim != 3 or array.shape[2] < 3: + raise ValueError(f"Expected an HxWx3 RGB frame, got shape {array.shape!r}") + rgb = array[:, :, :3] + if rgb.dtype != np.uint8: + rgb = np.clip(rgb, 0, 255).astype(np.uint8) + return np.ascontiguousarray(rgb.copy()) + + +def _write_video(frames: list[np.ndarray], path: Path, fps: int) -> None: + if not frames: + return + try: + import mediapy as media # noqa: PLC0415 + except ImportError as exc: # pragma: no cover - import-time gate + raise ImportError( + "Writing interactive-drive recordings needs mediapy. " + "Install the flashdreams-omnidreams package dependencies." + ) from exc + + media.write_video(str(path), np.stack(frames, axis=0), fps=fps) + + +def _slugify(value: str) -> str: + slug = re.sub(r"[^a-zA-Z0-9_.-]+", "-", value.strip()).strip("-._") + return slug or "scene" diff --git a/integrations/omnidreams/omnidreams/interactive_drive/runtime/loop.py b/integrations/omnidreams/omnidreams/interactive_drive/runtime/loop.py index 1696bc22d..2e9cc7293 100644 --- a/integrations/omnidreams/omnidreams/interactive_drive/runtime/loop.py +++ b/integrations/omnidreams/omnidreams/interactive_drive/runtime/loop.py @@ -127,6 +127,19 @@ def present_frame(self, frame: PresentedFrame, view_mode: str) -> None: ... def close(self) -> None: ... +class RecordingBackend(Protocol): + @property + def auto_start(self) -> bool: ... + + def start(self) -> None: ... + + def toggle(self) -> None: ... + + def record_frame(self, frame: PresentedFrame) -> None: ... + + def close(self, *, reason: str) -> None: ... + + class MainLoopState: """Mutable per-iteration counters and timestamps for :func:`run_main_loop`. @@ -397,6 +410,22 @@ def _drain_pipeline_frames( ready_frames.append(queued_frame) +def _consume_recording_toggle_request(runtime_controls: RuntimeControls) -> bool: + consume = getattr(runtime_controls, "consume_recording_toggle_request", None) + if not callable(consume): + return False + return bool(consume()) + + +def _close_recording( + recorder: RecordingBackend | None, + *, + reason: str, +) -> None: + if recorder is not None: + recorder.close(reason=reason) + + def run_main_loop( presenter: PresenterBackend, runtime_controls: RuntimeControls, @@ -406,6 +435,7 @@ def run_main_loop( pipeline: ChunkPipeline, config: LoopConfig, loading_status: Callable[[], str | None] | None = None, + recorder: RecordingBackend | None = None, ) -> bool: """Drive the request -> render -> present pipeline. @@ -426,13 +456,18 @@ def run_main_loop( chunk_history = ChunkHistory(config.history_capacity) if _profile_input_to_present_enabled(): reset_input_to_present_profile_window() + if recorder is not None and recorder.auto_start and not presenter.should_close: + recorder.start() while not presenter.should_close: presenter.process_events() if presenter.should_close: break if runtime_controls.consume_reset_request(): + _close_recording(recorder, reason="reset") return True + if _consume_recording_toggle_request(runtime_controls) and recorder is not None: + recorder.toggle() sampled = input_backend.sample() # Keep one chunk in flight. @@ -450,6 +485,7 @@ def run_main_loop( # OOB overlay from the new boundary frame and auto-respawn (same # ``return True`` as a manual reset) when far enough off-map. if update_oob_state(state, simulation, config): + _close_recording(recorder, reason="respawn") return True # Republish telemetry per chunk so read-side observers (e.g. the # presenter's ``/state`` endpoint) see the latest state. @@ -481,6 +517,8 @@ def run_main_loop( view_mode=view_mode, oob_message=state.oob_message, ) + if recorder is not None: + recorder.record_frame(queued_frame.frame) last_presented_frame = queued_frame.frame state.frame_count += 1 else: @@ -500,4 +538,5 @@ def run_main_loop( ) state.next_present_time += config.frame_interval_s + _close_recording(recorder, reason="loop-end") return False diff --git a/integrations/omnidreams/omnidreams/interactive_drive/slangpy_hud_presenter.py b/integrations/omnidreams/omnidreams/interactive_drive/slangpy_hud_presenter.py index 89f720ad0..8161cc37a 100644 --- a/integrations/omnidreams/omnidreams/interactive_drive/slangpy_hud_presenter.py +++ b/integrations/omnidreams/omnidreams/interactive_drive/slangpy_hud_presenter.py @@ -28,6 +28,7 @@ _CudaRGBInterop, _env_truthy, ) +from omnidreams.interactive_drive.recording import slangpy_key_name_candidates from omnidreams.interactive_drive.types import DriverCommand, PresentedFrame from PIL import Image, ImageDraw, ImageFont @@ -2083,6 +2084,8 @@ def _on_keyboard_event(self, event: Any) -> None: self._keyboard.request_reset() elif self._key_matches(key, "x"): self.exit_scene() + elif self._recording_hotkey_matches(key): + self._keyboard.request_recording_toggle() def _expire_pending_drive_releases(self) -> None: """Commit any debounced release whose grace window has passed. @@ -2134,6 +2137,15 @@ def _key_matches(self, event_key: Any, name: str) -> bool: code = self._key_codes.get(name) return code is not None and event_key == code + def _recording_hotkey_matches(self, event_key: Any) -> bool: + if not self._keyboard.recording_enabled: + return False + key_enum = self._spy.KeyCode + for name in slangpy_key_name_candidates(self._keyboard.recording_hotkey): + if event_key == _lookup_key(key_enum, name): + return True + return False + def _on_mouse_event(self, event: Any) -> None: spy = self._spy # ``pos`` is float2 in window-relative pixels. We round to int diff --git a/integrations/omnidreams/omnidreams/interactive_drive/streaming_presenter.py b/integrations/omnidreams/omnidreams/interactive_drive/streaming_presenter.py index 8de0cd46e..4d76d66e4 100644 --- a/integrations/omnidreams/omnidreams/interactive_drive/streaming_presenter.py +++ b/integrations/omnidreams/omnidreams/interactive_drive/streaming_presenter.py @@ -27,6 +27,7 @@ from omnidreams.interactive_drive.config import RasterConfig from omnidreams.interactive_drive.input.keyboard import KeyboardState from omnidreams.interactive_drive.loading_overlay import render_loading_overlay +from omnidreams.interactive_drive.recording import recording_hotkey_matches from omnidreams.interactive_drive.types import DriverCommand, PresentedFrame from PIL import Image @@ -823,6 +824,11 @@ def _apply_control(self, key: str, down: bool) -> None: # holding the key doesn't trigger a cascade of resets. if key in ("r", "R"): self._keyboard.request_reset() + return + if self._keyboard.recording_enabled and recording_hotkey_matches( + self._keyboard.recording_hotkey, key + ): + self._keyboard.request_recording_toggle() def _state_snapshot(self) -> dict[str, float | None]: """JSON-serialisable telemetry snapshot from ``KeyboardState.vehicle_state``. diff --git a/integrations/omnidreams/omnidreams/interactive_drive/world_model/manifest.py b/integrations/omnidreams/omnidreams/interactive_drive/world_model/manifest.py index abc60280d..8eb6b1a19 100644 --- a/integrations/omnidreams/omnidreams/interactive_drive/world_model/manifest.py +++ b/integrations/omnidreams/omnidreams/interactive_drive/world_model/manifest.py @@ -16,6 +16,11 @@ resolve_hf_org, rewrite_omni_dreams_urls, ) +from omnidreams.interactive_drive.recording import ( + RECORDING_HOTKEY_COLLISIONS, + normalize_recording_hotkey, + recording_hotkey_collides_with_controls, +) _HF_URL_PATTERN = re.compile( r"^https?://(?:www\.)?huggingface\.co/[^/]+/[^/]+/(?:blob|resolve)/[^/]+/.+$", @@ -143,6 +148,54 @@ def _parse_native_vae_encoder(raw: object) -> str: return encoder +def _parse_recording_dir(raw: object) -> Path | None: + if raw is None: + return None + value = str(raw).strip() + if not value: + return None + return Path(value).expanduser() + + +def _parse_recording_hotkey(raw: object, *, enabled: bool) -> str: + hotkey = normalize_recording_hotkey("f9" if raw is None else raw) + if enabled and recording_hotkey_collides_with_controls(hotkey): + logger.warning( + "[manifest] recording_hotkey={!r} overlaps an interactive-drive " + "control key {}; choose a non-control key such as F9 to avoid " + "shadowing reset/view/exit controls.", + hotkey, + sorted(RECORDING_HOTKEY_COLLISIONS), + ) + return hotkey + + +@dataclass(frozen=True) +class RecordingManifest: + enabled: bool = False + dir: Path | None = None + hotkey: str = "f9" + auto_start: bool = False + + +def parse_recording_manifest(data: dict[str, object]) -> RecordingManifest: + enabled = bool(data.get("recording_enabled", False)) + return RecordingManifest( + enabled=enabled, + dir=_parse_recording_dir(data.get("recording_dir")), + hotkey=_parse_recording_hotkey( + data.get("recording_hotkey"), + enabled=enabled, + ), + auto_start=bool(data.get("recording_auto_start", False)), + ) + + +def load_recording_manifest(path: str | Path) -> RecordingManifest: + data = yaml.safe_load(Path(path).read_text(encoding="utf-8")) or {} + return parse_recording_manifest(data) + + @dataclass(frozen=True) class WorldModelManifest: debug_condition_frame_dir: Path | None = None @@ -171,6 +224,10 @@ class WorldModelManifest: native_dit_sparge_hybrid_phase: int | None = None native_vae_encoder: str = "disabled" native_vae_fp8_state_path: Path | None = None + recording_enabled: bool = False + recording_dir: Path | None = None + recording_hotkey: str = "f9" + recording_auto_start: bool = False def load_world_model_manifest(path: str | Path) -> WorldModelManifest: @@ -191,6 +248,7 @@ def load_world_model_manifest(path: str | Path) -> WorldModelManifest: raw_yaml = rewritten data = yaml.safe_load(raw_yaml) or {} resolution = _parse_resolution_wh(data.get("resolution_wh")) + recording = parse_recording_manifest(data) return WorldModelManifest( debug_condition_frame_dir=_resolve_manifest_path( data.get("debug_condition_frame_dir"), @@ -248,4 +306,8 @@ def load_world_model_manifest(path: str | Path) -> WorldModelManifest: data.get("native_vae_fp8_state_path"), manifest_dir=manifest_dir, ), + recording_enabled=recording.enabled, + recording_dir=recording.dir, + recording_hotkey=recording.hotkey, + recording_auto_start=recording.auto_start, ) diff --git a/integrations/omnidreams/tests/interactive_drive/test_cli_manifest_resolution.py b/integrations/omnidreams/tests/interactive_drive/test_cli_manifest_resolution.py index 6bddcfbc0..e4561e75d 100644 --- a/integrations/omnidreams/tests/interactive_drive/test_cli_manifest_resolution.py +++ b/integrations/omnidreams/tests/interactive_drive/test_cli_manifest_resolution.py @@ -5,8 +5,10 @@ import os import tempfile +import textwrap import unittest from pathlib import Path +from unittest.mock import patch from omnidreams.interactive_drive import cli @@ -33,6 +35,87 @@ def test_cwd_relative_manifest_wins(self) -> None: self.assertEqual(resolved, manifest.resolve()) + def test_relative_recording_dir_resolves_from_flashdreams_root(self) -> None: + resolved = cli._resolve_recording_output_dir(Path("captures"), enabled=True) + + self.assertEqual(resolved, (cli._FLASHDREAMS_ROOT / "captures").resolve()) + + def test_absolute_recording_dir_is_preserved(self) -> None: + absolute = Path(tempfile.gettempdir()) / "interactive-drive-captures" + + resolved = cli._resolve_recording_output_dir(absolute, enabled=True) + + self.assertEqual(resolved, absolute) + + def test_default_recording_dir_is_flashdreams_root_recordings(self) -> None: + resolved = cli._resolve_recording_output_dir(None, enabled=True) + + self.assertEqual( + resolved, + (cli._FLASHDREAMS_ROOT / "recordings").resolve(), + ) + + def test_disabled_recording_has_no_output_dir(self) -> None: + resolved = cli._resolve_recording_output_dir(Path("captures"), enabled=False) + + self.assertIsNone(resolved) + + def test_flashdreams_root_fallback_warns(self) -> None: + with tempfile.TemporaryDirectory() as tmpdir: + start = Path(tmpdir) / "installed" / "omnidreams" + start.mkdir(parents=True) + old_cwd = Path.cwd() + try: + os.chdir(tmpdir) + with self.assertWarnsRegex( + RuntimeWarning, + "Could not locate the flashdreams repository root", + ): + resolved = cli._find_flashdreams_root(start) + finally: + os.chdir(old_cwd) + + self.assertEqual(resolved, Path(tmpdir).resolve()) + + def test_raster_backend_loads_recording_fields_from_optional_manifest( + self, + ) -> None: + class FakeRasterRenderBackend: + def __init__(self, *, chunk, raster, bev) -> None: + del chunk, raster, bev + + with tempfile.TemporaryDirectory() as tmpdir: + manifest = Path(tmpdir) / "manifest.yaml" + manifest.write_text( + textwrap.dedent( + """ + recording_enabled: true + recording_dir: raster-captures + recording_hotkey: F8 + recording_auto_start: true + """ + ).strip(), + encoding="utf-8", + ) + args = cli.build_parser().parse_args( + [ + "--backend", + "raster", + "--manifest", + str(manifest), + ] + ) + with patch.object(cli, "RasterRenderBackend", FakeRasterRenderBackend): + config, _backend = cli.prepare_config_and_backend(args) + + self.assertTrue(config.recording.enabled) + self.assertEqual( + config.recording.output_dir, + (cli._FLASHDREAMS_ROOT / "raster-captures").resolve(), + ) + self.assertEqual(config.recording.hotkey, "f8") + self.assertTrue(config.recording.auto_start) + if __name__ == "__main__": unittest.main() diff --git a/integrations/omnidreams/tests/interactive_drive/test_keyboard_state.py b/integrations/omnidreams/tests/interactive_drive/test_keyboard_state.py index eb08877a4..1363da9e6 100644 --- a/integrations/omnidreams/tests/interactive_drive/test_keyboard_state.py +++ b/integrations/omnidreams/tests/interactive_drive/test_keyboard_state.py @@ -62,3 +62,19 @@ def test_consume_exit_scene_request_returns_true_once_per_request() -> None: keyboard.request_exit_scene() assert keyboard.consume_exit_scene_request() is True assert keyboard.consume_exit_scene_request() is False + + +def test_recording_toggle_request_ignored_when_recording_disabled() -> None: + keyboard = KeyboardState() + keyboard.request_recording_toggle() + assert keyboard.consume_recording_toggle_request() is False + + +def test_recording_toggle_request_returns_true_once_when_enabled() -> None: + keyboard = KeyboardState(recording_enabled=True, recording_hotkey="F9") + assert keyboard.recording_enabled is True + assert keyboard.recording_hotkey == "f9" + keyboard.request_recording_toggle() + keyboard.request_recording_toggle() + assert keyboard.consume_recording_toggle_request() is True + assert keyboard.consume_recording_toggle_request() is False diff --git a/integrations/omnidreams/tests/interactive_drive/test_latency_loop.py b/integrations/omnidreams/tests/interactive_drive/test_latency_loop.py index 2cd1718fe..1abc4f8f4 100644 --- a/integrations/omnidreams/tests/interactive_drive/test_latency_loop.py +++ b/integrations/omnidreams/tests/interactive_drive/test_latency_loop.py @@ -149,8 +149,14 @@ def present_frame(self, frame: PresentedFrame, view_mode: str) -> None: class _FakeRuntimeControls: - def __init__(self, *, reset_after_present: int | None = None) -> None: + def __init__( + self, + *, + reset_after_present: int | None = None, + recording_toggle_after_present: int | None = None, + ) -> None: self._reset_after_present = reset_after_present + self._recording_toggle_after_present = recording_toggle_after_present self._presenter: _CountingPresenter | None = None self.view_mode = "rgb" @@ -165,6 +171,44 @@ def consume_reset_request(self) -> bool: return True return False + def consume_recording_toggle_request(self) -> bool: + if self._recording_toggle_after_present is None or self._presenter is None: + return False + if len(self._presenter.records) >= self._recording_toggle_after_present: + self._recording_toggle_after_present = None + return True + return False + + +class _RecorderProbe: + def __init__(self, *, auto_start: bool = False) -> None: + self.auto_start = auto_start + self.start_calls = 0 + self.toggle_calls = 0 + self.recorded_frames: list[PresentedFrame] = [] + self.close_reasons: list[str] = [] + self._active = False + + def start(self) -> None: + self.start_calls += 1 + self._active = True + + def toggle(self) -> None: + self.toggle_calls += 1 + if self._active: + self.close(reason="hotkey") + return + self.start() + + def record_frame(self, frame: PresentedFrame) -> None: + if self._active: + self.recorded_frames.append(frame) + + def close(self, *, reason: str) -> None: + if self._active: + self.close_reasons.append(reason) + self._active = False + class _FakeInputBackend: def sample(self) -> SampledInput: @@ -202,6 +246,7 @@ def _drive_loop( simulation: _FakeSimulation, initial: PresentedFrame, frame_interval_s: float, + recorder: _RecorderProbe | None = None, ) -> bool: pipeline = ChunkPipeline(backend) pipeline.request_scene(minimal_scene()) @@ -214,6 +259,7 @@ def _drive_loop( simulation=simulation, pipeline=pipeline, config=_loop_config(frame_interval_s=frame_interval_s), + recorder=recorder, ) finally: pipeline.shutdown() @@ -404,6 +450,55 @@ def test_loop_presents_backend_frames_when_available() -> None: assert any(record.frame is not initial for record in presenter.records) +def test_loop_auto_starts_recorder_and_records_backend_frames() -> None: + initial = _make_frame() + presenter = _CountingPresenter( + present_budget=_backend_frame_wait_budget(), close_on_frame=initial + ) + controls = _FakeRuntimeControls() + recorder = _RecorderProbe(auto_start=True) + + result = _drive_loop( + presenter=presenter, + controls=controls, + backend=FakeVideoModelBackend(frames_per_render=1, rgb_value=7), + simulation=_FakeSimulation(), + initial=initial, + frame_interval_s=0.001, + recorder=recorder, + ) + + assert result is False + assert recorder.start_calls == 1 + assert recorder.recorded_frames + assert all(frame is not initial for frame in recorder.recorded_frames) + assert recorder.close_reasons == ["loop-end"] + + +def test_loop_consumes_recording_hotkey_toggle() -> None: + initial = _make_frame() + presenter = _CountingPresenter( + present_budget=_backend_frame_wait_budget(), close_on_frame=initial + ) + controls = _FakeRuntimeControls(recording_toggle_after_present=0) + controls.bind_presenter(presenter) + recorder = _RecorderProbe() + + _drive_loop( + presenter=presenter, + controls=controls, + backend=FakeVideoModelBackend(frames_per_render=1, rgb_value=7), + simulation=_FakeSimulation(), + initial=initial, + frame_interval_s=0.001, + recorder=recorder, + ) + + assert recorder.toggle_calls == 1 + assert recorder.start_calls == 1 + assert recorder.recorded_frames + + def test_loop_prepares_backend_frames_before_presenting_them() -> None: initial = _make_frame() presenter = _PreparingPresenter( diff --git a/integrations/omnidreams/tests/interactive_drive/test_manifest.py b/integrations/omnidreams/tests/interactive_drive/test_manifest.py index b268ad31e..6dc9be9c7 100644 --- a/integrations/omnidreams/tests/interactive_drive/test_manifest.py +++ b/integrations/omnidreams/tests/interactive_drive/test_manifest.py @@ -7,6 +7,7 @@ import textwrap import unittest from pathlib import Path +from unittest.mock import patch from omnidreams.interactive_drive.world_model.manifest import load_world_model_manifest @@ -30,6 +31,7 @@ def test_loads_defaults(self) -> None: self.assertEqual(manifest.native_dit_acceleration, "disabled") self.assertEqual(manifest.native_vae_encoder, "disabled") self.assertIsNone(manifest.native_vae_fp8_state_path) + self.assertEqual(manifest.recording_hotkey, "f9") def test_loads_native_dit_knobs(self) -> None: with tempfile.TemporaryDirectory() as tmpdir: @@ -83,6 +85,46 @@ def test_loads_native_vae_knobs(self) -> None: (root / "native/lightvae-fp8-state.pt").resolve(), ) + def test_loads_recording_knobs(self) -> None: + with tempfile.TemporaryDirectory() as tmpdir: + path = Path(tmpdir) / "manifest.yaml" + path.write_text( + textwrap.dedent( + """ + recording_enabled: true + recording_dir: captures + recording_hotkey: F9 + recording_auto_start: true + """ + ).strip(), + encoding="utf-8", + ) + manifest = load_world_model_manifest(path) + self.assertTrue(manifest.recording_enabled) + self.assertEqual(manifest.recording_dir, Path("captures")) + self.assertEqual(manifest.recording_hotkey, "f9") + self.assertTrue(manifest.recording_auto_start) + + def test_warns_when_recording_hotkey_collides_with_controls(self) -> None: + with tempfile.TemporaryDirectory() as tmpdir: + path = Path(tmpdir) / "manifest.yaml" + path.write_text( + textwrap.dedent( + """ + recording_enabled: true + recording_hotkey: r + """ + ).strip(), + encoding="utf-8", + ) + with patch( + "omnidreams.interactive_drive.world_model.manifest.logger.warning" + ) as warning: + manifest = load_world_model_manifest(path) + + self.assertEqual(manifest.recording_hotkey, "r") + warning.assert_called_once() + def test_rejects_unaligned_resolution(self) -> None: with tempfile.TemporaryDirectory() as tmpdir: path = Path(tmpdir) / "manifest.yaml" diff --git a/integrations/omnidreams/tests/interactive_drive/test_recording.py b/integrations/omnidreams/tests/interactive_drive/test_recording.py new file mode 100644 index 000000000..c53ce9666 --- /dev/null +++ b/integrations/omnidreams/tests/interactive_drive/test_recording.py @@ -0,0 +1,240 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + +import json +from dataclasses import replace +from datetime import datetime +from pathlib import Path + +import numpy as np +import omnidreams.interactive_drive.recording as recording_module +from omnidreams.interactive_drive._pipeline_fakes import minimal_scene +from omnidreams.interactive_drive.recording import ( + InteractiveDriveRecorder, + RecordingConfig, + normalize_recording_hotkey, + recording_hotkey_collides_with_controls, + recording_hotkey_matches, + slangpy_key_name_candidates, +) +from omnidreams.interactive_drive.types import PresentedFrame +from PIL import Image + + +def test_hotkey_normalization_and_matching() -> None: + assert normalize_recording_hotkey("R") == "r" + assert normalize_recording_hotkey("ArrowUp") == "up" + assert recording_hotkey_matches("f9", "F9") is True + assert recording_hotkey_collides_with_controls("R") is True + assert recording_hotkey_collides_with_controls("F9") is False + assert slangpy_key_name_candidates("1") == ("key1", "digit1", "num_1") + + +def test_recorder_writes_bundle_on_stop(tmp_path, monkeypatch) -> None: + writes: list[tuple[Path, int, tuple[int, ...]]] = [] + + def fake_write_video(frames: list[np.ndarray], path: Path, fps: int) -> None: + writes.append((path, fps, np.stack(frames, axis=0).shape)) + path.write_bytes(b"fake mp4") + + monkeypatch.setattr(recording_module, "_write_video", fake_write_video) + scene = minimal_scene() + scene = replace( + scene, + scene_id="recording scene", + initial_rgb=np.full((4, 4, 3), 3, dtype=np.uint8), + prompt="drive through a bright test scene", + ) + recorder = InteractiveDriveRecorder( + RecordingConfig(enabled=True, output_dir=tmp_path, hotkey="r"), + scene=scene, + fps=30, + ) + + recorder.start() + recorder.record_frame( + PresentedFrame( + timestamp_us=0, + rgb_host_uint8=np.full((4, 4, 3), 7, dtype=np.uint8), + depth_host_f32=None, + model_rgb_host_uint8=np.full((4, 4, 3), 11, dtype=np.uint8), + ) + ) + output_dir = recorder.stop(reason="hotkey") + + assert output_dir is not None + assert (output_dir / "first_frame.png").exists() + assert np.array_equal( + np.asarray(Image.open(output_dir / "first_frame.png")), + np.full((4, 4, 3), 11, dtype=np.uint8), + ) + assert (output_dir / "prompt.txt").read_text(encoding="utf-8") == scene.prompt + assert (output_dir / "metadata.json").exists() + assert (output_dir / "hdmap.mp4").read_bytes() == b"fake mp4" + assert (output_dir / "inferred.mp4").read_bytes() == b"fake mp4" + assert writes == [ + (output_dir / "hdmap.mp4", 30, (1, 4, 4, 3)), + (output_dir / "inferred.mp4", 30, (1, 4, 4, 3)), + ] + + +def test_recorder_video_write_failure_is_nonfatal(tmp_path, monkeypatch) -> None: + def fail_write_video( + frames: list[np.ndarray], + path: Path, + fps: int, + ) -> None: + raise OSError("disk full") + + monkeypatch.setattr(recording_module, "_write_video", fail_write_video) + scene = replace(minimal_scene(), scene_id="recording scene") + recorder = InteractiveDriveRecorder( + RecordingConfig(enabled=True, output_dir=tmp_path, hotkey="F9"), + scene=scene, + fps=30, + ) + + recorder.start() + recorder.record_frame( + PresentedFrame( + timestamp_us=0, + rgb_host_uint8=np.full((4, 4, 3), 7, dtype=np.uint8), + depth_host_f32=None, + model_rgb_host_uint8=np.full((4, 4, 3), 11, dtype=np.uint8), + ) + ) + + assert recorder.stop(reason="hotkey") is None + assert recorder.is_recording is False + assert recorder.closed_paths == () + + +def test_recorder_start_failure_is_nonfatal(tmp_path) -> None: + output_root = tmp_path / "not-a-directory" + output_root.write_text("plain file", encoding="utf-8") + recorder = InteractiveDriveRecorder( + RecordingConfig(enabled=True, output_dir=output_root, hotkey="F9"), + scene=minimal_scene(), + fps=30, + ) + + recorder.start() + + assert recorder.is_recording is False + assert recorder.closed_paths == () + + +def test_recorder_caps_frame_buffers(tmp_path, monkeypatch) -> None: + writes: list[tuple[Path, list[int]]] = [] + + def fake_write_video(frames: list[np.ndarray], path: Path, fps: int) -> None: + del fps + writes.append((path, [int(frame[0, 0, 0]) for frame in frames])) + path.write_bytes(b"fake mp4") + + monkeypatch.setattr(recording_module, "_write_video", fake_write_video) + scene = replace(minimal_scene(), scene_id="recording scene") + recorder = InteractiveDriveRecorder( + RecordingConfig( + enabled=True, + output_dir=tmp_path, + hotkey="F9", + max_buffer_frames=2, + ), + scene=scene, + fps=30, + ) + + recorder.start() + for value in (1, 2, 3): + recorder.record_frame( + PresentedFrame( + timestamp_us=value, + rgb_host_uint8=np.full((4, 4, 3), value, dtype=np.uint8), + depth_host_f32=None, + model_rgb_host_uint8=np.full((4, 4, 3), value + 10, dtype=np.uint8), + ) + ) + output_dir = recorder.stop(reason="hotkey") + + assert output_dir is not None + assert writes == [ + (output_dir / "hdmap.mp4", [2, 3]), + (output_dir / "inferred.mp4", [12, 13]), + ] + assert np.array_equal( + np.asarray(Image.open(output_dir / "first_frame.png")), + np.full((4, 4, 3), 12, dtype=np.uint8), + ) + metadata = json.loads((output_dir / "metadata.json").read_text(encoding="utf-8")) + assert metadata["hdmap_frames"] == 2 + assert metadata["inferred_frames"] == 2 + assert metadata["dropped_hdmap_frames"] == 1 + assert metadata["dropped_inferred_frames"] == 1 + assert metadata["max_buffer_frames"] == 2 + + +def test_recorder_saves_first_frame_from_hdmap_when_no_inferred_frames( + tmp_path, + monkeypatch, +) -> None: + writes: list[tuple[Path, int]] = [] + + def fake_write_video(frames: list[np.ndarray], path: Path, fps: int) -> None: + del fps + writes.append((path, len(frames))) + if frames: + path.write_bytes(b"fake mp4") + + monkeypatch.setattr(recording_module, "_write_video", fake_write_video) + scene = replace(minimal_scene(), scene_id="raster scene") + recorder = InteractiveDriveRecorder( + RecordingConfig(enabled=True, output_dir=tmp_path, hotkey="F9"), + scene=scene, + fps=30, + ) + + recorder.start() + recorder.record_frame( + PresentedFrame( + timestamp_us=0, + rgb_host_uint8=np.full((4, 4, 3), 17, dtype=np.uint8), + depth_host_f32=None, + model_rgb_host_uint8=None, + ) + ) + output_dir = recorder.stop(reason="hotkey") + + assert output_dir is not None + assert np.array_equal( + np.asarray(Image.open(output_dir / "first_frame.png")), + np.full((4, 4, 3), 17, dtype=np.uint8), + ) + assert (output_dir / "hdmap.mp4").read_bytes() == b"fake mp4" + assert not (output_dir / "inferred.mp4").exists() + assert writes == [ + (output_dir / "hdmap.mp4", 1), + (output_dir / "inferred.mp4", 0), + ] + + +def test_recorder_collision_suffix_starts_at_one(tmp_path, monkeypatch) -> None: + class FixedDateTime: + @classmethod + def now(cls, tz): + return datetime(2026, 6, 10, 1, 2, 3, tzinfo=tz) + + monkeypatch.setattr(recording_module, "datetime", FixedDateTime) + scene = replace(minimal_scene(), scene_id="recording scene") + recorder = InteractiveDriveRecorder( + RecordingConfig(enabled=True, output_dir=tmp_path, hotkey="F9"), + scene=scene, + fps=30, + ) + base_name = "20260610-010203-recording-scene-001" + (tmp_path / base_name).mkdir() + + recorder.start() + + assert recorder.is_recording is True + assert (tmp_path / f"{base_name}-01").is_dir()