""" Video Processor - Handles video I/O using FFmpeg Supports frame extraction, video assembly, and audio preservation """ import os import cv2 import json import shutil import subprocess import numpy as np from pathlib import Path from typing import Generator, Tuple, Optional, Dict from dataclasses import dataclass @dataclass class VideoInfo: """Video metadata""" width: int height: int fps: float total_frames: int duration: float has_audio: bool codec: str filepath: str class VideoProcessor: """Handles all video I/O operations""" def __init__(self, config): self.config = config self._verify_ffmpeg() def _verify_ffmpeg(self): """Verify FFmpeg is installed""" try: result = subprocess.run( ["ffmpeg", "-version"], capture_output=True, text=True, timeout=5 ) if result.returncode != 0: raise RuntimeError("FFmpeg not working properly") except FileNotFoundError: raise RuntimeError( "FFmpeg not found! Install it:\n" " Ubuntu: sudo apt install ffmpeg\n" " Mac: brew install ffmpeg\n" " Windows: choco install ffmpeg" ) def get_video_info(self, video_path: str) -> VideoInfo: """Extract video metadata using ffprobe""" cmd = [ "ffprobe", "-v", "quiet", "-print_format", "json", "-show_format", "-show_streams", video_path ] result = subprocess.run(cmd, capture_output=True, text=True) if result.returncode != 0: raise ValueError(f"Cannot read video: {video_path}") probe = json.loads(result.stdout) # Find video stream video_stream = None has_audio = False for stream in probe.get("streams", []): if stream["codec_type"] == "video" and video_stream is None: video_stream = stream elif stream["codec_type"] == "audio": has_audio = True if not video_stream: raise ValueError("No video stream found") # Parse FPS fps_parts = video_stream.get("r_frame_rate", "30/1").split("/") fps = float(fps_parts[0]) / float(fps_parts[1]) if len(fps_parts) == 2 else 30.0 # Parse frame count nb_frames = int(video_stream.get("nb_frames", 0)) duration = float(probe.get("format", {}).get("duration", 0)) if nb_frames == 0 and duration > 0: nb_frames = int(duration * fps) return VideoInfo( width=int(video_stream["width"]), height=int(video_stream["height"]), fps=fps, total_frames=nb_frames, duration=duration, has_audio=has_audio, codec=video_stream.get("codec_name", "unknown"), filepath=video_path ) def extract_frames_to_dir( self, video_path: str, output_dir: str, max_height: Optional[int] = None, progress_callback=None ) -> Tuple[str, VideoInfo]: """ Extract all frames as numbered JPEGs (required by SAM 2) Args: video_path: Path to input video output_dir: Directory to save frames max_height: Optional max height for resizing progress_callback: Optional callback(current, total) Returns: (frames_dir, video_info) """ os.makedirs(output_dir, exist_ok=True) info = self.get_video_info(video_path) # Build FFmpeg command cmd = ["ffmpeg", "-y", "-i", video_path] # Add scaling if needed if max_height and info.height > max_height: cmd.extend(["-vf", f"scale=-2:{max_height}"]) # Output as numbered JPEGs cmd.extend([ "-qscale:v", "2", # High quality JPEG os.path.join(output_dir, "%06d.jpg") ]) process = subprocess.Popen( cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE ) _, stderr = process.communicate() if process.returncode != 0: raise RuntimeError(f"Frame extraction failed: {stderr.decode()}") # Count extracted frames frame_files = sorted([ f for f in os.listdir(output_dir) if f.endswith('.jpg') ]) info.total_frames = len(frame_files) # Update dimensions if resized if frame_files: sample = cv2.imread(os.path.join(output_dir, frame_files[0])) if sample is not None: info.height, info.width = sample.shape[:2] print(f"📹 Extracted {info.total_frames} frames ({info.width}x{info.height} @ {info.fps:.1f} FPS)") return output_dir, info def read_frames_generator( self, video_path: str, max_height: Optional[int] = None ) -> Generator[Tuple[int, np.ndarray], None, None]: """ Stream frames from video using OpenCV (memory efficient) Yields: (frame_index, frame_bgr) """ cap = cv2.VideoCapture(video_path) if not cap.isOpened(): raise ValueError(f"Cannot open video: {video_path}") frame_idx = 0 while True: ret, frame = cap.read() if not ret: break # Resize if needed if max_height and frame.shape[0] > max_height: scale = max_height / frame.shape[0] new_w = int(frame.shape[1] * scale) frame = cv2.resize(frame, (new_w, max_height)) yield frame_idx, frame frame_idx += 1 cap.release() def assemble_video( self, frames_dir: str, output_path: str, fps: float, original_video: Optional[str] = None, progress_callback=None ) -> str: """ Assemble processed frames back into video with optional audio Args: frames_dir: Directory with numbered JPEG frames output_path: Output video path fps: Frame rate original_video: Original video to copy audio from progress_callback: Optional callback Returns: Path to output video """ temp_video = output_path + ".temp.mp4" # Step 1: Encode frames to video cmd = [ "ffmpeg", "-y", "-framerate", str(fps), "-i", os.path.join(frames_dir, "%06d.jpg"), "-c:v", self.config.video.output_codec, "-crf", str(self.config.video.output_crf), "-preset", self.config.video.output_preset, "-pix_fmt", self.config.video.pixel_format, "-movflags", "+faststart", ] if original_video: cmd.append(temp_video) else: cmd.append(output_path) process = subprocess.Popen( cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE ) _, stderr = process.communicate() if process.returncode != 0: raise RuntimeError(f"Video encoding failed: {stderr.decode()}") # Step 2: Mux audio from original video if original_video: info = self.get_video_info(original_video) if info.has_audio: mux_cmd = [ "ffmpeg", "-y", "-i", temp_video, "-i", original_video, "-c:v", "copy", "-c:a", "aac", "-map", "0:v:0", "-map", "1:a:0?", "-shortest", output_path ] process = subprocess.Popen( mux_cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE ) _, stderr = process.communicate() if process.returncode != 0: # Fallback: just use video without audio shutil.move(temp_video, output_path) print("⚠️ Audio muxing failed, output has no audio") else: os.remove(temp_video) print("🔊 Audio preserved from original video") else: shutil.move(temp_video, output_path) file_size_mb = os.path.getsize(output_path) / (1024 * 1024) print(f"✅ Output video: {output_path} ({file_size_mb:.1f} MB)") return output_path def cleanup_temp(self, temp_dir: str): """Remove temporary files""" if os.path.exists(temp_dir): shutil.rmtree(temp_dir, ignore_errors=True)