import torch import cv2 import numpy as np from pathlib import Path from collections.abc import Iterator from tqdm import tqdm from enum import Enum class ClipSamplingStrategy(Enum): """Clip sampling strategies for FVD evaluation.""" BEGINNING = 'beginning' # Take first N frames (most common) RANDOM = 'random' # Random N consecutive frames UNIFORM = 'uniform' # Uniformly spaced frames across video MIDDLE = 'middle' # Middle N frames SLIDING = 'sliding' # Multiple sliding windows ALL = 'all' # All possible clips def _load_video_cv2(video_path: str | Path, num_frames: int | None = 16, sample_strategy: str = 'uniform') -> torch.Tensor: """ Load video from video file using OpenCV. Args: video_path: Path to video file (MP4, AVI, MOV, MKV) num_frames: Number of frames to extract sample_strategy: 'uniform' or 'random' Returns: video: [T, C, H, W] """ video_path = str(video_path) cap = cv2.VideoCapture(video_path) if not cap.isOpened(): raise RuntimeError(f"Cannot open video: {video_path}") frames = [] total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT)) if num_frames is None: # Read all available frames while True: ret, frame = cap.read() if not ret: break frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB) frames.append(frame) cap.release() if len(frames) == 0: raise RuntimeError(f"Video has 0 frames: {video_path}") frames = np.stack(frames) # [T, H, W, C] frames = torch.from_numpy(frames).permute(0, 3, 1, 2).float() # [T, C, H, W] return frames if total_frames == 0: raise RuntimeError(f"Video has 0 frames: {video_path}") # Determine frame indices for sampling if total_frames < num_frames: frame_indices = list(range(total_frames)) + [total_frames - 1] * (num_frames - total_frames) elif sample_strategy == 'uniform': frame_indices = np.linspace(0, total_frames - 1, num_frames, dtype=int).tolist() elif sample_strategy == 'random': frame_indices = sorted(np.random.choice(total_frames, num_frames, replace=False)) else: raise ValueError(f"Unknown sample_strategy: {sample_strategy}") # Extract frames for idx in frame_indices: cap.set(cv2.CAP_PROP_POS_FRAMES, idx) ret, frame = cap.read() if not ret: if len(frames) > 0: frames.append(frames[-1].copy()) else: h = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT)) w = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH)) frames.append(np.zeros((h, w, 3), dtype=np.uint8)) continue frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB) frames.append(frame) cap.release() frames = np.stack(frames) # [T, H, W, C] frames = torch.from_numpy(frames).permute(0, 3, 1, 2).float() # [T, C, H, W] return frames def _load_video_from_frames(frame_dir: str | Path, num_frames: int | None = 16, sample_strategy: str = 'uniform', frame_extensions: list[str] | None = None) -> torch.Tensor: """ Load video from directory of frame images. Args: frame_dir: Directory containing frames num_frames: Number of frames to sample sample_strategy: 'uniform' or 'random' frame_extensions: Image file extensions to look for Returns: video: [T, C, H, W] """ if frame_extensions is None: frame_extensions = ['.jpg', '.png', '.jpeg', '.bmp'] frame_dir = Path(frame_dir) if not frame_dir.exists(): raise FileNotFoundError(f"Frame directory not found: {frame_dir}") # Find all frames frame_files: list[Path] = [] for ext in frame_extensions: frame_files.extend(frame_dir.glob(f"*{ext}")) if len(frame_files) == 0: raise ValueError(f"No frames found in {frame_dir} with extensions {frame_extensions}") frame_files = sorted(frame_files, key=lambda x: x.name) total_frames = len(frame_files) # Determine frame indices if num_frames is None: frame_indices = list(range(total_frames)) else: if total_frames < num_frames: frame_indices = list(range(total_frames)) + [total_frames - 1] * (num_frames - total_frames) elif sample_strategy == 'uniform': frame_indices = np.linspace(0, total_frames - 1, num_frames, dtype=int).tolist() elif sample_strategy == 'random': frame_indices = sorted(np.random.choice(total_frames, num_frames, replace=False)) else: raise ValueError(f"Unknown sample_strategy: {sample_strategy}") # Load frames frames = [] for idx in frame_indices: frame_path = frame_files[idx] frame = cv2.imread(str(frame_path)) if frame is None: raise RuntimeError(f"Failed to load frame: {frame_path}") frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB) frames.append(frame) # Stack and convert to tensor frames = np.stack(frames) # [T, H, W, C] frames = torch.from_numpy(frames).permute(0, 3, 1, 2).float() # [T, C, H, W] return frames def _detect_video_format(path: str | Path) -> str: """ Detect if path is a video file or frame directory. Returns: 'video_file', 'frame_directory', or 'unknown' """ path = Path(path) if path.is_file(): return 'video_file' elif path.is_dir(): # Check if contains image files image_extensions = ['.jpg', '.jpeg', '.png', '.bmp'] for ext in image_extensions: if list(path.glob(f"*{ext}")): return 'frame_directory' return 'unknown' else: raise ValueError(f"Path does not exist: {path}") def load_video_auto(video_path: str | Path, num_frames: int | None = 16, sample_strategy: str = 'uniform') -> torch.Tensor: """ Automatically detect format and load video. Supports: - Video files (MP4, AVI, MOV, MKV) - Frame directories (JPG, PNG) Args: video_path: Path to video file or frame directory num_frames: Number of frames to extract sample_strategy: 'uniform' or 'random' Returns: video: [T, C, H, W] """ format_type = _detect_video_format(video_path) if format_type == 'video_file': return _load_video_cv2(video_path, num_frames, sample_strategy) elif format_type == 'frame_directory': return _load_video_from_frames(video_path, num_frames, sample_strategy) else: raise ValueError(f"Unknown video format at {video_path}") def sample_clips_from_video(video: torch.Tensor, num_frames_per_clip: int = 16, num_clips: int = 1, strategy: str | ClipSamplingStrategy = ClipSamplingStrategy.BEGINNING, frame_stride: int = 1, temporal_stride: int = 1) -> list[torch.Tensor]: """ Sample clips from a video with various strategies. Args: video: [T, C, H, W] full video num_frames_per_clip: Frames per clip num_clips: Number of clips to extract strategy: ClipSamplingStrategy or string ('beginning', 'random', etc.) frame_stride: Skip frames (FPS control: 1=all, 2=every 2nd, 8=every 8th) temporal_stride: Stride between clips for sliding window Returns: List of clips, each [num_frames_per_clip, C, H, W] Examples: >>> # Beginning clip (most common for FVD) >>> clips = sample_clips_from_video(video, 16, strategy='beginning') >>> # Multiple random clips >>> clips = sample_clips_from_video(video, 16, num_clips=4, strategy='random') >>> # Subsample FPS by 2x (every 2nd frame) >>> clips = sample_clips_from_video(video, 16, frame_stride=2) >>> # Sliding window with overlap >>> clips = sample_clips_from_video(video, 16, strategy='sliding', temporal_stride=8) """ # Convert string to enum if needed if isinstance(strategy, str): strategy = ClipSamplingStrategy(strategy) T, C, H, W = video.shape # Apply frame stride (FPS subsampling) if frame_stride > 1: video = video[::frame_stride] T = len(video) effective_clip_length = num_frames_per_clip # Handle videos shorter than clip length if effective_clip_length > T: pad_length = effective_clip_length - T last_frame = video[-1:].repeat(pad_length, 1, 1, 1) video = torch.cat([video, last_frame], dim=0) T = len(video) clips = [] if strategy == ClipSamplingStrategy.BEGINNING: # Take first clip (most common for FVD evaluation) clip = video[:effective_clip_length] clips.append(clip) elif strategy == ClipSamplingStrategy.MIDDLE: # Take middle clip start = (T - effective_clip_length) // 2 clip = video[start:start + effective_clip_length] clips.append(clip) elif strategy == ClipSamplingStrategy.RANDOM: # Sample N random clips for _ in range(num_clips): start = 0 if effective_clip_length == T else np.random.randint(0, T - effective_clip_length + 1) clip = video[start:start + effective_clip_length] clips.append(clip) elif strategy == ClipSamplingStrategy.UNIFORM: # Uniformly spaced clips if num_clips == 1: # Single clip from middle start = (T - effective_clip_length) // 2 clip = video[start:start + effective_clip_length] clips.append(clip) else: # Multiple uniformly spaced clips step = (T - effective_clip_length) / (num_clips - 1) if num_clips > 1 else 0 for i in range(num_clips): start = int(i * step) start = min(start, T - effective_clip_length) clip = video[start:start + effective_clip_length] clips.append(clip) elif strategy == ClipSamplingStrategy.SLIDING: # Sliding window with stride for start in range(0, T - effective_clip_length + 1, temporal_stride): clip = video[start:start + effective_clip_length] clips.append(clip) if len(clips) >= num_clips: break elif strategy == ClipSamplingStrategy.ALL: # All possible clips (overlapping) for start in range(T - effective_clip_length + 1): clip = video[start:start + effective_clip_length] clips.append(clip) else: raise ValueError(f"Unknown strategy: {strategy}") return clips def load_video_clips_streaming(directory: str | Path, num_frames: int = 16, max_videos: int | None = None, clip_strategy: str | ClipSamplingStrategy = 'beginning', frame_stride: int = 1, num_clips_per_video: int = 1, video_extensions: list[str] | None = None, support_frame_dirs: bool = True, target_size: tuple[int, int] | None = (224, 224), verbose: bool = True) -> Iterator[torch.Tensor]: """ This generator yields clips one-by-one instead of loading all videos into RAM. Perfect for large datasets where memory is limited. Args: directory: Path to directory with videos num_frames: Frames per clip max_videos: Max videos to load clip_strategy: 'beginning', 'random', 'uniform', etc. frame_stride: Frame skip (1=all, 2=every 2nd, 8=every 8th) num_clips_per_video: Number of clips per video video_extensions: Video file extensions support_frame_dirs: Also load frame directories target_size: Resize clips to (H, W). If None, keep original size. verbose: Show progress Yields: clip: [T, C, H, W] individual clips Example: >>> for clip in load_video_clips_streaming('data/videos/', num_frames=16): >>> features = model.extract_features(clip.unsqueeze(0)) >>> # Process one clip at a time - low memory usage! """ if video_extensions is None: video_extensions = ['.mp4', '.avi', '.mov', '.mkv'] directory = Path(directory) if not directory.exists(): raise FileNotFoundError(f"Directory not found: {directory}") # Find video paths video_paths: list[Path] = [] # Find video files for ext in video_extensions: video_paths.extend(directory.glob(f"**/*{ext}")) # Find frame directories if enabled if support_frame_dirs: for subdir in directory.iterdir(): if subdir.is_dir(): # Check if it contains frames image_extensions = ['.jpg', '.jpeg', '.png', '.bmp'] for ext in image_extensions: if list(subdir.glob(f"*{ext}")): video_paths.append(subdir) break if len(video_paths) == 0: raise ValueError(f"No videos found in {directory}") video_paths = sorted(video_paths) if max_videos is not None: video_paths = video_paths[:max_videos] if verbose: print(f"Found {len(video_paths)} videos in {directory}") if num_clips_per_video > 1: print(f"Extracting {num_clips_per_video} clips per video...") if frame_stride > 1: print(f"Subsampling frames with stride {frame_stride}...") if target_size: print(f"Resizing clips to {target_size}...") # Track statistics failed_count = 0 total_clips = 0 iterator = tqdm(video_paths, desc="Loading videos") if verbose else video_paths for video_path in iterator: try: # Load full video video = load_video_auto(video_path, num_frames=None, sample_strategy='uniform') # Sample clips from video clips = sample_clips_from_video(video, num_frames_per_clip=num_frames, num_clips=num_clips_per_video, strategy=clip_strategy, frame_stride=frame_stride) if target_size is not None: resized_clips = [] for clip in clips: T, C, H, W = clip.shape if target_size != (H, W): # Resize to target size clip = clip.contiguous() # Fix non-contiguous tensors first clip_flat = clip.view(T * C, H, W).unsqueeze(0) # [1, T*C, H, W] clip_resized = torch.nn.functional.interpolate(clip_flat, size=target_size, mode='bilinear', align_corners=False) clip = clip_resized.squeeze(0).view(T, C, target_size[0], target_size[1]) # Back to [T, C, H, W] resized_clips.append(clip) clips = resized_clips # Yield clips one by one for clip in clips: yield clip total_clips += 1 # Free memory del video, clips except Exception as e: failed_count += 1 if verbose: print(f"\nWarning: Failed to load {video_path}: {e}") continue # Validate if total_clips == 0: raise RuntimeError(f"Failed to load any videos from {directory}") failure_rate = failed_count / len(video_paths) if failure_rate > 0.1: # More than 10% failed print(f"\nWARNING: {failure_rate:.1%} of videos failed to load ({failed_count}/{len(video_paths)})") if verbose: print(f"\nSuccessfully loaded {total_clips} clips from {len(video_paths) - failed_count} videos")