fix video encoding/opencv

src/lerobot/scripts/lerobot_create_advantage_video.py

@@ -425,6 +425,56 @@ def _encode_h264(temp_path: Path, output_path: Path) -> None:
     temp_path.unlink(missing_ok=True)
 
 
+def _create_decode_proxy(
+    source_path: Path,
+    proxy_path: Path,
+    *,
+    start_timestamp: float,
+    end_timestamp: float,
+    fps: float,
+    expected_frames: int,
+) -> Path:
+    """Decode an episode segment with system FFmpeg into an OpenCV-safe proxy."""
+    ffmpeg = shutil.which("ffmpeg")
+    if ffmpeg is None:
+        raise RuntimeError(
+            "System ffmpeg is required to decode this dataset's video codec. "
+            "Install ffmpeg and rerun the command."
+        )
+    duration = end_timestamp - start_timestamp
+    result = subprocess.run(  # nosec B603
+        [
+            ffmpeg,
+            "-y",
+            "-ss",
+            f"{start_timestamp:.9f}",
+            "-i",
+            str(source_path),
+            "-t",
+            f"{duration:.9f}",
+            "-an",
+            "-vf",
+            f"fps={fps:.9f}",
+            "-frames:v",
+            str(expected_frames),
+            "-c:v",
+            "libx264",
+            "-preset",
+            "ultrafast",
+            "-crf",
+            "12",
+            "-pix_fmt",
+            "yuv420p",
+            str(proxy_path),
+        ],
+        capture_output=True,
+        text=True,
+    )
+    if result.returncode != 0 or not proxy_path.is_file():
+        raise RuntimeError(f"FFmpeg episode decode failed:\n{result.stderr[-1500:]}")
+    return proxy_path
+
+
 def create_advantage_video(
     *,
     repo_id: str,
@@ -473,19 +523,29 @@ def create_advantage_video(
     tasks = episode_metadata.get("tasks") or []
     task = str(tasks[0]) if len(tasks) else ""
 
-    capture = cv2.VideoCapture(str(video_path))
-    capture.set(cv2.CAP_PROP_POS_MSEC, start_timestamp * 1000)
+    output_dir.mkdir(parents=True, exist_ok=True)
+    camera_name = camera_key.replace(".", "_").replace("/", "_")
+    output_path = output_dir / f"episode_{episode_index:06d}_{camera_name}_advantage.mp4"
+    decode_proxy_path = output_path.with_name(output_path.stem + "_decode_proxy.mp4")
+    _create_decode_proxy(
+        video_path,
+        decode_proxy_path,
+        start_timestamp=start_timestamp,
+        end_timestamp=end_timestamp,
+        fps=fps,
+        expected_frames=len(episode),
+    )
+
+    capture = cv2.VideoCapture(str(decode_proxy_path))
     width = int(capture.get(cv2.CAP_PROP_FRAME_WIDTH))
     frame_height = int(capture.get(cv2.CAP_PROP_FRAME_HEIGHT))
     if width <= 0 or frame_height <= 0:
         capture.release()
+        decode_proxy_path.unlink(missing_ok=True)
         raise RuntimeError(f"Could not read video dimensions from {video_path}")
     dashboard_height = max(132, round(frame_height * 0.27))
     output_height = frame_height + dashboard_height
 
-    output_dir.mkdir(parents=True, exist_ok=True)
-    camera_name = camera_key.replace(".", "_").replace("/", "_")
-    output_path = output_dir / f"episode_{episode_index:06d}_{camera_name}_advantage.mp4"
     temp_path = output_path.with_name(output_path.stem + "_temp.mp4")
     writer = cv2.VideoWriter(
         str(temp_path),
@@ -495,6 +555,7 @@ def create_advantage_video(
     )
     if not writer.isOpened():
         capture.release()
+        decode_proxy_path.unlink(missing_ok=True)
         raise RuntimeError(f"Could not open video writer for {temp_path}")
 
     expected_frames = len(episode)
@@ -534,6 +595,7 @@ def create_advantage_video(
     finally:
         writer.release()
         capture.release()
+        decode_proxy_path.unlink(missing_ok=True)
 
     if written != expected_frames:
         raise RuntimeError(

tests/scripts/test_create_advantage_video.py

@@ -1,11 +1,14 @@
+import shutil
+import subprocess
+
+import cv2
 import numpy as np
 import pandas as pd
 import pytest
 
-pytest.importorskip("cv2")
-
 from lerobot.scripts.lerobot_create_advantage_video import (
     _contiguous_segments,
+    _create_decode_proxy,
     _draw_dashboard,
     _draw_status_badge,
 )
@@ -55,3 +58,43 @@ def test_status_badge_modifies_frame():
     _draw_status_badge(frame, "positive", intervention=False, episode_index=12)
 
     assert frame.sum() > 0
+
+
+@pytest.mark.skipif(shutil.which("ffmpeg") is None, reason="ffmpeg is required")
+def test_ffmpeg_decode_proxy_handles_av1(tmp_path):
+    source = tmp_path / "source.mkv"
+    proxy = tmp_path / "proxy.mp4"
+    encode = subprocess.run(
+        [
+            "ffmpeg",
+            "-y",
+            "-f",
+            "lavfi",
+            "-i",
+            "testsrc=size=64x64:rate=10:duration=1",
+            "-c:v",
+            "libaom-av1",
+            "-cpu-used",
+            "8",
+            "-crf",
+            "40",
+            str(source),
+        ],
+        capture_output=True,
+    )
+    if encode.returncode != 0:
+        pytest.skip("ffmpeg does not provide the libaom-av1 encoder")
+
+    _create_decode_proxy(
+        source,
+        proxy,
+        start_timestamp=0,
+        end_timestamp=1,
+        fps=10,
+        expected_frames=10,
+    )
+
+    capture = cv2.VideoCapture(str(proxy))
+    frame_count = int(capture.get(cv2.CAP_PROP_FRAME_COUNT))
+    capture.release()
+    assert frame_count == 10