"""Video decoding and processing helpers for Yasa2.""" from __future__ import annotations import io import urllib.request from typing import Callable, Dict, List, Union import imageio.v3 as iio import numpy as np import torch from transformers.video_processing_utils import BaseVideoProcessor from .image_processing_yasa2 import Yasa2ImageProcessor def frame_sampling_uniform(num_frames: int, total_frames: int) -> List[int]: """Sample frames uniformly across the video timeline. Args: num_frames: Number of frames to sample. total_frames: Total number of frames. Returns: List[int]: Equally spaced frame indices (clamped to total_frames). """ if num_frames >= total_frames: return list(range(total_frames)) interval = total_frames / num_frames start_point = interval / 2 return ( (np.arange(start_point, total_frames, interval)).astype(int).tolist() ) def frame_sampling_random(num_frames: int, total_frames: int) -> List[int]: """Sample frames randomly without replacement. Args: num_frames: Number of frames to sample. total_frames: Total number of frames. Returns: List[int]: Random unique frame indices. """ if total_frames <= num_frames: return list(range(total_frames)) return np.random.choice( np.arange(total_frames), num_frames, replace=False ).tolist() def frame_sampling_chunked(num_frames: int, total_frames: int) -> List[int]: """Sample frames by dividing the video into chunks and picking one index per chunk. Args: num_frames: Number of frames to sample. total_frames: Total number of frames. Returns: List[int]: One randomly chosen index per chunk. """ if total_frames <= num_frames: return list(range(total_frames)) chunk_size = total_frames // num_frames extra_frames = total_frames % num_frames sampled_frames = [] for i in range(num_frames): start = i * chunk_size + min(i, extra_frames) end = start + chunk_size + (1 if i < extra_frames else 0) sampled_frames.append(np.random.randint(start, end)) return sampled_frames def _read_bytes_from_uri(uri: str) -> bytes: """Read bytes from a local path or HTTP(S) URL. Args: uri: Local file path or HTTP(S) URL. Returns: Raw bytes content. """ if uri.startswith("http://") or uri.startswith("https://"): with urllib.request.urlopen(uri) as response: return response.read() with open(uri, "rb") as f: return f.read() def video_rgb_decoder_iio( video_bytes: bytes, num_frames: int, frame_sampler: Callable[[int, int], List[int]], plugin: str = "pyav", skip_errors: bool = False, ) -> Dict[str, Union[np.ndarray, float, List[int]]]: """Decode video bytes into sampled RGB frames together with metadata. Args: video_bytes: Raw video bytes. num_frames: Number of frames to sample. frame_sampler: Frame sampling function. plugin: ImageIO plugin name. skip_errors: Whether to return error info instead of raising. Returns: Dict[str, Union[np.ndarray, float, List[int]]]: Pixel values, fps, taken indices, and sampled count. """ try: with io.BytesIO(video_bytes) as video_io: properties = iio.improps(video_io, plugin=plugin) total_frames, height, width, channels = properties.shape if channels != 3: raise NotImplementedError( f"Video with {channels} channels not supported." ) video_io.seek(0) metadata = iio.immeta(video_io, plugin=plugin) fps = metadata["fps"] if total_frames == 0: total_frames = int(fps * metadata["duration"]) # Mirror training-time sampling behavior (total_frames - 1). frame_idxs = set(frame_sampler(num_frames, total_frames - 1)) frame_idxs_actual = [] pixel_values = [] video_io.seek(0) for idx, frame in enumerate( iio.imiter(video_io, plugin=plugin, thread_type="FRAME") ): if idx in frame_idxs: frame_idxs.remove(idx) frame_idxs_actual.append(idx) pixel_values.append(frame) if not frame_idxs: break if frame_idxs and not skip_errors: raise ValueError(f"Failed to read frames {frame_idxs}.") pixel_values = np.stack(pixel_values, axis=0) return { "pixel_values": pixel_values, "fps": fps, "frame_idxs": frame_idxs_actual, "num_frames": len(frame_idxs_actual), } except Exception as exc: if not skip_errors: raise return {"error": str(exc)} def video_rgb_decoder_factory( num_frames: int, sampling: str = "uniform", skip_errors: bool = False ) -> Callable[[bytes], Dict[str, Union[np.ndarray, float, List[int]]]]: """Create a decoder that samples frames according to the chosen strategy. Args: num_frames: Number of frames to sample. sampling: Sampling strategy name. skip_errors: Whether to return error info instead of raising. Returns: Callable[[bytes], Dict[str, Union[np.ndarray, float, List[int]]]]: Decoder that maps raw bytes to decoded frames/metadata. """ if sampling == "uniform": frame_sampler_fn = frame_sampling_uniform elif sampling == "random": frame_sampler_fn = frame_sampling_random elif sampling == "chunk": frame_sampler_fn = frame_sampling_chunked else: raise NotImplementedError( f"Frame sampling method {sampling} not implemented." ) return lambda video_bytes: video_rgb_decoder_iio( video_bytes, num_frames=num_frames, frame_sampler=frame_sampler_fn, skip_errors=skip_errors, ) class Yasa2VideoProcessor(BaseVideoProcessor): """Video processor that samples frames and applies the ConvNeXt image processor.""" model_input_names = ["pixel_values", "patch_attention_mask"] def __init__( self, num_frames: int = 6, frame_sample_mode: str = "chunk", patch_size: int = 14, size: int = 512, vision_patch_stride: int = 32, image_mean: List[float] | None = None, image_std: List[float] | None = None, max_num_frames: int | None = None, **kwargs, ) -> None: """Initialize the video processor. Args: num_frames: Number of frames to sample per video. frame_sample_mode: Sampling strategy for frames. patch_size: Vision patch size for attention mask. size: Input resolution for the image processor. vision_patch_stride: Effective stride of the vision encoder. image_mean: Mean values for normalization. image_std: Std values for normalization. max_num_frames: Optional padding target for frames. **kwargs: Passed to BaseVideoProcessor. """ super().__init__(**kwargs) self.num_frames = num_frames self.frame_sample_mode = frame_sample_mode self.patch_size = patch_size self.size = size self.vision_patch_stride = vision_patch_stride self.image_mean = image_mean or [0.485, 0.456, 0.406] self.image_std = image_std or [0.229, 0.224, 0.225] self.max_num_frames = max_num_frames self.image_processor = Yasa2ImageProcessor( size={"shortest_edge": size}, crop_size={"height": size, "width": size}, do_resize=True, do_normalize=True, image_mean=self.image_mean, image_std=self.image_std, patch_size=patch_size, ) def decode_video( self, video: Union[str, bytes] ) -> Dict[str, Union[np.ndarray, float, List[int]]]: """Decode a video path or raw bytes into sampled frames. Args: video: Video path/URL or raw bytes. Returns: Dict[str, Union[np.ndarray, float, List[int]]]: Decoded frames, fps, sampled indices, and frame count. """ if isinstance(video, str): video_bytes = _read_bytes_from_uri(video) else: video_bytes = video decoder = video_rgb_decoder_factory( num_frames=self.num_frames, sampling=self.frame_sample_mode ) return decoder(video_bytes) def to_dict( self, ) -> Dict[str, Union[int, str, float, List[float], None, Dict[str, str]]]: """Return a JSON-serializable config for logging and saving. Returns: Dict[str, Union[int, str, float, List[float], None, Dict[str, str]]]: Processor attributes without None values. """ output = super().to_dict() output.pop("image_processor", None) # Ensure the encoder stride used for patch masking is serialized. output["vision_patch_stride"] = self.vision_patch_stride # Drop unset values to avoid serializing nulls into the config. return { key: value for key, value in output.items() if value is not None } def preprocess( self, videos: Union[str, bytes, np.ndarray, List[np.ndarray]], return_tensors: str | None = "pt", **kwargs, ) -> Dict[str, torch.Tensor]: """Preprocess videos into pixel values and patch attention masks. Args: videos: Video path/URL, raw bytes, or frame array(s). return_tensors: Tensor type to return. **kwargs: Unused extra arguments. Returns: Dict[str, torch.Tensor]: `pixel_values` and `patch_attention_mask` tensors, padded as needed. """ if isinstance(videos, (str, bytes)): video_datum = self.decode_video(videos) pixel_values = video_datum["pixel_values"] else: pixel_values = videos image_outputs = self.image_processor( images=pixel_values, return_tensors="pt" ) img_tensor = image_outputs["pixel_values"] if "patch_attention_mask" in image_outputs: patch_attention_mask = image_outputs["patch_attention_mask"] else: # ConvNeXt outputs features at stride 32 (512 -> 16 grid), so build a grid mask at that resolution. grid_size = max(1, self.size // self.vision_patch_stride) patch_attention_mask = torch.ones( ( img_tensor.shape[0], grid_size, grid_size, ), dtype=torch.bool, ) if ( self.max_num_frames is not None and img_tensor.shape[0] < self.max_num_frames ): pad_frames = self.max_num_frames - img_tensor.shape[0] img_tensor = torch.cat( [ img_tensor, torch.zeros( ( pad_frames, img_tensor.shape[1], img_tensor.shape[2], img_tensor.shape[3], ) ), ], dim=0, ) patch_attention_mask = torch.cat( [ patch_attention_mask, # Mask out padded frames to avoid leaking zero-padded inputs. torch.zeros( ( pad_frames, patch_attention_mask.shape[1], patch_attention_mask.shape[2], ), dtype=torch.bool, ), ], dim=0, ) return { "pixel_values": img_tensor, "patch_attention_mask": patch_attention_mask, } Yasa2VideoProcessor.register_for_auto_class()