From 47010cf91c3222a0c60054d8ddd01d1b2963d648 Mon Sep 17 00:00:00 2001 From: Raghav K Date: Thu, 5 Nov 2026 05:26:45 -0800 Subject: [PATCH] [bugfix] Fall back to PyAV when torchvision.io.read_video is unavailable in preprocessing (#2068) Co-authored-by: Claude Opus 5.5 --- fastvideo/dataset/preprocessing_datasets.py | 5 +- fastvideo/dataset/utils.py | 29 +++++++ .../pipelines/preprocess/preprocess_stages.py | 5 +- .../tests/dataset/test_read_video_tchw.py | 84 +++++++++++++++++++ .../dataset_preparation/prepare_json_file.py | 10 ++- 5 files changed, 123 insertions(+), 10 deletions(-) create mode 100644 fastvideo/tests/dataset/test_read_video_tchw.py diff --git a/fastvideo/dataset/preprocessing_datasets.py b/fastvideo/dataset/preprocessing_datasets.py index 85696e53a..888aae21d 100644 --- a/fastvideo/dataset/preprocessing_datasets.py +++ b/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") diff --git a/fastvideo/dataset/utils.py b/fastvideo/dataset/utils.py index f99a1e816..660837fef 100644 --- a/fastvideo/dataset/utils.py +++ b/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.""" diff --git a/fastvideo/pipelines/preprocess/preprocess_stages.py b/fastvideo/pipelines/preprocess/preprocess_stages.py index 0c649e256..da53857bc 100644 --- a/fastvideo/pipelines/preprocess/preprocess_stages.py +++ b/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) diff --git a/fastvideo/tests/dataset/test_read_video_tchw.py b/fastvideo/tests/dataset/test_read_video_tchw.py new file mode 100644 index 000000000..1e25f6935 --- /dev/null +++ b/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) + diff --git a/scripts/dataset_preparation/prepare_json_file.py b/scripts/dataset_preparation/prepare_json_file.py index cb29494f3..6b1435890 100644 --- a/scripts/dataset_preparation/prepare_json_file.py +++ b/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