Source code for rppg.roi

"""
Face ROI extraction for remote heart-rate estimation (rPPG) — MediaPipe
Tasks FaceLandmarker with output_face_blendshapes=False.

Same model family/file as expression/detector.py, facemask/detectors.py,
and (in the separate python/ project) python/pose/gaze_estimator.py — but
this module constructs its OWN FaceLandmarkerOptions instance rather than
reusing any of them, for the same reason expression/detector.py's own
docstring gives for not reaching into a sibling plugin folder. Owns its own
model cache under analysis/rppg/models/ (gitignored), re-downloading the
same face_landmarker.task file the other plugins already have as an
accepted tradeoff (see expression/detector.py's identical note).

ROI landmark indices
---------------------
The 9-point "lower face" polygon below is a real, cited, working reference
(NOT invented/recalled from memory — no prior art for this existed anywhere
in this codebase, so it was independently verified against a real GitHub
implementation before being written down here, matching this project's own
established discipline of checking algorithmic claims against primary
sources rather than trusting memory): SamProell/yarppg's
FaceMeshDetector._lower_face (src/yarppg/roi/facemesh_segmenter.py), itself
attributed to:

    X. Li, J. Chen, G. Zhao, and M. Pietikainen, "Remote Heart Rate
    Measurement From Face Videos Under Realistic Situations", CVPR 2014.
    https://doi.org/10.1109/CVPR.2014.543

This single polygon (cheeks + the area between them, avoiding eyes,
eyebrows, and the mouth) was chosen over a separate forehead+cheek pair
specifically because every index in it traces to a real cited source —
a forehead-specific index set could not be independently verified the same
way and was deliberately not invented from partial memory.
"""

from __future__ import annotations

import urllib.request
from dataclasses import dataclass
from pathlib import Path

import cv2
import numpy as np

_MODELS_DIR = Path(__file__).parent / "models"

_MEDIAPIPE_MODEL_URL = (
    "https://storage.googleapis.com/mediapipe-models/"
    "face_landmarker/face_landmarker/float16/1/face_landmarker.task"
)

#: MediaPipe FaceMesh landmark indices bounding the lower-face skin region
#: (both cheeks + the area between them) used for rPPG signal extraction —
#: see the module docstring for the verified source/attribution. Order
#: matters: these are used directly as a polygon contour, not sorted.
LOWER_FACE_LANDMARK_INDICES: list[int] = [200, 431, 411, 340, 349, 120, 111, 187, 211]




@dataclass
class FaceRoiSample:
    """One frame's detected face ROI and its mean RGB color.

    Attributes
    ----------
    roi_bbox_px : tuple of int
        ``(x, y, w, h)`` pixel bounding box of the sampled ROI polygon —
        kept for the debug overlay, not used by the signal-processing math.
    rgb_mean : tuple of float
        Mean ``(R, G, B)`` pixel value inside the ROI polygon, in ``[0,255]``.
    """

    roi_bbox_px: tuple[int, int, int, int]
    rgb_mean: tuple[float, float, float]





class MediaPipeFaceRoiExtractor:
    """MediaPipe Tasks ``FaceLandmarker``-based skin-ROI RGB sampler.

    Parameters
    ----------
    min_confidence : float, default 0.5
        Minimum face-detection/-presence confidence to keep a detection.
    """

    def __init__(self, min_confidence: float = 0.5) -> None:
        import mediapipe as mp
        from mediapipe.tasks import python as mp_python
        from mediapipe.tasks.python import vision as mp_vision

        model_path = _ensure_download(_MODELS_DIR / "face_landmarker.task", _MEDIAPIPE_MODEL_URL)

        options = mp_vision.FaceLandmarkerOptions(
            base_options=mp_python.BaseOptions(model_asset_path=str(model_path)),
            running_mode=mp_vision.RunningMode.IMAGE,
            num_faces=1,  # rPPG only ever tracks one subject's pulse per camera
            min_face_detection_confidence=min_confidence,
            min_face_presence_confidence=min_confidence,
            output_face_blendshapes=False,
            output_facial_transformation_matrixes=False,
        )
        self._mp = mp
        self._landmarker = mp_vision.FaceLandmarker.create_from_options(options)



    def extract(self, frame_bgr: np.ndarray) -> FaceRoiSample | None:
        """Detect a face and sample its lower-face ROI's mean RGB color.

        Parameters
        ----------
        frame_bgr : numpy.ndarray
            BGR frame, as returned by ``cv2.VideoCapture``.

        Returns
        -------
        FaceRoiSample or None
            ``None`` if no face was detected this frame — callers must treat
            this as a real gap, not fabricate a sample (matches this
            codebase's established "skip missing samples" discipline, e.g.
            pose_kinematics.hpp).
        """
        h, w = frame_bgr.shape[:2]
        rgb = cv2.cvtColor(frame_bgr, cv2.COLOR_BGR2RGB)
        mp_image = self._mp.Image(image_format=self._mp.ImageFormat.SRGB, data=rgb)
        result = self._landmarker.detect(mp_image)

        if not result.face_landmarks:
            return None

        landmarks = result.face_landmarks[0]
        pts = np.array(
            [(landmarks[i].x * w, landmarks[i].y * h) for i in LOWER_FACE_LANDMARK_INDICES],
            dtype=np.int32,
        )

        mask = np.zeros((h, w), dtype=np.uint8)
        cv2.fillPoly(mask, [pts], 255)
        mean_bgr = cv2.mean(frame_bgr, mask=mask)  # (B, G, R, alpha)
        rgb_mean = (mean_bgr[2], mean_bgr[1], mean_bgr[0])

        x, y, bw, bh = cv2.boundingRect(pts)
        return FaceRoiSample(roi_bbox_px=(x, y, bw, bh), rgb_mean=rgb_mean)




def _ensure_download(dest: Path, url: str) -> Path:
    if not dest.exists():
        dest.parent.mkdir(parents=True, exist_ok=True)
        print(f"[rppg] Downloading {dest.name} …", flush=True)
        urllib.request.urlretrieve(url, dest)
        print(f"[rppg] Downloaded to {dest}", flush=True)
    return dest