image_processing_edgeface.py

"""
Image processor for EdgeFace.

Faithful port of the alignment used by the Idiap EdgeFace Space (utils.py):
MediaPipe FaceMesh landmarks -> 5 points -> reflective similarity transform onto
the ArcFace 112x112 template (custom MATLAB cp2tform-style solver).

Works with both MediaPipe backends:
  * "tasks"     -> latest API (mp.tasks.vision.FaceLandmarker + .task bundle)
  * "solutions" -> legacy mp.solutions.face_mesh.FaceMesh (older installs)
The default backend="auto" tries tasks first and falls back to solutions.

Pipeline: (optional) align -> rescale to [0,1] -> normalize mean/std=0.5.
If do_align=False the input is treated as an already-aligned crop and only
resized to image_size.
"""

import os
import weakref
from typing import List, Optional, Union

import numpy as np
from numpy.linalg import inv, lstsq, matrix_rank, norm
from transformers.image_processing_utils import BaseImageProcessor, BatchFeature
from transformers.image_utils import ImageInput, make_list_of_images, to_numpy_array

# ArcFace 5-point reference template for a 112x112 crop.
# order matches the 5 source points: [reye, leye, nose, mouthright, mouthleft]
REFERENCE_FACIAL_POINTS = np.array(
    [
        [38.2946, 51.6963],
        [73.5318, 51.5014],
        [56.0252, 71.7366],
        [41.5493, 92.3655],
        [70.7299, 92.2041],
    ],
    dtype=np.float32,
)

# MediaPipe FaceMesh indices (from the Space's utils.py). Valid for both the
# 468-point legacy mesh and the 478-point tasks mesh (extra points are irises).
IDX_REYE = (362, 263)   # eye on the image-left  (subject's right)
IDX_LEYE = (33, 243)    # eye on the image-right (subject's left)
IDX_NOSE = 1
IDX_MOUTH_RIGHT = 287   # mouth corner on the image-left
IDX_MOUTH_LEFT = 57     # mouth corner on the image-right

# Official Google model bundle for the tasks API.
_TASK_MODEL_URL = (
    "https://storage.googleapis.com/mediapipe-models/face_landmarker/"
    "face_landmarker/float16/1/face_landmarker.task"
)

# Live MediaPipe detectors are not JSON/deepcopy-safe, so keep them off the
# instance __dict__ (which save_pretrained serializes) via a weak cache.
_RUNTIME: "weakref.WeakKeyDictionary" = weakref.WeakKeyDictionary()


# --------------------------------------------------------------------------
# Similarity transform utilities (ported from the Space's utils.py)
# --------------------------------------------------------------------------
def _tformfwd(trans, uv):
    uv_h = np.hstack((uv, np.ones((uv.shape[0], 1))))
    xy = uv_h @ trans
    return xy[:, :-1]


def _find_nonreflective_similarity(uv, xy, K=2):
    M = xy.shape[0]
    x, y = xy[:, 0:1], xy[:, 1:2]
    u, v = uv[:, 0:1], uv[:, 1:2]

    X = np.vstack((
        np.hstack((x, y, np.ones((M, 1)), np.zeros((M, 1)))),
        np.hstack((y, -x, np.zeros((M, 1)), np.ones((M, 1)))),
    ))
    U = np.vstack((u, v))

    if matrix_rank(X) >= 2 * K:
        r, _, _, _ = lstsq(X, U, rcond=None)
    else:
        raise ValueError("cp2tform:twoUniquePointsReq")

    sc, ss, tx, ty = r.flatten()
    Tinv = np.array([[sc, -ss, 0], [ss, sc, 0], [tx, ty, 1]])
    T = inv(Tinv)
    T[:, 2] = [0, 0, 1]
    return T, Tinv


def _find_similarity(uv, xy):
    trans1, trans1_inv = _find_nonreflective_similarity(uv, xy)

    xyR = xy.copy()
    xyR[:, 0] *= -1
    trans2r, _ = _find_nonreflective_similarity(uv, xyR)
    TreflectY = np.array([[-1, 0, 0], [0, 1, 0], [0, 0, 1]])
    trans2 = trans2r @ TreflectY

    norm1 = norm(_tformfwd(trans1, uv) - xy)
    norm2 = norm(_tformfwd(trans2, uv) - xy)
    return (trans1, trans1_inv) if norm1 <= norm2 else (trans2, inv(trans2))


def _get_cv2_affine(src_pts, dst_pts):
    trans, _ = _find_similarity(src_pts, dst_pts)
    return trans[:, :2].T  # 2x3 for cv2.warpAffine


def _warp_and_crop_face(src_img, facial_pts, reference_pts=REFERENCE_FACIAL_POINTS,
                        crop_size=(112, 112), scale=1):
    import cv2

    ref_pts = reference_pts * scale
    ref_pts = ref_pts + (np.mean(reference_pts, axis=0) - np.mean(ref_pts, axis=0))

    src_pts = np.array(facial_pts, dtype=np.float32)
    if src_pts.shape != ref_pts.shape:
        raise ValueError("facial_pts and reference_pts must have the same shape")

    tfm = _get_cv2_affine(src_pts, ref_pts)
    return cv2.warpAffine(src_img, tfm, crop_size)


class EdgeFaceImageProcessor(BaseImageProcessor):
    model_input_names = ["pixel_values"]

    def __init__(
        self,
        do_align: bool = True,
        image_size: int = 112,
        do_rescale: bool = True,
        rescale_factor: float = 1 / 255,
        do_normalize: bool = True,
        image_mean: Optional[List[float]] = None,
        image_std: Optional[List[float]] = None,
        mp_backend: str = "auto",          # "auto" | "tasks" | "solutions"
        mp_model_path: Optional[str] = None,  # path to a .task bundle (tasks backend)
        **kwargs,
    ):
        super().__init__(**kwargs)
        self.do_align = do_align
        self.image_size = image_size
        self.do_rescale = do_rescale
        self.rescale_factor = rescale_factor
        self.do_normalize = do_normalize
        self.image_mean = image_mean if image_mean is not None else [0.5, 0.5, 0.5]
        self.image_std = image_std if image_std is not None else [0.5, 0.5, 0.5]
        self.mp_backend = mp_backend
        self.mp_model_path = mp_model_path

    # -- runtime (non-serialized) cache ------------------------------------
    def _runtime(self):
        d = _RUNTIME.get(self)
        if d is None:
            d = {}
            _RUNTIME[self] = d
        return d

    # -- model bundle for the tasks backend --------------------------------
    def _resolve_model_path(self) -> str:
        if self.mp_model_path:
            return self.mp_model_path
        env = os.environ.get("EDGEFACE_MP_MODEL")
        if env:
            return env
        cache_dir = os.path.join(os.path.expanduser("~"), ".cache", "edgeface")
        os.makedirs(cache_dir, exist_ok=True)
        path = os.path.join(cache_dir, "face_landmarker.task")
        if not os.path.exists(path):
            import urllib.request
            urllib.request.urlretrieve(_TASK_MODEL_URL, path)
        return path

    # -- backend builders: each returns fn(rgb_uint8) -> (N,2) norm or None -
    def _build_tasks_detector(self):
        import mediapipe as mp
        from mediapipe.tasks import python as mp_python
        from mediapipe.tasks.python import vision as mp_vision

        options = mp_vision.FaceLandmarkerOptions(
            base_options=mp_python.BaseOptions(model_asset_path=self._resolve_model_path()),
            running_mode=mp_vision.RunningMode.IMAGE,
            num_faces=1,
        )
        landmarker = mp_vision.FaceLandmarker.create_from_options(options)

        def detect(rgb):
            mp_img = mp.Image(image_format=mp.ImageFormat.SRGB, data=np.ascontiguousarray(rgb))
            res = landmarker.detect(mp_img)
            if not res.face_landmarks:
                return None
            return np.array([[p.x, p.y] for p in res.face_landmarks[0]], dtype=np.float32)

        return detect

    def _build_solutions_detector(self):
        import mediapipe as mp

        face_mesh = mp.solutions.face_mesh.FaceMesh(
            static_image_mode=True, refine_landmarks=True, min_detection_confidence=0.5,
        )

        def detect(rgb):
            res = face_mesh.process(rgb)
            if not res.multi_face_landmarks:
                return None
            return np.array([[p.x, p.y] for p in res.multi_face_landmarks[0].landmark],
                            dtype=np.float32)

        return detect

    def _get_detect_fn(self):
        runtime = self._runtime()
        if "detect_fn" in runtime:
            return runtime["detect_fn"]

        order = {
            "auto": ["tasks", "solutions"],
            "tasks": ["tasks"],
            "solutions": ["solutions"],
        }.get(self.mp_backend)
        if order is None:
            raise ValueError(f"Unknown mp_backend={self.mp_backend!r}")

        errors = []
        for backend in order:
            try:
                fn = (self._build_tasks_detector() if backend == "tasks"
                      else self._build_solutions_detector())
                runtime["detect_fn"] = fn
                return fn
            except Exception as e:  # noqa: BLE001 - try next backend
                errors.append(f"{backend}: {type(e).__name__}: {e}")

        raise ImportError(
            "Could not initialize a MediaPipe face detector. Install mediapipe "
            "(`pip install mediapipe`) and ensure network access for the .task "
            "bundle, or pass do_align=False / precomputed landmarks.\n"
            + "\n".join(errors)
        )

    # -- landmark extraction -----------------------------------------------
    def _detect_landmarks(self, image_rgb: np.ndarray) -> Optional[np.ndarray]:
        """Return the 5 source points in [reye, leye, nose, mouthright, mouthleft] order."""
        h, w = image_rgb.shape[:2]
        norm_pts = self._get_detect_fn()(image_rgb)
        if norm_pts is None or len(norm_pts) <= max(*IDX_REYE, *IDX_LEYE, IDX_MOUTH_RIGHT):
            return None

        px = norm_pts * np.array([w, h], dtype=np.float32)

        reye = (px[IDX_REYE[0]] + px[IDX_REYE[1]]) / 2.0
        leye = (px[IDX_LEYE[0]] + px[IDX_LEYE[1]]) / 2.0
        return np.stack([reye, leye, px[IDX_NOSE], px[IDX_MOUTH_RIGHT], px[IDX_MOUTH_LEFT]]).astype(np.float32)

    def _align_one(self, image_rgb: np.ndarray, landmarks: Optional[np.ndarray]) -> np.ndarray:
        if landmarks is None:
            landmarks = self._detect_landmarks(image_rgb)
        if landmarks is None:
            import cv2  # detection failed -> plain resize so the batch still runs
            return cv2.resize(image_rgb, (self.image_size, self.image_size))
        return _warp_and_crop_face(image_rgb, landmarks, crop_size=(self.image_size, self.image_size))

    # -- main entry point --------------------------------------------------
    def preprocess(
        self,
        images: ImageInput,
        do_align: Optional[bool] = None,
        landmarks: Optional[Union[np.ndarray, List[np.ndarray]]] = None,
        return_tensors: Optional[str] = "pt",
        **kwargs,
    ) -> BatchFeature:
        do_align = self.do_align if do_align is None else do_align
        images = make_list_of_images(images)

        if landmarks is not None and not isinstance(landmarks, list):
            landmarks = [landmarks]

        processed = []
        for i, img in enumerate(images):
            arr = to_numpy_array(img)  # RGB, HxWxC
            if arr.ndim == 2:
                arr = np.stack([arr] * 3, axis=-1)
            if arr.shape[-1] == 4:
                arr = arr[..., :3]
            arr = arr.astype(np.uint8)

            if do_align:
                lmk = landmarks[i] if landmarks is not None else None
                arr = self._align_one(arr, lmk)
            else:
                import cv2
                arr = cv2.resize(arr, (self.image_size, self.image_size))

            arr = arr.astype(np.float32)
            if self.do_rescale:
                arr = arr * self.rescale_factor
            if self.do_normalize:
                arr = (arr - np.array(self.image_mean)) / np.array(self.image_std)

            processed.append(arr.transpose(2, 0, 1))  # CxHxW

        pixel_values = np.stack(processed, axis=0).astype(np.float32)
        return BatchFeature(data={"pixel_values": pixel_values}, tensor_type=return_tensors)