45 lines
1.2 KiB
Python
45 lines
1.2 KiB
Python
from __future__ import annotations
|
|
|
|
from typing import Tuple
|
|
|
|
import numpy as np
|
|
|
|
|
|
_DST_5PTS_112 = 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,
|
|
)
|
|
|
|
|
|
def align_face_5pts(img_rgb: np.ndarray, landmarks5: np.ndarray, out_size: Tuple[int, int] = (112, 112)) -> np.ndarray:
|
|
"""Return aligned RGB image (H,W,3) uint8."""
|
|
|
|
import cv2
|
|
|
|
if img_rgb is None or img_rgb.ndim != 3 or img_rgb.shape[2] != 3:
|
|
raise ValueError("img_rgb must be HxWx3")
|
|
lmk = np.asarray(landmarks5, dtype=np.float32)
|
|
if lmk.shape != (5, 2):
|
|
raise ValueError("landmarks5 must be shape (5,2)")
|
|
|
|
dst = _DST_5PTS_112.copy()
|
|
if out_size != (112, 112):
|
|
sx = out_size[0] / 112.0
|
|
sy = out_size[1] / 112.0
|
|
dst[:, 0] *= sx
|
|
dst[:, 1] *= sy
|
|
|
|
M, inliers = cv2.estimateAffinePartial2D(lmk, dst, method=cv2.LMEDS)
|
|
if M is None or not np.isfinite(M).all():
|
|
raise ValueError("estimateAffinePartial2D failed")
|
|
|
|
w, h = out_size
|
|
aligned = cv2.warpAffine(img_rgb, M, (w, h), flags=cv2.INTER_LINEAR, borderMode=cv2.BORDER_CONSTANT, borderValue=0)
|
|
return aligned
|