scripts/camera_recorder.py

"""Record the simulator's fixed RGB camera and a synchronized activity timeline."""
import json
from pathlib import Path
import shutil
import subprocess


class CameraRecorder:
    def __init__(self, path, simulation_time, fps=30):
        self.path = Path(path).resolve()
        self.path.parent.mkdir(parents=True, exist_ok=True)
        self.fps = int(fps)
        if not 1 <= self.fps <= 60:
            raise ValueError("Recording frame rate must be between 1 and 60")
        self.start_time = float(simulation_time)
        self.next_time = self.start_time
        self.frames = 0
        self.events = []
        executable = shutil.which("ffmpeg")
        if not executable:
            raise RuntimeError("ffmpeg is required for camera recording")
        self.log = self.path.with_suffix(".ffmpeg.log").open("w")
        self.process = subprocess.Popen([
            executable, "-hide_banner", "-loglevel", "warning", "-y",
            "-f", "rawvideo", "-pixel_format", "rgb24", "-video_size", "960x720",
            "-framerate", str(self.fps), "-i", "pipe:0", "-an",
            "-c:v", "libx264", "-preset", "veryfast", "-crf", "18",
            "-pix_fmt", "yuv420p", "-movflags", "+faststart", str(self.path),
        ], stdin=subprocess.PIPE, stdout=subprocess.DEVNULL, stderr=self.log)

    def event(self, simulation_time, kind, label, **extra):
        self.events.append({"time_seconds": max(0., float(simulation_time) - self.start_time),
                            "type": kind, "label": label, **extra})

    def due(self, simulation_time):
        return float(simulation_time) + 1e-8 >= self.next_time

    def write(self, rgb, simulation_time):
        frame = rgb.tobytes()
        while self.due(simulation_time):
            self.process.stdin.write(frame)
            self.frames += 1
            self.next_time = self.start_time + self.frames / self.fps

    def close(self):
        self.process.stdin.close()
        code = self.process.wait(timeout=30)
        self.log.close()
        if code:
            raise RuntimeError(f"Video encoder failed; see {self.path.with_suffix('.ffmpeg.log')}")
        result = {"video": str(self.path), "fps": self.fps, "frames": self.frames,
                  "duration_seconds": self.frames / self.fps, "events": self.events,
                  "source": "MuJoCo front RGB camera; continuous frames from the actual run"}
        timeline = self.path.parent / "timeline.json"
        timeline.write_text(json.dumps(result, indent=2))
        return {"video": str(self.path), "timeline": str(timeline),
                "frames": self.frames, "duration_seconds": result["duration_seconds"]}