kos_sim.video_recorder

  1"""Video recording functionality for the KOS simulator."""
  2
  3import asyncio
  4import datetime
  5import os
  6import time
  7from pathlib import Path
  8
  9import mediapy as mp
 10import numpy as np
 11
 12from kos_sim import logger
 13from kos_sim.simulator import MujocoSimulator
 14
 15
 16class VideoRecorder:
 17    """Manages video recording for a MuJoCo simulator."""
 18
 19    def __init__(
 20        self,
 21        simulator: MujocoSimulator,
 22        output_dir: str | Path,
 23        fps: int = 30,
 24        frame_width: int = 640,
 25        frame_height: int = 480,
 26    ) -> None:
 27        self.simulator = simulator
 28        self.output_dir = Path(output_dir)
 29        self.fps = fps
 30        self.frame_width = frame_width
 31        self.frame_height = frame_height
 32
 33        self.last_frame_time = 0.0
 34        self.writer = None
 35        self.is_recording = False
 36        self.recording_task: asyncio.Task[None] | None = None
 37        self.frames: list[np.ndarray] = []
 38
 39        os.makedirs(self.output_dir, exist_ok=True)
 40
 41    def start_recording(self, custom_filename: str | None = None) -> str:
 42        """Start recording video.
 43
 44        Args:
 45            custom_filename: Optional custom filename (without extension)
 46
 47        Returns:
 48            Path to the output video file
 49        """
 50        if self.is_recording:
 51            self.stop_recording()
 52
 53        # Generate filename based on datetime if not provided
 54        if custom_filename is None:
 55            timestamp = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
 56            filename = f"sim_recording_{timestamp}.mp4"
 57        else:
 58            filename = f"{custom_filename}.mp4"
 59
 60        self.filepath = self.output_dir / filename
 61
 62        self.frames = []
 63        self.last_frame_time = time.time()
 64        self.is_recording = True
 65
 66        logger.info(f"Started video recording to {self.filepath}")
 67        return str(self.filepath)
 68
 69    def stop_recording(self) -> None:
 70        """Stop recording and save the video to disk."""
 71        if not self.is_recording:
 72            return
 73
 74        self.is_recording = False
 75
 76        if self.recording_task and not self.recording_task.done():
 77            self.recording_task.cancel()
 78
 79        # This will introduce a delay server-side... there is probably a better way to do this
 80        if self.frames:
 81            try:
 82                logger.info(f"Writing {len(self.frames)} frames to {self.filepath}")
 83                mp.write_video(str(self.filepath), self.frames, fps=self.fps)
 84                logger.info("Video recording stopped and saved")
 85            except Exception as e:
 86                logger.error(f"Error saving video: {e}")
 87        else:
 88            logger.warning("No frames captured, no video file created")
 89
 90        self.frames = []
 91
 92    async def capture_frame(self) -> None:
 93        """Capture and add a single frame to the recording."""
 94        if not self.is_recording:
 95            return
 96
 97        try:
 98            frame, _ = await self.simulator.capture_frame()
 99            self.frames.append(frame)
100
101        except Exception as e:
102            logger.error(f"Error capturing frame: {e}")