preprocessing.py

"""
Video pre-processing pipeline for the demo.

This mirrors 01_preprocessing.ipynb conceptually, but we use OpenCV's bundled
YuNet face detector instead of MTCNN because facenet-pytorch is incompatible
with the Python 3.13 + numpy 2.x + torch 2.11 stack on HF Spaces.

Pipeline:
    video file
        -> sample 24 evenly-spaced frames
        -> detect the largest face per frame with YuNet
        -> crop with 30% margin and resize to 224x224
        -> stack into [24, 3, 224, 224] float tensor in [0, 1]
        -> ImageNet-normalize for the model
"""

from typing import List, Optional

import cv2
import numpy as np
import torch
from huggingface_hub import hf_hub_download
from PIL import Image


# Same constants as the training-time preprocessing
FRAMES_PER_CLIP = 24
FACE_CROP_SIZE = 224
FACE_MARGIN_RATIO = 0.3

IMAGENET_MEAN = [0.485, 0.456, 0.406]
IMAGENET_STD = [0.229, 0.224, 0.225]

# YuNet ONNX model. OpenCV hosts the official copy on the Hugging Face Hub.
# Using hf_hub_download keeps the cache logic and Xet handling consistent
# with how we pull the main model checkpoint.
YUNET_HF_REPO_ID = "opencv/face_detection_yunet"
YUNET_HF_FILENAME = "face_detection_yunet_2023mar.onnx"


def build_yunet_face_detector():
    """Create a YuNet face detector. Input size is reset per frame."""
    model_path = hf_hub_download(
        repo_id=YUNET_HF_REPO_ID,
        filename=YUNET_HF_FILENAME,
    )
    detector = cv2.FaceDetectorYN.create(
        model=model_path,
        config="",
        input_size=(320, 320),   # placeholder; setInputSize is called per frame
        score_threshold=0.6,
        nms_threshold=0.3,
        top_k=5000,
    )
    return detector


def sample_evenly_spaced_frame_indices(total_frame_count: int, num_frames_wanted: int) -> List[int]:
    """Return num_frames_wanted indices spread evenly across a video."""
    if total_frame_count <= num_frames_wanted:
        indices = list(range(total_frame_count))
        while len(indices) < num_frames_wanted:
            indices.append(total_frame_count - 1)
        return indices
    return np.linspace(0, total_frame_count - 1, num=num_frames_wanted, dtype=int).tolist()


def read_frames_with_opencv(video_path: str, frame_indices: List[int]) -> Optional[np.ndarray]:
    """Load the requested frames using OpenCV. Returns [N, H, W, 3] uint8 in RGB."""
    video_capture = cv2.VideoCapture(video_path)
    if not video_capture.isOpened():
        return None

    indices_to_grab = set(frame_indices)
    frames_by_index = {}
    current_frame_index = 0

    # We just iterate through frames sequentially. Seeking with CAP_PROP_POS_FRAMES
    # is unreliable on many codecs, so a linear scan is safer.
    max_index_needed = max(frame_indices)
    while current_frame_index <= max_index_needed:
        read_success, frame_bgr = video_capture.read()
        if not read_success:
            break
        if current_frame_index in indices_to_grab:
            # OpenCV gives BGR. Convert to RGB.
            frame_rgb = cv2.cvtColor(frame_bgr, cv2.COLOR_BGR2RGB)
            frames_by_index[current_frame_index] = frame_rgb
        current_frame_index += 1

    video_capture.release()

    if not frames_by_index:
        return None

    # Build the output array in the requested order. If we ran out of frames
    # before hitting all indices, pad with the last one we did get.
    last_available_frame = frames_by_index[max(frames_by_index.keys())]
    ordered_frames = [
        frames_by_index.get(idx, last_available_frame) for idx in frame_indices
    ]
    return np.stack(ordered_frames, axis=0)


def crop_face_with_margin(frame_rgb: np.ndarray, face_box: np.ndarray, margin_ratio: float) -> np.ndarray:
    """Crop with extra margin around a face bounding box [x1, y1, x2, y2]."""
    frame_height, frame_width = frame_rgb.shape[:2]
    x1, y1, x2, y2 = face_box
    box_width = x2 - x1
    box_height = y2 - y1

    margin_x = box_width * margin_ratio
    margin_y = box_height * margin_ratio

    x1_expanded = max(0, int(x1 - margin_x))
    y1_expanded = max(0, int(y1 - margin_y))
    x2_expanded = min(frame_width, int(x2 + margin_x))
    y2_expanded = min(frame_height, int(y2 + margin_y))

    return frame_rgb[y1_expanded:y2_expanded, x1_expanded:x2_expanded]


def resize_to_square(image_rgb: np.ndarray, target_size: int) -> np.ndarray:
    """Resize an RGB image to (target_size, target_size)."""
    pil_image = Image.fromarray(image_rgb)
    pil_resized = pil_image.resize((target_size, target_size), Image.BILINEAR)
    return np.array(pil_resized)


def detect_largest_face_per_frame(frames_rgb_uint8: np.ndarray, face_detector) -> List[Optional[np.ndarray]]:
    """Run YuNet on each frame and return the largest detected box per frame.

    Returns a list of [x1, y1, x2, y2] numpy arrays, with None where no face
    was found.
    """
    largest_box_per_frame = []

    for frame_rgb in frames_rgb_uint8:
        frame_height, frame_width = frame_rgb.shape[:2]
        # YuNet needs the input size set before each detect() call.
        face_detector.setInputSize((frame_width, frame_height))

        # YuNet expects BGR
        frame_bgr = cv2.cvtColor(frame_rgb, cv2.COLOR_RGB2BGR)
        _, detections = face_detector.detect(frame_bgr)

        if detections is None or len(detections) == 0:
            largest_box_per_frame.append(None)
            continue

        # Each detection row is [x, y, w, h, 5 landmark x/y pairs, score] = 15 values.
        # Convert (x, y, w, h) -> (x1, y1, x2, y2) and pick the largest by area.
        boxes_xyxy = []
        for detection in detections:
            x, y, width, height = detection[:4]
            boxes_xyxy.append([x, y, x + width, y + height])

        box_areas = [(b[2] - b[0]) * (b[3] - b[1]) for b in boxes_xyxy]
        largest_index = int(np.argmax(box_areas))
        largest_box_per_frame.append(np.array(boxes_xyxy[largest_index], dtype=np.float32))

    return largest_box_per_frame


def process_video_to_clip_tensor(video_path: str, face_detector, progress_callback=None):
    """
    Run the full preprocessing pipeline on a single uploaded video.

    Returns a dict with:
        - clip_tensor: [24, 3, 224, 224] float32 in [0, 1] (un-normalized for display)
        - clip_normalized: [24, 3, 224, 224] float32, ImageNet-normalized (model input)
        - face_crops_rgb: list of 24 uint8 RGB images of the cropped faces (for display)
        - error: optional error message string

    progress_callback is an optional function(message, fraction) the caller can pass
    in to update a Streamlit progress bar.
    """
    def update_progress(message, fraction):
        if progress_callback is not None:
            progress_callback(message, fraction)

    update_progress("Opening video", 0.05)
    video_capture = cv2.VideoCapture(video_path)
    if not video_capture.isOpened():
        return {"error": "Could not open the video file. Is it a valid MP4?"}
    total_frames = int(video_capture.get(cv2.CAP_PROP_FRAME_COUNT))
    video_capture.release()

    if total_frames <= 0:
        return {"error": "The video appears to have zero frames."}

    update_progress("Sampling frames", 0.15)
    frame_indices = sample_evenly_spaced_frame_indices(total_frames, FRAMES_PER_CLIP)
    sampled_frames = read_frames_with_opencv(video_path, frame_indices)
    if sampled_frames is None or len(sampled_frames) == 0:
        return {"error": "Failed to decode any frames from the video."}

    update_progress("Detecting faces", 0.40)
    face_boxes = detect_largest_face_per_frame(sampled_frames, face_detector)

    update_progress("Cropping and resizing", 0.70)
    cropped_resized_frames = []
    last_valid_box = None
    for frame_rgb, face_box in zip(sampled_frames, face_boxes):
        # If YuNet missed this frame, fall back to the most recent valid box.
        if face_box is None:
            face_box = last_valid_box
        else:
            last_valid_box = face_box

        if face_box is None:
            return {"error": "No face was detected in any of the sampled frames."}

        face_crop = crop_face_with_margin(frame_rgb, face_box, FACE_MARGIN_RATIO)
        if face_crop.size == 0:
            return {"error": "An empty crop was produced. The face box may be invalid."}
        face_crop_resized = resize_to_square(face_crop, FACE_CROP_SIZE)
        cropped_resized_frames.append(face_crop_resized)

    update_progress("Building tensor", 0.90)
    # [T, H, W, C] uint8 -> [T, C, H, W] float32 in [0, 1]
    clip_array = np.stack(cropped_resized_frames, axis=0)
    clip_tensor = torch.from_numpy(clip_array).permute(0, 3, 1, 2).float() / 255.0

    # ImageNet normalization for model input
    mean_tensor = torch.tensor(IMAGENET_MEAN).view(1, 3, 1, 1)
    std_tensor = torch.tensor(IMAGENET_STD).view(1, 3, 1, 1)
    clip_normalized = (clip_tensor - mean_tensor) / std_tensor

    update_progress("Preprocessing done", 1.0)
    return {
        "clip_tensor": clip_tensor,
        "clip_normalized": clip_normalized,
        "face_crops_rgb": cropped_resized_frames,
        "error": None,
    }