code/src/preprocess.py

"""Turn a raw video clip into a grayscale mouth-ROI array of shape (T, H, W).

Mouth localization tries, in order:
  1. MediaPipe legacy Solutions API (``mp.solutions.face_mesh``, mediapipe <= ~0.10.20)
  2. MediaPipe Tasks API (``FaceLandmarker``, modern mediapipe; downloads its
     ~4 MB model file to models/ on first use)
  3. Fixed center crop of the lower-middle of the frame (no mediapipe at all).

The fallback (3) is acceptable for LRW, whose clips are already face-centered,
but landmark cropping matters for your own recordings.
"""
from __future__ import annotations

import os
import urllib.request
from typing import List, Optional

import cv2
import numpy as np

# Mouth landmark indices — identical topology in legacy FaceMesh and the
# Tasks-API FaceLandmarker (478-point mesh).
_MOUTH_IDX = [61, 291, 0, 17, 40, 270, 39, 269, 37, 267, 84, 314, 91, 321]

_MODEL_URL = ("https://storage.googleapis.com/mediapipe-models/face_landmarker/"
              "face_landmarker/float16/1/face_landmarker.task")
_MODEL_PATH = os.environ.get(
    "MP_FACE_MODEL",
    os.path.join(os.path.dirname(os.path.dirname(os.path.abspath(__file__))),
                 "models", "face_landmarker.task"))

# Unified detector: callable(frame_bgr) -> [(x, y) normalized] or None.
_detector = None
_detector_failed = False


def _init_detector() -> None:
    global _detector, _detector_failed
    if _detector is not None or _detector_failed:
        return
    try:
        import mediapipe as mp

        if hasattr(mp, "solutions"):  # ---- legacy Solutions API ----
            mesh = mp.solutions.face_mesh.FaceMesh(
                static_image_mode=False, max_num_faces=1,
                refine_landmarks=True, min_detection_confidence=0.5,
                min_tracking_confidence=0.5)

            def detect(frame_bgr: np.ndarray):
                res = mesh.process(cv2.cvtColor(frame_bgr, cv2.COLOR_BGR2RGB))
                if not res.multi_face_landmarks:
                    return None
                lm = res.multi_face_landmarks[0].landmark
                return [(lm[i].x, lm[i].y) for i in _MOUTH_IDX]

        else:  # ---- modern Tasks API ----
            from mediapipe.tasks import python as mp_tasks
            from mediapipe.tasks.python import vision

            if not os.path.isfile(_MODEL_PATH):
                os.makedirs(os.path.dirname(_MODEL_PATH), exist_ok=True)
                print(f"[preprocess] downloading FaceLandmarker model "
                      f"-> {_MODEL_PATH}")
                urllib.request.urlretrieve(_MODEL_URL, _MODEL_PATH)

            landmarker = vision.FaceLandmarker.create_from_options(
                vision.FaceLandmarkerOptions(
                    base_options=mp_tasks.BaseOptions(
                        model_asset_path=_MODEL_PATH),
                    running_mode=vision.RunningMode.IMAGE, num_faces=1))

            def detect(frame_bgr: np.ndarray):
                img = mp.Image(image_format=mp.ImageFormat.SRGB,
                               data=cv2.cvtColor(frame_bgr, cv2.COLOR_BGR2RGB))
                res = landmarker.detect(img)
                if not res.face_landmarks:
                    return None
                lm = res.face_landmarks[0]
                return [(lm[i].x, lm[i].y) for i in _MOUTH_IDX]

        _detector = detect
    except Exception as e:  # pragma: no cover - environment dependent
        print(f"[preprocess] mediapipe unavailable "
              f"({type(e).__name__}: {e}) — using center-crop fallback")
        _detector_failed = True


def _read_frames(video_path: str) -> List[np.ndarray]:
    cap = cv2.VideoCapture(video_path)
    frames = []
    while True:
        ok, frame = cap.read()
        if not ok:
            break
        frames.append(frame)  # BGR
    cap.release()
    if not frames:
        raise RuntimeError(f"Could not read any frames from {video_path}")
    return frames


def _mouth_box(frame_bgr: np.ndarray, pad: float = 0.6):
    _init_detector()
    if _detector is None:
        return None
    pts = _detector(frame_bgr)
    if pts is None:
        return None
    h, w = frame_bgr.shape[:2]
    xs = [x * w for x, _ in pts]
    ys = [y * h for _, y in pts]
    cx, cy = np.mean(xs), np.mean(ys)
    half = max(max(xs) - min(xs), max(ys) - min(ys)) * (0.5 + pad)
    return cx, cy, half


def _crop(frame_bgr: np.ndarray, box, out_size: int) -> np.ndarray:
    h, w = frame_bgr.shape[:2]
    if box is None:  # center-crop fallback (lower-middle of frame)
        side = min(h, w) // 2
        cx, cy = w // 2, int(h * 0.62)
    else:
        cx, cy, half = box
        side = int(half)
    x0 = int(np.clip(cx - side, 0, w - 1))
    y0 = int(np.clip(cy - side, 0, h - 1))
    x1 = int(np.clip(cx + side, 1, w))
    y1 = int(np.clip(cy + side, 1, h))
    roi = frame_bgr[y0:y1, x0:x1]
    if roi.size == 0:
        roi = frame_bgr
    gray = cv2.cvtColor(roi, cv2.COLOR_BGR2GRAY)
    return cv2.resize(gray, (out_size, out_size), interpolation=cv2.INTER_AREA)


def _resample_time(frames: List[np.ndarray], n: int) -> List[np.ndarray]:
    if len(frames) == n:
        return frames
    idx = np.linspace(0, len(frames) - 1, n).round().astype(int)
    return [frames[i] for i in idx]


def extract_mouth_clip(video_path: str, out_size: int = 96,
                       num_frames: Optional[int] = 29) -> np.ndarray:
    """Return a float32 array (T, out_size, out_size) with values in [0, 1]."""
    frames = _read_frames(video_path)
    # Detect the mouth box once per clip (first frame) for temporal stability;
    # fall back per-frame if the first-frame detection fails.
    box = _mouth_box(frames[0])
    crops = []
    for f in frames:
        b = box if box is not None else _mouth_box(f)
        crops.append(_crop(f, b, out_size))
    if num_frames is not None:
        crops = _resample_time(crops, num_frames)
    clip = np.stack(crops).astype(np.float32) / 255.0
    return clip