""" MultiSense-DF — Preprocessing Pipeline Face detection, frame extraction, audio processing, mouth-crop extraction """ import cv2 import torch import numpy as np import librosa import soundfile as sf from pathlib import Path from torchvision import transforms # ── Image Transforms ───────────────────────────────────────────────────────── VISUAL_TRAIN_TRANSFORM = transforms.Compose([ transforms.ToPILImage(), transforms.RandomHorizontalFlip(), transforms.ColorJitter(brightness=0.2, contrast=0.2, saturation=0.1), transforms.RandomRotation(10), transforms.Resize((224, 224)), transforms.ToTensor(), transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]), ]) VISUAL_VAL_TRANSFORM = transforms.Compose([ transforms.ToPILImage(), transforms.Resize((224, 224)), transforms.ToTensor(), transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]), ]) MOUTH_TRANSFORM = transforms.Compose([ transforms.ToPILImage(), transforms.Resize((96, 96)), transforms.ToTensor(), transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5]), ]) # ── Video Processing ────────────────────────────────────────────────────────── def extract_frames(video_path: str, num_frames: int = 125, transform=None) -> torch.Tensor: """ Uniformly sample `num_frames` frames from a video. Returns tensor of shape (num_frames, 3, 224, 224). """ cap = cv2.VideoCapture(str(video_path)) total = int(cap.get(cv2.CAP_PROP_FRAME_COUNT)) indices = np.linspace(0, total - 1, num_frames, dtype=int) frames = [] for idx in indices: cap.set(cv2.CAP_PROP_POS_FRAMES, idx) ret, frame = cap.read() if not ret: # Pad with last frame if read fails frame = frames[-1].numpy().transpose(1, 2, 0) if frames else np.zeros((224, 224, 3), dtype=np.uint8) else: frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB) tf = transform or VISUAL_VAL_TRANSFORM frames.append(tf(frame)) cap.release() return torch.stack(frames) # (T, 3, 224, 224) def extract_audio_waveform(video_path: str, sr: int = 16000, duration: float = 5.0) -> torch.Tensor: """ Extract audio from video, resample to `sr` Hz, normalise. Returns tensor of shape (sr * duration,). """ import subprocess, tempfile, os with tempfile.NamedTemporaryFile(suffix='.wav', delete=False) as f: tmp_path = f.name try: subprocess.run([ 'ffmpeg', '-y', '-i', str(video_path), '-ar', str(sr), '-ac', '1', '-t', str(duration), tmp_path ], stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL, check=True) wav, _ = librosa.load(tmp_path, sr=sr) target_len = int(sr * duration) if len(wav) < target_len: wav = np.pad(wav, (0, target_len - len(wav))) else: wav = wav[:target_len] wav = (wav - wav.mean()) / (wav.std() + 1e-8) return torch.from_numpy(wav).float() finally: os.remove(tmp_path) def extract_mel_spectrogram(waveform: torch.Tensor, sr: int = 16000, n_mels: int = 80, frame_duration: float = 5.0, num_frames: int = 125) -> torch.Tensor: """ Compute mel-spectrogram and split into per-frame windows. Returns (num_frames, 1, 80, W) where W is the window width in mel bins. """ wav_np = waveform.numpy() mel = librosa.feature.melspectrogram(y=wav_np, sr=sr, n_mels=n_mels, hop_length=160, n_fft=512) log_mel = librosa.power_to_db(mel, ref=np.max) # (80, T_mel) # Normalise log_mel = (log_mel - log_mel.mean()) / (log_mel.std() + 1e-8) # Split into num_frames windows total_mel_cols = log_mel.shape[1] window_w = max(1, total_mel_cols // num_frames) mel_frames = [] for i in range(num_frames): start = i * window_w end = start + window_w window = log_mel[:, start:end] if end <= total_mel_cols \ else log_mel[:, -window_w:] mel_frames.append(torch.from_numpy(window).float().unsqueeze(0)) # (1, 80, W) return torch.stack(mel_frames) # (T, 1, 80, W) def extract_mouth_crops(video_path: str, num_frames: int = 125, transform=None) -> torch.Tensor: """ Extract mouth-region crops from uniformly sampled frames. Uses a simple face detector + landmark estimation. Falls back to centre-crop if no face found. Returns (num_frames, 3, 96, 96). """ cap = cv2.VideoCapture(str(video_path)) total = int(cap.get(cv2.CAP_PROP_FRAME_COUNT)) indices = np.linspace(0, total - 1, num_frames, dtype=int) # Attempt to use dlib or mediapipe; fall back to centre-crop try: import mediapipe as mp face_mesh = mp.solutions.face_mesh.FaceMesh( static_image_mode=True, max_num_faces=1 ) use_landmarks = True except ImportError: use_landmarks = False crops = [] tf = transform or MOUTH_TRANSFORM for idx in indices: cap.set(cv2.CAP_PROP_POS_FRAMES, idx) ret, frame = cap.read() if not ret: crops.append(torch.zeros(3, 96, 96)) continue rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB) H, W = rgb.shape[:2] crop = None if use_landmarks: results = face_mesh.process(rgb) if results.multi_face_landmarks: lm = results.multi_face_landmarks[0].landmark # Approximate mouth centre from MediaPipe landmarks mouth_y = int(np.mean([lm[13].y, lm[14].y]) * H) mouth_x = int(np.mean([lm[61].x, lm[291].x]) * W) half = int(min(H, W) * 0.15) y1, y2 = max(0, mouth_y - half), min(H, mouth_y + half) x1, x2 = max(0, mouth_x - half), min(W, mouth_x + half) crop = rgb[y1:y2, x1:x2] if crop is None or crop.size == 0: # Fallback: lower centre quarter y1, y2 = int(H * 0.55), int(H * 0.85) x1, x2 = int(W * 0.25), int(W * 0.75) crop = rgb[y1:y2, x1:x2] crops.append(tf(crop)) cap.release() return torch.stack(crops) # (T, 3, 96, 96) def preprocess_video(video_path: str, num_frames: int = 125, sr: int = 16000, duration: float = 5.0, split: str = 'train') -> dict: """ Full preprocessing pipeline for a single video. Returns dict with all tensors ready for model input. """ transform = VISUAL_TRAIN_TRANSFORM if split == 'train' else VISUAL_VAL_TRANSFORM frames = extract_frames(video_path, num_frames, transform) waveform = extract_audio_waveform(video_path, sr, duration) mels = extract_mel_spectrogram(waveform, sr, num_frames=num_frames) mouths = extract_mouth_crops(video_path, num_frames) return { 'frames': frames, # (T, 3, 224, 224) 'waveform': waveform, # (sr*duration,) 'mel_specs': mels, # (T, 1, 80, W) 'mouth_crops': mouths, # (T, 3, 96, 96) }