[bugfix] Fall back to PyAV when torchvision.io.read_video is unavailable in preprocessing (#2068)

Co-authored-by: Claude Opus 5.5 <noreply@anthropic.com>

fastvideo/dataset/preprocessing_datasets.py

@@ -11,11 +11,11 @@ from typing import Any
 
 import numpy as np
 import torch
-import torchvision
 from einops import rearrange
 from PIL import Image
 from transformers import AutoTokenizer
 
+from fastvideo.dataset.utils import read_video_tchw
 from fastvideo.logger import init_logger
 
 logger = init_logger(__name__)
@@ -258,8 +258,7 @@ class VideoTransformStage(DatasetStage):
         assert os.path.exists(batch.path), f"file {batch.path} do not exist!"
         assert batch.sample_frame_index is not None, "Frame indices must be set before transformation"
 
-        torchvision_video, _, metadata = torchvision.io.read_video(batch.path, output_format="TCHW")
-        video = torchvision_video[batch.sample_frame_index]
+        video = read_video_tchw(batch.path)[batch.sample_frame_index]
         if self.transform is not None:
             video = self.transform(video)
         video = rearrange(video, "t c h w -> c t h w")

fastvideo/dataset/utils.py

@@ -18,6 +18,35 @@ _DTYPE_ALIASES = {
 }
 
 
+
+def read_video_tchw(path: str) -> torch.Tensor:
+    """Read every frame of a video as a ``(T, C, H, W)`` uint8 tensor.
+
+    Tries ``torchvision.io.read_video`` first (available in older
+    torchvision) and falls back to PyAV when the function has been removed
+    (torchvision >= 0.26), the same order as ``fastvideo/tests/utils.py``.
+    """
+    try:
+        from torchvision.io import read_video
+
+        frames, _, _ = read_video(path, output_format="TCHW")
+        return frames
+    except (ImportError, AttributeError):
+        pass
+
+    import av
+
+    with av.open(path) as container:
+        if not container.streams.video:
+            raise RuntimeError(f"No video stream found in {path}")
+        frames = [
+            torch.from_numpy(frame.to_ndarray(format="rgb24")).permute(2, 0, 1) for frame in container.decode(video=0)
+        ]
+    if not frames:
+        raise RuntimeError(f"No video frames decoded from {path}")
+    return torch.stack(frames)
+
+
 @functools.lru_cache(maxsize=None)
 def _normalize_tensor_dtype(dtype_value: Any) -> str:
     """Normalize dtype strings written by NumPy or PyTorch record creators."""

fastvideo/pipelines/preprocess/preprocess_stages.py

@@ -4,12 +4,12 @@ from typing import cast
 
 import numpy as np
 import torch
-import torchvision
 from einops import rearrange
 from torchvision import transforms
 
 from fastvideo.configs.configs import VideoLoaderType
 from fastvideo.dataset.transform import (CenterCropResizeVideo, TemporalRandomCrop)
+from fastvideo.dataset.utils import read_video_tchw
 from fastvideo.fastvideo_args import FastVideoArgs, WorkloadType
 from fastvideo.pipelines.pipeline_batch_info import (ForwardBatch, PreprocessBatch)
 from fastvideo.pipelines.stages.base import PipelineStage
@@ -60,8 +60,7 @@ class VideoTransformStage(PipelineStage):
             if fastvideo_args.preprocess_config.video_loader_type == VideoLoaderType.TORCHCODEC:
                 video = batch.video_loader[i].get_frames_at(frame_indices).data
             elif fastvideo_args.preprocess_config.video_loader_type == VideoLoaderType.TORCHVISION:
-                video, _, _ = torchvision.io.read_video(batch.video_loader[i], output_format="TCHW")
-                video = video[frame_indices]
+                video = read_video_tchw(batch.video_loader[i])[frame_indices]
             else:
                 raise ValueError(f"Invalid video loader type: {fastvideo_args.preprocess_config.video_loader_type}")
             video = self.video_transform(video)

fastvideo/tests/dataset/test_read_video_tchw.py

@@ -0,0 +1,84 @@
+# SPDX-License-Identifier: Apache-2.0
+
+import wave
+
+import av
+import numpy as np
+import pytest
+import torch
+
+from fastvideo.dataset.utils import read_video_tchw
+
+
+def _write_video(path, frames: np.ndarray, fps: int = 8) -> None:
+    with av.open(str(path), mode="w") as container:
+        stream = container.add_stream("libx264rgb", rate=fps)  # lossless RGB H.264
+        stream.width, stream.height = frames.shape[2], frames.shape[1]
+        stream.pix_fmt = "rgb24"
+        stream.options = {"qp": "0"}
+        for frame in frames:
+            for packet in stream.encode(av.VideoFrame.from_ndarray(frame, format="rgb24")):
+                container.mux(packet)
+        for packet in stream.encode():
+            container.mux(packet)
+
+
+@pytest.fixture
+def pyav_only(monkeypatch):
+    """Simulate torchvision >= 0.26, where torchvision.io.read_video no longer exists."""
+    import torchvision.io
+
+    monkeypatch.delattr(torchvision.io, "read_video", raising=False)
+
+
+def test_decodes_every_frame_as_uint8_tchw_exactly(tmp_path, pyav_only) -> None:
+    rng = np.random.default_rng(0)
+    frames = rng.integers(0, 256, size=(9, 32, 48, 3), dtype=np.uint8)
+    path = tmp_path / "clip.mkv"
+    _write_video(path, frames)
+
+    video = read_video_tchw(str(path))
+
+    assert video.dtype == torch.uint8
+    assert tuple(video.shape) == (9, 3, 32, 48)
+    torch.testing.assert_close(video, torch.from_numpy(frames).permute(0, 3, 1, 2))
+
+
+def test_frame_index_selection_matches_the_old_call_pattern(tmp_path, pyav_only) -> None:
+    frames = np.stack([np.full((16, 16, 3), i * 20, dtype=np.uint8) for i in range(6)])
+    path = tmp_path / "ramp.mkv"
+    _write_video(path, frames)
+
+    picked = read_video_tchw(str(path))[[0, 2, 5]]
+
+    assert picked[:, 0, 0, 0].tolist() == [0, 40, 100]
+
+
+def test_missing_video_stream_raises(tmp_path, pyav_only) -> None:
+    path = tmp_path / "audio_only.wav"
+    with wave.open(str(path), "wb") as audio:
+        audio.setnchannels(1)
+        audio.setsampwidth(2)
+        audio.setframerate(8000)
+        audio.writeframes(b"\x00\x00" * 800)
+
+    with pytest.raises(RuntimeError, match="No video stream"):
+        read_video_tchw(str(path))
+
+
+def test_prefers_torchvision_when_available(tmp_path, monkeypatch) -> None:
+    import torchvision.io
+
+    calls = []
+
+    def fake_read_video(path, output_format):
+        calls.append((path, output_format))
+        return torch.zeros(2, 3, 4, 4, dtype=torch.uint8), torch.empty(0), {}
+
+    monkeypatch.setattr(torchvision.io, "read_video", fake_read_video, raising=False)
+
+    video = read_video_tchw("clip.mp4")
+
+    assert calls == [("clip.mp4", "TCHW")]
+    assert tuple(video.shape) == (2, 3, 4, 4)
+

scripts/dataset_preparation/prepare_json_file.py

@@ -6,19 +6,21 @@ import time
 from multiprocessing import Pool, cpu_count
 from pathlib import Path
 
-import torchvision
 from tqdm import tqdm
 
+from fastvideo.dataset.utils import read_video_tchw
+from fastvideo.eval.io.video import probe_video_fps
+
 
 def get_video_info(video_path):
-    """Get video information using torchvision."""
+    """Get video information."""
     # Read video tensor (T, C, H, W)
-    video_tensor, _, info = torchvision.io.read_video(str(video_path), output_format="TCHW", pts_unit="sec")
+    video_tensor = read_video_tchw(str(video_path))
 
     num_frames = video_tensor.shape[0]
     height = video_tensor.shape[2]
     width = video_tensor.shape[3]
-    fps = info.get("video_fps", 0)
+    fps = probe_video_fps(video_path) or 0
     duration = num_frames / fps if fps > 0 else 0
 
     # Extract name