File size: 2,503 Bytes
58034d3 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 | """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)
|