# coding=utf-8 # Copyright 2025 The FNLP Vision Team and The HuggingFace Inc. team. All rights reserved. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. """video processor class for Moss-VL.""" import json import logging as system_logging import math import os import re import subprocess import traceback from functools import lru_cache from typing import Any, Dict, List, Optional, Union import numpy as np import torch from joblib import Parallel, delayed from torchcodec.decoders import VideoDecoder from transformers.feature_extraction_utils import BatchFeature from transformers.image_utils import ChannelDimension, PILImageResampling, SizeDict, get_image_size, validate_kwargs from transformers.processing_utils import Unpack, VideosKwargs from transformers.utils import TensorType, add_start_docstrings, logging from transformers.video_processing_utils import BASE_VIDEO_PROCESSOR_DOCSTRING, BaseVideoProcessor from transformers.video_utils import VideoMetadata, group_videos_by_shape, reorder_videos logger = logging.get_logger(__name__) TORCHCODEC_TIMESTAMP_EPSILON = 1e-6 # ----------------------------------------------------------------------------- # Torchcodec video frame extraction utilities # ----------------------------------------------------------------------------- def check_video_for_extra_streams_and_errors(video_path: str) -> dict: """ Check if video file has abnormal streams or errors reported by ffprobe. Args: video_path: Path to the video file. Returns: A dictionary containing: - 'has_extra_streams': bool, whether there are streams other than video and audio. - 'unsupported_codec_errors': list, all "Unsupported codec" error messages. - 'ffprobe_output_error': str, other errors/warnings from ffprobe stderr. - 'ffprobe_successful': bool, whether ffprobe command executed successfully (return code 0). - 'stream_details': list, codec_type and index for each stream. - 'num_streams': int, total number of streams identified in the video file. """ result = { 'has_extra_streams': False, 'unsupported_codec_errors': [], 'ffprobe_output_error': '', 'ffprobe_successful': False, 'stream_details': [], 'num_streams': 0 } command = [ "ffprobe", "-v", "error", "-show_streams", "-show_format", "-of", "json", video_path ] try: process = subprocess.run( command, capture_output=True, text=True, check=False ) result['ffprobe_successful'] = (process.returncode == 0) if process.stderr: result['ffprobe_output_error'] = process.stderr unsupported_codec_pattern = re.compile(r"Unsupported codec with id \d+ for input stream \d+") result['unsupported_codec_errors'] = unsupported_codec_pattern.findall(process.stderr) if process.stdout: ffprobe_data = json.loads(process.stdout) if 'streams' in ffprobe_data: result['num_streams'] = len(ffprobe_data['streams']) for stream in ffprobe_data['streams']: stream_type = stream.get('codec_type') stream_index = stream.get('index') result['stream_details'].append({'index': stream_index, 'codec_type': stream_type}) if stream_type not in ['video', 'audio']: result['has_extra_streams'] = True if 'format' in ffprobe_data and 'nb_streams' in ffprobe_data['format']: if result['num_streams'] == 0: result['num_streams'] = ffprobe_data['format']['nb_streams'] elif result['num_streams'] != ffprobe_data['format']['nb_streams']: logger.warning( f"Number of streams in 'streams' list ({result['num_streams']}) " f"differs from 'nb_streams' in 'format' ({ffprobe_data['format']['nb_streams']})." ) except FileNotFoundError: result['ffprobe_output_error'] = "ffprobe command not found. Please ensure FFmpeg is installed and in your PATH." result['ffprobe_successful'] = False except json.JSONDecodeError: result['ffprobe_output_error'] = "Failed to parse ffprobe JSON output. Check ffprobe installation or video file." result['ffprobe_successful'] = False except Exception as e: result['ffprobe_output_error'] = f"An unexpected error occurred: {e}" result['ffprobe_successful'] = False return result def remove_video_extra_stream_ffmpeg(input_video: str, output_video: str) -> bool: """ Remove extra streams from video using ffmpeg. Args: input_video: Path to input video. output_video: Path to output video. Returns: bool: True if successful, False otherwise. """ command_list = [ "ffmpeg", "-y", "-i", input_video, "-map", "0:v:0", "-c", "copy", "-an", "-sn", "-dn", "-map_metadata", "-1", "-map_chapters", "-1", "-movflags", "faststart", output_video, ] try: subprocess.run(command_list, shell=False, check=True, capture_output=True) return True except subprocess.CalledProcessError as e: system_logging.error(f"Command execution failed with return code: {e.returncode}, video: {input_video}") system_logging.error(f"Error output:\n{e.stderr}") return False except FileNotFoundError: system_logging.error("Error: ffmpeg command not found. Please ensure ffmpeg is installed and in PATH.") return False except Exception as e: system_logging.error(f"Unexpected error executing command: {e}, video: {input_video}", exc_info=True) return False def clean_video_streams(video_path: str) -> str: """ Clean video streams if extra streams are detected. Args: video_path: Path to the video file. Returns: str: Path to cleaned video (or original if no cleaning needed). """ ffprobe_res = check_video_for_extra_streams_and_errors(video_path) if ffprobe_res['has_extra_streams']: base_name = os.path.basename(video_path) output_folder = os.path.dirname(video_path) file_name_without_ext, file_ext = os.path.splitext(base_name) new_base_name = f"{file_name_without_ext}_fix{file_ext}" video_path_output = os.path.join(output_folder, new_base_name) process_flag = remove_video_extra_stream_ffmpeg(video_path, video_path_output) if not process_flag: logger.warning("Failed to remove extra streams with ffmpeg") return video_path return video_path_output return video_path @lru_cache(maxsize=8192) def cached_clean_video_streams(video_path: str) -> str: return clean_video_streams(video_path) def clamp_timestamps_for_torchcodec(timestamps: List[float], torchcodec_metadata) -> List[float]: if not timestamps: return timestamps min_pts = torchcodec_metadata.begin_stream_seconds_from_content if min_pts is None: min_pts = 0.0 max_pts_candidates = [] if torchcodec_metadata.num_frames_from_content and torchcodec_metadata.average_fps: max_pts_candidates.append( (torchcodec_metadata.num_frames_from_content - 1) / torchcodec_metadata.average_fps + min_pts ) if torchcodec_metadata.end_stream_seconds_from_content is not None: # TorchCodec requires requested PTS to be strictly smaller than the content end. max_pts_candidates.append(torchcodec_metadata.end_stream_seconds_from_content - TORCHCODEC_TIMESTAMP_EPSILON) if not max_pts_candidates and torchcodec_metadata.duration_seconds is not None: max_pts_candidates.append(torchcodec_metadata.duration_seconds - TORCHCODEC_TIMESTAMP_EPSILON) if max_pts_candidates: max_pts = max(min_pts, min(max_pts_candidates)) return [max(min_pts, min(float(t), max_pts)) for t in timestamps] if min_pts > 0: return [max(min_pts, float(t)) for t in timestamps] return [float(t) for t in timestamps] def split_indices(indices: List[Union[int, float]], num_chunks: int) -> List[List[Union[int, float]]]: """ Split an index list into roughly equal chunks. Args: indices: List of indices to split. num_chunks: Number of chunks to create. Returns: List of index chunks. """ chunk_size = len(indices) // num_chunks chunks = [] for i in range(num_chunks - 1): chunks.append(indices[i * chunk_size:(i + 1) * chunk_size]) chunks.append(indices[(num_chunks - 1) * chunk_size:]) return chunks def decode_sequentially(indices: List[int], video_path: str, ffmpeg_threads: int = 0): """ Decode frames sequentially from a video. Args: indices: List of frame indices to decode. video_path: Path to the video file. ffmpeg_threads: Number of ffmpeg threads to use. Returns: FrameBatch from torchcodec. """ decoder = VideoDecoder(video_path, num_ffmpeg_threads=ffmpeg_threads) try: return decoder.get_frames_at(indices) finally: del decoder def decode_with_multithreading(indices: List[int], num_threads: int, video_path: str) -> dict: """ Decode frames using multithreading with joblib. Args: indices: List of frame indices to decode. num_threads: Number of threads to use. video_path: Path to the video file. Returns: dict: Contains 'data', 'duration_seconds', 'pts_seconds' tensors. """ chunks = split_indices(indices, num_chunks=num_threads) results = Parallel(n_jobs=num_threads, prefer="threads", verbose=0)( delayed(decode_sequentially)(chunk, video_path) for chunk in chunks ) return { "data": torch.cat([frame_batch.data for frame_batch in results], dim=0), "duration_seconds": torch.cat([frame_batch.duration_seconds for frame_batch in results], dim=0), "pts_seconds": torch.cat([frame_batch.pts_seconds for frame_batch in results], dim=0) } def decode_sequentially_timestamp(timestamp_list: List[float], video_path: str, ffmpeg_threads: int = 0): """ Decode frames sequentially from a video based on timestamps. Args: timestamp_list: List of timestamps (in seconds) to decode. video_path: Path to the video file. ffmpeg_threads: Number of ffmpeg threads to use. Returns: FrameBatch from torchcodec. """ decoder = VideoDecoder(video_path, num_ffmpeg_threads=ffmpeg_threads) try: metadata = decoder.metadata timestamp_list = clamp_timestamps_for_torchcodec(timestamp_list, metadata) return decoder.get_frames_played_at(timestamp_list) finally: del decoder def timestamp_decode_with_multithreading(timestamp_list: List[float], num_threads: int, video_path: str) -> dict: """ Decode frames using multithreading based on timestamps. Args: timestamp_list: List of timestamps (in seconds) to decode. num_threads: Number of threads to use. video_path: Path to the video file. Returns: dict: Contains 'data', 'duration_seconds', 'pts_seconds' tensors. """ chunks = split_indices(timestamp_list, num_chunks=num_threads) results = Parallel(n_jobs=num_threads, prefer="threads", verbose=0)( delayed(decode_sequentially_timestamp)(chunk, video_path) for chunk in chunks ) # Concatenate results from all threads data_list = [frame_batch.data for frame_batch in results] duration_list = [frame_batch.duration_seconds for frame_batch in results] pts_list = [frame_batch.pts_seconds for frame_batch in results] if not data_list: logger.warning("No frames were successfully decoded.") return {"data": torch.empty(0), "duration_seconds": torch.empty(0), "pts_seconds": torch.empty(0)} return { "data": torch.cat(data_list, dim=0), "duration_seconds": torch.cat(duration_list, dim=0), "pts_seconds": torch.cat(pts_list, dim=0) } def extract_frames_with_torchcodec( video_path: str, sample_frames_count: int, num_threads: int = 4, ) -> Optional[dict]: """ Extract frames from video using torchcodec with multithreading. Args: video_path: Path to the video file. sample_frames_count: Number of frames to sample. num_threads: Number of threads to use for extraction. sampling_method: Sampling method, either "index" (uniform frame indices) or "timestamp" (uniform timestamps). Returns: dict: Contains 'data' (N, C, H, W), 'duration_seconds' (N,), 'pts_seconds' (N,) tensors. Returns None if extraction fails. """ try: video_path = cached_clean_video_streams(video_path) decoder = VideoDecoder(video_path, num_ffmpeg_threads=0) metadata = decoder.metadata total_frames_in_video = metadata.num_frames_from_content effective_sample_count = min(sample_frames_count, total_frames_in_video) if effective_sample_count == 0: logger.error("Cannot extract frames: video has 0 frames or specified frame count is 0") return None # Generate uniform frame indices frame_indices = np.linspace(0, total_frames_in_video - 1, effective_sample_count).astype(np.int32) # Ensure indices are valid and remove duplicates frame_indices = np.unique(np.clip(frame_indices, 0, total_frames_in_video - 1)) result = decode_with_multithreading(frame_indices.tolist(), num_threads=num_threads, video_path=video_path) # Add frame_indices to the result for later use result["frame_indices"] = frame_indices return result except Exception: traceback.print_exc() return None def smart_resize( num_frames: int, height: int, width: int, temporal_factor: int = 1, factor: int = 32, min_pixels: int = 128 * 128, max_pixels: int = 16 * 16 * 2 * 2 * 2 * 6144, per_frame_min_pixels: int = None, per_frame_max_pixels: int = None, ): if num_frames < temporal_factor: raise ValueError(f"t:{num_frames} must be larger than temporal_factor:{temporal_factor}") if height < factor or width < factor: raise ValueError(f"height:{height} or width:{width} must be larger than factor:{factor}") elif max(height, width) / min(height, width) > 200: raise ValueError( f"absolute aspect ratio must be smaller than 200, got {max(height, width) / min(height, width)}" ) h_bar = round(height / factor) * factor w_bar = round(width / factor) * factor t_bar = round(num_frames / temporal_factor) * temporal_factor # Step 1: Apply per-frame upper limit constraint if per_frame_max_pixels is not None and h_bar * w_bar > per_frame_max_pixels: beta = math.sqrt((height * width) / per_frame_max_pixels) h_bar = max(factor, math.floor(height / beta / factor) * factor) w_bar = max(factor, math.floor(width / beta / factor) * factor) # Step 2: Apply 3D volume constraints (frames * height * width) if t_bar * h_bar * w_bar > max_pixels: beta = math.sqrt((num_frames * height * width) / max_pixels) h_bar = max(factor, math.floor(height / beta / factor) * factor) w_bar = max(factor, math.floor(width / beta / factor) * factor) elif t_bar * h_bar * w_bar < min_pixels: beta = math.sqrt(min_pixels / (num_frames * height * width)) h_bar = math.ceil(height * beta / factor) * factor w_bar = math.ceil(width * beta / factor) * factor # Step 3: Ensure per-frame lower limit is respected (after volume constraint) # This guarantees single frame stays within [per_frame_min_pixels, per_frame_max_pixels] if per_frame_min_pixels is not None and h_bar * w_bar < per_frame_min_pixels: beta = math.sqrt(per_frame_min_pixels / (height * width)) h_bar = math.ceil(height * beta / factor) * factor w_bar = math.ceil(width * beta / factor) * factor return h_bar, w_bar class MossVLVideoProcessorInitKwargs(VideosKwargs): patch_size: Optional[int] temporal_patch_size: Optional[int] merge_size: Optional[int] min_frames: Optional[int] max_frames: Optional[int] video_fps: Optional[Union[int, float]] num_extract_threads: Optional[int] # Total 3D volume budget across all videos; distributed proportionally per video by T*H*W video_max_pixels: Optional[int] @add_start_docstrings( "Constructs a fast Moss-VL video processor that dynamically resizes videos based on the original videos.", BASE_VIDEO_PROCESSOR_DOCSTRING, """ patch_size (`int`, *optional*, defaults to 16): The spacial patch size of the vision encoder. temporal_patch_size (`int`, *optional*, defaults to 1): The temporal patch size of the vision encoder. merge_size (`int`, *optional*, defaults to 2): The merge size of the vision encoder to llm encoder. video_fps (`float`, *optional*, defaults to 1.0): Target frames per second for video sampling. min_frames (`int`, *optional*, defaults to 1): Minimum number of frames to sample from a video. max_frames (`int`, *optional*, defaults to 256): Maximum number of frames to sample from a video. num_extract_threads (`int`, *optional*, defaults to 4): Number of threads to use for frame extraction. """, ) class MossVLVideoProcessor(BaseVideoProcessor): resample = PILImageResampling.BICUBIC size = {"shortest_edge": 128 * 32 * 32, "longest_edge": 32 * 32 * 768} image_mean = [0.5, 0.5, 0.5] image_std = [0.5, 0.5, 0.5] do_resize = True do_rescale = True do_normalize = True do_convert_rgb = True patch_size = 16 temporal_patch_size = 1 merge_size = 2 video_fps = 1.0 min_frames = 1 max_frames = 256 num_extract_threads = 4 do_sample_frames = True # Total 3D volume budget across all videos; distributed proportionally per video by T*H*W video_max_pixels = None # read from config valid_kwargs = MossVLVideoProcessorInitKwargs model_input_names = ["pixel_values_videos", "video_grid_thw"] def __init__(self, **kwargs: Unpack[MossVLVideoProcessorInitKwargs]): super().__init__(**kwargs) if self.size is not None and ( self.size.get("shortest_edge", None) is None or self.size.get("longest_edge", None) is None ): raise ValueError("size must contain 'shortest_edge' and 'longest_edge' keys.") def _further_process_kwargs( self, size: Optional[SizeDict] = None, **kwargs, ) -> dict: """ Update kwargs that need further processing before being validated Can be overridden by subclasses to customize the processing of kwargs. """ if size is not None and ("shortest_edge" not in size or "longest_edge" not in size): raise ValueError("size must contain 'shortest_edge' and 'longest_edge' keys.") return super()._further_process_kwargs(size=size, **kwargs) def _get_video_path_from_input(self, video_input: Union[str, Dict[str, Any]]) -> str: """Normalize a video input into a video path.""" if isinstance(video_input, dict): return video_input["video_path"] return video_input def _get_video_duration_seconds(self, video_input: Union[str, Dict[str, Any]]) -> float: """Get video duration in seconds for weighted frame-budget allocation.""" video_path = cached_clean_video_streams(self._get_video_path_from_input(video_input)) decoder = VideoDecoder(video_path, num_ffmpeg_threads=0) try: metadata = decoder.metadata duration = None if ( metadata.end_stream_seconds_from_content is not None and metadata.begin_stream_seconds_from_content is not None ): duration = metadata.end_stream_seconds_from_content - metadata.begin_stream_seconds_from_content if duration is None or duration <= 0: duration = metadata.duration_seconds return max(0.0, float(duration or 0.0)) finally: del decoder def _allocate_max_frames_for_multiple_videos( self, video_inputs: List[Union[str, Dict[str, Any]]], total_max_frames: Optional[int], ) -> List[Optional[int]]: """ Treat max_frames as a total budget for multi-video input and allocate it by duration. The returned values are per-video max_frames. Segment dict inputs still keep their existing per-segment weighting logic after receiving the video-level allocation. """ if not video_inputs: return [] if total_max_frames is None or len(video_inputs) == 1: return [total_max_frames] * len(video_inputs) total_max_frames = int(total_max_frames) num_videos = len(video_inputs) if total_max_frames < num_videos: logger.warning( "Received max_frames=%s for %s videos. At least one frame per video is required, " "so falling back to 1 frame per video.", total_max_frames, num_videos, ) return [1] * num_videos video_durations = [self._get_video_duration_seconds(video_input) for video_input in video_inputs] total_duration = sum(video_durations) # Reserve one frame per video first, then distribute the remaining budget by duration. allocations = [1] * num_videos remaining_budget = total_max_frames - num_videos if remaining_budget == 0: return allocations if total_duration <= 0: raw_extra_allocations = [remaining_budget / num_videos] * num_videos else: raw_extra_allocations = [ remaining_budget * (duration / total_duration) for duration in video_durations ] base_extra_allocations = [int(math.floor(value)) for value in raw_extra_allocations] allocations = [base + extra for base, extra in zip(allocations, base_extra_allocations)] remainder = remaining_budget - sum(base_extra_allocations) if remainder > 0: fractional_parts = [ (raw_value - base_value, index) for index, (raw_value, base_value) in enumerate(zip(raw_extra_allocations, base_extra_allocations)) ] fractional_parts.sort(key=lambda item: (-item[0], item[1])) for _, index in fractional_parts[:remainder]: allocations[index] += 1 return allocations def calculate_num_frames( self, metadata: VideoMetadata, num_frames: Optional[int] = None, fps: Optional[Union[int, float]] = None, min_frames: Optional[int] = None, max_frames: Optional[int] = None, **kwargs, ) -> int: """ Calculate the number of frames to sample using fps-based logic with min/max constraints. Logic: 1. Calculate target_frames based on fps and video duration 2. Apply min_frames and max_frames constraints 3. Apply max_allowed_frames protection (rough cap from total video_max_pixels budget) 4. Return the number of frames to sample Args: metadata (`VideoMetadata`): Metadata of the video containing information about total duration, fps and total number of frames. num_frames (`int`, *optional*): Maximum number of frames to sample. If provided, overrides fps-based calculation. fps (`int` or `float`, *optional*): Target frames to sample per second. Defaults to `self.video_fps`. min_frames (`int`, *optional*): Minimum number of frames to sample. If None, uses self.min_frames. max_frames (`int`, *optional*): Maximum number of frames to sample. If None, uses self.max_frames. Returns: int: Number of frames to sample. """ if fps is not None and num_frames is not None: raise ValueError("`num_frames` and `fps` are mutually exclusive arguments, please use only one!") total_num_frames = metadata.total_num_frames # Use provided min/max or fall back to defaults effective_min_frames = min_frames if min_frames is not None else self.min_frames effective_max_frames = max_frames if max_frames is not None else self.max_frames # Rough per-video frame cap derived from the multi-video total budget # (exact allocation happens later in _preprocess via weighted distribution) per_frame_min_pixels = self.size.get("shortest_edge", None) if self.size else None video_max_pixels = getattr(self, "video_max_pixels", None) if per_frame_min_pixels is not None and video_max_pixels is not None and per_frame_min_pixels > 0: max_allowed_frames = video_max_pixels // per_frame_min_pixels effective_max_frames = min(effective_max_frames, max_allowed_frames) # Get video duration if hasattr(metadata, 'duration') and metadata.duration is not None: duration = metadata.duration else: video_fps = metadata.fps if video_fps is not None and video_fps > 0: duration = total_num_frames / video_fps else: # Fallback: assume 24 fps video_fps = 24.0 duration = total_num_frames / video_fps logger.warning_once( "Could not determine video fps from metadata, defaulting to 24 fps for duration calculation." ) # Use provided fps or default target_fps = fps if fps is not None else self.video_fps # Calculate target frames based on fps and duration if num_frames is None: # Calculate how many frames we should sample based on target fps target_total_frames = int(math.ceil(duration * target_fps - 1e-6)) # Apply min/max constraints sample_frames = max(target_total_frames, effective_min_frames) sample_frames = min(sample_frames, effective_max_frames, total_num_frames) else: # If num_frames is explicitly provided, use it directly with constraints sample_frames = min(max(num_frames, effective_min_frames), effective_max_frames, total_num_frames) return sample_frames def _decode_timestamps_with_decoder( self, decoder: VideoDecoder, timestamps: List[float], chunk_size: int = 128, ) -> torch.Tensor: if not timestamps: return torch.empty(0) frame_chunks = [] for start in range(0, len(timestamps), chunk_size): frame_batch = decoder.get_frames_played_at(timestamps[start:start + chunk_size]) frame_chunks.append(frame_batch.data) if len(frame_chunks) == 1: return frame_chunks[0] return torch.cat(frame_chunks, dim=0) def _clamp_timestamps_for_decoder( self, timestamps: List[float], torchcodec_metadata, ) -> List[float]: return clamp_timestamps_for_torchcodec(timestamps, torchcodec_metadata) def _fetch_video_segments_batched( self, video_path: str, segments: List[List[float]], min_frames: Optional[int] = None, max_frames: Optional[int] = None, video_fps: Optional[float] = None, ): min_frames = max(1, min_frames if min_frames is not None else self.min_frames) max_frames = max(1, max_frames if max_frames is not None else self.max_frames) target_video_fps = video_fps if video_fps is not None else self.video_fps video_path = cached_clean_video_streams(video_path) decoder = VideoDecoder(video_path, num_ffmpeg_threads=0) try: torchcodec_metadata = decoder.metadata source_video_fps = torchcodec_metadata.average_fps duration = None if ( torchcodec_metadata.end_stream_seconds_from_content is not None and torchcodec_metadata.begin_stream_seconds_from_content is not None ): duration = ( torchcodec_metadata.end_stream_seconds_from_content - torchcodec_metadata.begin_stream_seconds_from_content ) if duration is None or duration <= 0: duration = torchcodec_metadata.duration_seconds segment_durations = [ segment[1] - segment[0] if len(segment) == 2 else None for segment in segments ] total_segment_duration = sum(d for d in segment_durations if d is not None) num_range_segments = sum(1 for d in segment_durations if d is not None) segment_timestamps = [] decode_timestamps = [] for i, segment in enumerate(segments): if len(segment) == 1: actual_timestamps = self._clamp_timestamps_for_decoder([segment[0]], torchcodec_metadata) segment_timestamps.append(actual_timestamps) decode_timestamps.extend(actual_timestamps) continue start_time, end_time = segment segment_duration = end_time - start_time target_frames = int(math.ceil(segment_duration * target_video_fps)) if total_segment_duration > 0: weight = segment_durations[i] / total_segment_duration else: weight = 1.0 / num_range_segments if num_range_segments > 0 else 1.0 weighted_min_frames = max(1, int(round(min_frames * weight))) weighted_max_frames = max(1, int(round(max_frames * weight))) target_frames = max(target_frames, weighted_min_frames) target_frames = min(target_frames, weighted_max_frames) if target_frames == 1: actual_timestamps = [start_time] else: actual_timestamps = np.linspace( start_time, end_time, target_frames, endpoint=False, ).tolist() actual_timestamps = self._clamp_timestamps_for_decoder(actual_timestamps, torchcodec_metadata) segment_timestamps.append(actual_timestamps) decode_timestamps.extend(actual_timestamps) flat_frames = self._decode_timestamps_with_decoder(decoder, decode_timestamps) videos = [] metadata = [] frame_offset = 0 for actual_timestamps in segment_timestamps: sample_count = len(actual_timestamps) video_tensor = flat_frames[frame_offset:frame_offset + sample_count] frame_offset += sample_count video_metadata = VideoMetadata( total_num_frames=sample_count, fps=source_video_fps, duration=duration, video_backend="torchcodec", height=torchcodec_metadata.height, width=torchcodec_metadata.width, frames_indices=None, ) video_metadata.actual_timestamps = actual_timestamps videos.append(video_tensor) metadata.append(video_metadata) return videos, metadata finally: del decoder def _fetch_video_segment( self, video_path: str, segment: List[float], min_frames: Optional[int] = None, max_frames: Optional[int] = None, video_fps: Optional[float] = None, ): """ Fetch video frames for a specific segment. Args: video_path: Path to the video file segment: [start, end] for a segment (left-closed, right-open) or [time] for a single frame min_frames: Minimum frames for this segment (weighted). Defaults to self.min_frames. Must be >= 1. max_frames: Maximum frames for this segment (weighted). Defaults to self.max_frames. Must be >= 1. video_fps: Target frames per second for video sampling. If None, uses self.video_fps. Returns: Tuple of (video_tensor, video_metadata) """ # Use provided min/max or fall back to defaults, ensure >= 1 min_frames = max(1, min_frames if min_frames is not None else self.min_frames) max_frames = max(1, max_frames if max_frames is not None else self.max_frames) # Use provided video_fps or fall back to self.video_fps target_video_fps = video_fps if video_fps is not None else self.video_fps video_path = clean_video_streams(video_path) decoder = VideoDecoder(video_path, num_ffmpeg_threads=0) try: torchcodec_metadata = decoder.metadata video_fps = torchcodec_metadata.average_fps # Calculate duration duration = None if torchcodec_metadata.end_stream_seconds_from_content is not None and torchcodec_metadata.begin_stream_seconds_from_content is not None: duration = torchcodec_metadata.end_stream_seconds_from_content - torchcodec_metadata.begin_stream_seconds_from_content if duration is None or duration <= 0: duration = torchcodec_metadata.duration_seconds if len(segment) == 1: # Single frame at specified time actual_timestamps = self._clamp_timestamps_for_decoder([segment[0]], torchcodec_metadata) frame_batch = decoder.get_frames_played_at(actual_timestamps) video_tensor = frame_batch.data sample_count = 1 else: # Segment [start, end) - left-closed, right-open interval start_time, end_time = segment segment_duration = end_time - start_time # Calculate number of frames to sample for this segment target_frames = int(math.ceil(segment_duration * target_video_fps)) target_frames = max(target_frames, min_frames) target_frames = min(target_frames, max_frames) # Generate timestamps for uniform sampling within segment if target_frames == 1: actual_timestamps = [start_time] # Use start_time for single frame else: # Sample uniformly within [start, end), endpoint=False for left-closed right-open actual_timestamps = np.linspace(start_time, end_time, target_frames, endpoint=False).tolist() actual_timestamps = self._clamp_timestamps_for_decoder(actual_timestamps, torchcodec_metadata) # Use multithreading for extraction result = timestamp_decode_with_multithreading(actual_timestamps, self.num_extract_threads, video_path) video_tensor = result["data"] sample_count = len(actual_timestamps) # Create VideoMetadata video_metadata = VideoMetadata( total_num_frames=sample_count, fps=video_fps, duration=duration, video_backend="torchcodec", height=torchcodec_metadata.height, width=torchcodec_metadata.width, frames_indices=None ) # Store actual timestamps as a custom attribute for _calculate_timestamps to use video_metadata.actual_timestamps = actual_timestamps return video_tensor, video_metadata finally: del decoder def fetch_videos( self, video_url_or_urls: Union[str, Dict[str, Any], List[Union[str, Dict[str, Any]]]], sample_indices_fn=None, video_fps: Optional[float] = None, min_frames: Optional[int] = None, max_frames: Optional[int] = None, ): """ Override fetch_videos to use torchcodec for frame extraction. This method uses torchcodec with multithreading for efficient frame extraction. Frame count is calculated by the calculate_num_frames method (fps-based with min/max constraints). Args: video_url_or_urls: Can be one of: - str: Single video path - Dict: Video with segments {"video_path": str, "segments": List[List[float]]} - List[Union[str, Dict]]: List of video paths or segment dicts sample_indices_fn: (Not used) Kept for compatibility with base class signature. video_fps: Target frames per second for video sampling. If None, uses self.video_fps. min_frames: Minimum number of frames to sample. If None, uses self.min_frames. max_frames: Maximum number of frames to sample. If None, uses self.max_frames. Returns: Tuple of (videos, metadata) where videos are torch.Tensors and metadata are VideoMetadata objects. """ # Use provided values or fall back to self defaults effective_video_fps = video_fps if video_fps is not None else self.video_fps effective_min_frames = min_frames if min_frames is not None else self.min_frames effective_max_frames = max_frames if max_frames is not None else self.max_frames # Handle recursive calls for lists if isinstance(video_url_or_urls, list): all_videos = [] all_metadata = [] if len(video_url_or_urls) == 1: per_video_max_frames = [effective_max_frames] else: per_video_max_frames = self._allocate_max_frames_for_multiple_videos( video_url_or_urls, effective_max_frames, ) for x, allocated_max_frames in zip(video_url_or_urls, per_video_max_frames): result = self.fetch_videos( x, video_fps=effective_video_fps, min_frames=effective_min_frames, max_frames=allocated_max_frames, ) # Check if result is from segment expansion (returns lists) or single item if isinstance(result[0], list): all_videos.extend(result[0]) all_metadata.extend(result[1]) else: all_videos.append(result[0]) all_metadata.append(result[1]) return all_videos, all_metadata # Handle dict with segments - returns lists (one per segment) if isinstance(video_url_or_urls, dict): video_path = video_url_or_urls["video_path"] segments = video_url_or_urls["segments"] return self._fetch_video_segments_batched( video_path, segments, min_frames=effective_min_frames, max_frames=effective_max_frames, video_fps=effective_video_fps, ) # Single video path video_path = video_url_or_urls # Clean video streams first (remove extra streams if needed) video_path = cached_clean_video_streams(video_path) decoder = None try: # Create VideoDecoder only once for both metadata and frame extraction decoder = VideoDecoder(video_path, num_ffmpeg_threads=0) torchcodec_metadata = decoder.metadata duration = None if torchcodec_metadata.end_stream_seconds_from_content is not None and torchcodec_metadata.begin_stream_seconds_from_content is not None: duration = torchcodec_metadata.end_stream_seconds_from_content - torchcodec_metadata.begin_stream_seconds_from_content if duration is None or duration <= 0: duration = torchcodec_metadata.duration_seconds # Use num_frames_from_content for accurate frame count (consistent with extraction) total_frames_in_video = torchcodec_metadata.num_frames_from_content # Create VideoMetadata object for sample_frames method temp_metadata = VideoMetadata( total_num_frames=total_frames_in_video, fps=torchcodec_metadata.average_fps, duration=duration, video_backend="torchcodec", height=torchcodec_metadata.height, width=torchcodec_metadata.width, frames_indices=None ) # Use calculate_num_frames method to get the number of frames to sample sample_frames_count = self.calculate_num_frames( temp_metadata, fps=effective_video_fps, min_frames=effective_min_frames, max_frames=effective_max_frames, ) # Ensure sample count is valid effective_sample_count = min(sample_frames_count, total_frames_in_video) if effective_sample_count == 0: raise ValueError(f"Cannot extract frames: video has 0 frames or specified frame count is 0") # Generate uniform frame indices frame_indices = np.linspace(0, total_frames_in_video - 1, effective_sample_count).astype(np.int32) # Ensure indices are valid and remove duplicates frame_indices = np.unique(np.clip(frame_indices, 0, total_frames_in_video - 1)) # Extract frames using multithreading (decoder is created inside each thread for thread safety) result = decode_with_multithreading(frame_indices.tolist(), num_threads=self.num_extract_threads, video_path=video_path) # Extract frame tensor (N, C, H, W) frames_tensor = result["data"] # Create final VideoMetadata object video_metadata = VideoMetadata( total_num_frames=len(frame_indices), fps=torchcodec_metadata.average_fps, duration=duration, video_backend="torchcodec", height=torchcodec_metadata.height, width=torchcodec_metadata.width, frames_indices=frame_indices ) # Ensure frames are in (T, C, H, W) format if frames_tensor.dim() == 4: # (N, C, H, W) video_tensor = frames_tensor else: raise ValueError(f"Unexpected frame tensor shape: {frames_tensor.shape}") return video_tensor, video_metadata except Exception as e: logger.error(f"Error loading video {video_path}: {e}") traceback.print_exc() raise ValueError(f"Failed to load video {video_path}: {e}") finally: if decoder is not None: del decoder def _preprocess( self, videos: list[torch.Tensor], do_convert_rgb: bool = True, do_resize: bool = True, size: Optional[SizeDict] = None, interpolation: PILImageResampling = PILImageResampling.BICUBIC, do_rescale: bool = True, rescale_factor: float = 1 / 255.0, do_normalize: bool = True, image_mean: Optional[Union[float, list[float]]] = None, image_std: Optional[Union[float, list[float]]] = None, patch_size: Optional[int] = None, temporal_patch_size: Optional[int] = None, merge_size: Optional[int] = None, return_tensors: Optional[Union[str, TensorType]] = None, **kwargs, ): grouped_videos, grouped_videos_index = group_videos_by_shape(videos) resized_videos_grouped = {} video_max_pixels = getattr(self, "video_max_pixels", None) if video_max_pixels is not None: total_volume = sum( sv.shape[0] * sv.shape[1] * sv.shape[3] * sv.shape[4] for sv in grouped_videos.values() ) else: total_volume = 0 for shape, stacked_videos in grouped_videos.items(): B, T, C, H, W = stacked_videos.shape num_frames, height, width = T, H, W # Convert to RGB if needed (reuse from base class) if do_convert_rgb: stacked_videos = self.convert_to_rgb(stacked_videos) if do_resize: if video_max_pixels is not None and total_volume > 0: allocated_max_pixels = int(video_max_pixels * (T * H * W) / total_volume) else: allocated_max_pixels = size.longest_edge resized_height, resized_width = smart_resize( num_frames=num_frames, height=height, width=width, temporal_factor=temporal_patch_size, factor=patch_size * merge_size, min_pixels=size.shortest_edge, max_pixels=allocated_max_pixels, per_frame_min_pixels=size.shortest_edge, per_frame_max_pixels=size.longest_edge, ) stacked_videos = stacked_videos.view(B * T, C, H, W) stacked_videos = self.resize( stacked_videos, size=SizeDict(height=resized_height, width=resized_width), interpolation=interpolation, ) stacked_videos = stacked_videos.view(B, T, C, resized_height, resized_width) resized_videos_grouped[shape] = stacked_videos resized_videos = reorder_videos(resized_videos_grouped, grouped_videos_index) # Group videos by size for further processing # Needed in case do_resize is False, or resize returns videos with different sizes grouped_videos, grouped_videos_index = group_videos_by_shape(resized_videos) processed_videos_grouped = {} processed_grids = {} for shape, stacked_videos in grouped_videos.items(): resized_height, resized_width = get_image_size(stacked_videos[0], channel_dim=ChannelDimension.FIRST) # Fused rescale and normalize stacked_videos = self.rescale_and_normalize( stacked_videos, do_rescale, rescale_factor, do_normalize, image_mean, image_std ) patches = stacked_videos # Check that videos have `num_frames` divisible by `temporal_patch_size` if patches.shape[1] % temporal_patch_size != 0: repeats = patches[:, -1:].repeat(1, temporal_patch_size - 1, 1, 1, 1) patches = torch.cat([patches, repeats], dim=1) batch_size, grid_t, channel = patches.shape[:3] grid_t = grid_t // temporal_patch_size grid_h, grid_w = resized_height // patch_size, resized_width // patch_size patches = patches.view( batch_size, grid_t, temporal_patch_size, channel, grid_h // merge_size, merge_size, patch_size, grid_w // merge_size, merge_size, patch_size, ) patches = patches.permute(0, 1, 4, 7, 5, 8, 3, 2, 6, 9) flatten_patches = patches.reshape( batch_size, grid_t * grid_h * grid_w, channel * temporal_patch_size * patch_size * patch_size, ) processed_videos_grouped[shape] = flatten_patches processed_grids[shape] = [[grid_t, grid_h, grid_w]] * batch_size processed_videos = reorder_videos(processed_videos_grouped, grouped_videos_index) processed_grids = reorder_videos(processed_grids, grouped_videos_index) pixel_values_videos = torch.cat(processed_videos, dim=0) video_grid_thw = torch.tensor(processed_grids) data = { "pixel_values_videos": pixel_values_videos, "video_grid_thw": video_grid_thw, } return BatchFeature(data=data, tensor_type=return_tensors) def preprocess( self, videos: Union[str, Dict[str, Any], List[Union[str, Dict[str, Any]]]], **kwargs, ) -> BatchFeature: """ Preprocess videos for the model. This method overrides the base class to handle two video input formats: 1. String path: "path/to/video.mp4" 2. Dict with segments: {"video_path": "...", "segment": [[start, end], [time], ...]} Args: videos: Video input(s) in one of the supported formats. **kwargs: Additional arguments passed to _preprocess. Returns: BatchFeature with pixel_values_videos, video_grid_thw, and optionally video_metadata. """ # Validate kwargs validate_kwargs( captured_kwargs=kwargs.keys(), valid_processor_keys=list(self.valid_kwargs.__annotations__.keys()) + ["return_tensors"], ) # Set default kwargs from self for kwarg_name in self.valid_kwargs.__annotations__: kwargs.setdefault(kwarg_name, getattr(self, kwarg_name, None)) # Pop kwargs that are handled separately return_tensors = kwargs.pop("return_tensors", None) return_metadata = kwargs.pop("return_metadata", False) input_data_format = kwargs.pop("input_data_format", None) device = kwargs.pop("device", None) kwargs.pop("video_metadata", None) # We generate our own metadata kwargs.pop("do_sample_frames", None) # We handle sampling ourselves kwargs.pop("data_format", None) # Not used # Normalize input to list format if not isinstance(videos, list): videos = [videos] # Get video processing params from kwargs (may be passed explicitly for per-batch configuration) video_fps = kwargs.pop("video_fps", None) min_frames = kwargs.pop("min_frames", None) max_frames = kwargs.pop("max_frames", None) # Use fetch_videos to handle both string and dict formats video_tensors, video_metadata = self.fetch_videos( videos, video_fps=video_fps, min_frames=min_frames, max_frames=max_frames, ) # Prepare video tensors using _prepare_input_videos prepared_videos = self._prepare_input_videos( videos=video_tensors, input_data_format=input_data_format, device=device, ) # Process kwargs for _preprocess kwargs = self._further_process_kwargs(**kwargs) self._validate_preprocess_kwargs(**kwargs) # Call _preprocess with prepared videos result = self._preprocess(videos=prepared_videos, return_tensors=return_tensors, **kwargs) # Add metadata if requested if return_metadata: result["video_metadata"] = video_metadata return result __all__ = ["MossVLVideoProcessor"]