src/data/preprocessing.py

"""
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)
    }