| """FFHQ-512 face alignment, identical to ALADDIN's preprocessing. |
| |
| Matches the face preprocessing of the ALADDIN code (https://github.com/isdevnull/aladdin), |
| which wraps facexlib 0.3.0 ``FaceRestoreHelper``: |
| RetinaFace-R50 detection, largest face only, 5-point similarity transform to the FFHQ 512 |
| template (cv2.LMEDS), bilinear warpAffine with a constant gray border (135, 133, 132). |
| |
| Needs ``pip install facexlib opencv-python``. Without ``model_root`` facexlib downloads |
| ``detection_Resnet50_Final.pth`` from its GitHub release on first use. |
| """ |
|
|
| import cv2 |
| import numpy as np |
|
|
| SIZE = 512 |
| BORDER_VALUE = (135, 133, 132) |
|
|
|
|
| def read_bgr(path): |
| """cv2.imread as ALADDIN does (applies EXIF orientation) -> uint8 BGR [H, W, 3].""" |
| return cv2.imread(str(path), cv2.IMREAD_COLOR) |
|
|
|
|
| def warp(image, affine): |
| """Crop with a known affine (original -> crop), exactly as facexlib's align_warp_face. |
| |
| Args: |
| image: uint8 BGR [H, W, 3] original frame. |
| affine: [2, 3] matrix. |
| |
| Returns: |
| uint8 BGR [SIZE, SIZE, 3]. |
| """ |
| return cv2.warpAffine(image, np.asarray(affine, np.float64), (SIZE, SIZE), |
| borderMode=cv2.BORDER_CONSTANT, borderValue=BORDER_VALUE) |
|
|
|
|
| class FaceAligner: |
| """RetinaFace-R50 + FFHQ-512 five-point alignment (state: the facexlib helper). |
| |
| Attributes: |
| helper: facexlib FaceRestoreHelper, face_size 512, no parsing. |
| min_eye_distance: float, detections with fewer pixels between the eyes are dropped. |
| """ |
|
|
| def __init__(self, model_root=None, device="cuda", min_eye_distance=5.0): |
| import torch |
| from facexlib.utils.face_restoration_helper import FaceRestoreHelper |
|
|
| self.min_eye_distance = min_eye_distance |
| kwargs = {"model_rootpath": str(model_root)} if model_root else {} |
| self.helper = FaceRestoreHelper(upscale_factor=1, face_size=SIZE, det_model="retinaface_resnet50", |
| use_parse=False, device=torch.device(device), **kwargs) |
|
|
| def align(self, image): |
| """Largest face of a BGR frame -> (crop uint8 BGR [SIZE, SIZE, 3], affine [2, 3]) or None.""" |
| h = self.helper |
| h.clean_all() |
| h.read_image(image) |
| h.get_face_landmarks_5(only_keep_largest=True, eye_dist_threshold=self.min_eye_distance) |
| if not h.all_landmarks_5: |
| return None |
| h.align_warp_face(border_mode="constant") |
| return h.cropped_faces[0], np.asarray(h.affine_matrices[0], np.float64) |
|
|
