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