src/FaceAlign/face_align.py

import cv2
import numpy as np
import torch
from skimage import transform as skt
from typing import Iterable, Tuple

src1 = np.array(
    [
        [51.642, 50.115],
        [57.617, 49.990],
        [35.740, 69.007],
        [51.157, 89.050],
        [57.025, 89.702],
    ],
    dtype=np.float32,
)
# <--left
src2 = np.array(
    [
        [45.031, 50.118],
        [65.568, 50.872],
        [39.677, 68.111],
        [45.177, 86.190],
        [64.246, 86.758],
    ],
    dtype=np.float32,
)

# ---frontal
src3 = np.array(
    [
        [39.730, 51.138],
        [72.270, 51.138],
        [56.000, 68.493],
        [42.463, 87.010],
        [69.537, 87.010],
    ],
    dtype=np.float32,
)

# -->right
src4 = np.array(
    [
        [46.845, 50.872],
        [67.382, 50.118],
        [72.737, 68.111],
        [48.167, 86.758],
        [67.236, 86.190],
    ],
    dtype=np.float32,
)

# -->right profile
src5 = np.array(
    [
        [54.796, 49.990],
        [60.771, 50.115],
        [76.673, 69.007],
        [55.388, 89.702],
        [61.257, 89.050],
    ],
    dtype=np.float32,
)

src = np.array([src1, src2, src3, src4, src5])
src_map = src

ffhq_src = np.array(
    [
        [192.98138, 239.94708],
        [318.90277, 240.1936],
        [256.63416, 314.01935],
        [201.26117, 371.41043],
        [313.08905, 371.15118],
    ]
)
ffhq_src = np.expand_dims(ffhq_src, axis=0)


# arcface_src = 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)

# arcface_src = np.expand_dims(arcface_src, axis=0)

# In[66]:


# lmk is prediction; src is template
def estimate_norm(lmk, image_size=112, mode="ffhq"):
    assert lmk.shape == (5, 2)
    tform = skt.SimilarityTransform()
    lmk_tran = np.insert(lmk, 2, values=np.ones(5), axis=1)
    min_M = []
    min_index = []
    min_error = float("inf")
    if mode == "ffhq":
        # assert image_size == 112
        src = ffhq_src * image_size / 512
    else:
        src = src_map * image_size / 112
    for i in np.arange(src.shape[0]):
        tform.estimate(lmk, src[i])
        M = tform.params[0:2, :]
        results = np.dot(M, lmk_tran.T)
        results = results.T
        error = np.sum(np.sqrt(np.sum((results - src[i]) ** 2, axis=1)))
        if error < min_error:
            min_error = error
            min_M = M
            min_index = i
    return min_M, min_index


def norm_crop(img, landmark, image_size=112, mode="ffhq"):
    if mode == "Both":
        M_None, _ = estimate_norm(landmark, image_size, mode="newarc")
        M_ffhq, _ = estimate_norm(landmark, image_size, mode="ffhq")
        warped_None = cv2.warpAffine(
            img, M_None, (image_size, image_size), borderValue=0.0
        )
        warped_ffhq = cv2.warpAffine(
            img, M_ffhq, (image_size, image_size), borderValue=0.0
        )
        return warped_ffhq, warped_None
    else:
        M, pose_index = estimate_norm(landmark, image_size, mode)
        warped = cv2.warpAffine(img, M, (image_size, image_size), borderValue=0.0)
        return warped


def square_crop(im, S):
    if im.shape[0] > im.shape[1]:
        height = S
        width = int(float(im.shape[1]) / im.shape[0] * S)
        scale = float(S) / im.shape[0]
    else:
        width = S
        height = int(float(im.shape[0]) / im.shape[1] * S)
        scale = float(S) / im.shape[1]
    resized_im = cv2.resize(im, (width, height))
    det_im = np.zeros((S, S, 3), dtype=np.uint8)
    det_im[: resized_im.shape[0], : resized_im.shape[1], :] = resized_im
    return det_im, scale


def transform(data, center, output_size, scale, rotation):
    scale_ratio = scale
    rot = float(rotation) * np.pi / 180.0
    # translation = (output_size/2-center[0]*scale_ratio, output_size/2-center[1]*scale_ratio)
    t1 = skt.SimilarityTransform(scale=scale_ratio)
    cx = center[0] * scale_ratio
    cy = center[1] * scale_ratio
    t2 = skt.SimilarityTransform(translation=(-1 * cx, -1 * cy))
    t3 = skt.SimilarityTransform(rotation=rot)
    t4 = skt.SimilarityTransform(translation=(output_size / 2, output_size / 2))
    t = t1 + t2 + t3 + t4
    M = t.params[0:2]
    cropped = cv2.warpAffine(data, M, (output_size, output_size), borderValue=0.0)
    return cropped, M


def trans_points2d(pts, M):
    new_pts = np.zeros(shape=pts.shape, dtype=np.float32)
    for i in range(pts.shape[0]):
        pt = pts[i]
        new_pt = np.array([pt[0], pt[1], 1.0], dtype=np.float32)
        new_pt = np.dot(M, new_pt)
        # print('new_pt', new_pt.shape, new_pt)
        new_pts[i] = new_pt[0:2]

    return new_pts


def trans_points3d(pts, M):
    scale = np.sqrt(M[0][0] * M[0][0] + M[0][1] * M[0][1])
    # print(scale)
    new_pts = np.zeros(shape=pts.shape, dtype=np.float32)
    for i in range(pts.shape[0]):
        pt = pts[i]
        new_pt = np.array([pt[0], pt[1], 1.0], dtype=np.float32)
        new_pt = np.dot(M, new_pt)
        # print('new_pt', new_pt.shape, new_pt)
        new_pts[i][0:2] = new_pt[0:2]
        new_pts[i][2] = pts[i][2] * scale

    return new_pts


def trans_points(pts, M):
    if pts.shape[1] == 2:
        return trans_points2d(pts, M)
    else:
        return trans_points3d(pts, M)


def inverse_transform(mat: np.ndarray) -> np.ndarray:
    # inverse the Affine transformation matrix
    inv_mat = np.zeros([2, 3])
    div1 = mat[0][0] * mat[1][1] - mat[0][1] * mat[1][0]
    inv_mat[0][0] = mat[1][1] / div1
    inv_mat[0][1] = -mat[0][1] / div1
    inv_mat[0][2] = -(mat[0][2] * mat[1][1] - mat[0][1] * mat[1][2]) / div1
    div2 = mat[0][1] * mat[1][0] - mat[0][0] * mat[1][1]
    inv_mat[1][0] = mat[1][0] / div2
    inv_mat[1][1] = -mat[0][0] / div2
    inv_mat[1][2] = -(mat[0][2] * mat[1][0] - mat[0][0] * mat[1][2]) / div2
    return inv_mat


def inverse_transform_batch(mat: torch.Tensor) -> torch.Tensor:
    # inverse the Affine transformation matrix
    inv_mat = torch.zeros_like(mat)
    div1 = mat[:, 0, 0] * mat[:, 1, 1] - mat[:, 0, 1] * mat[:, 1, 0]
    inv_mat[:, 0, 0] = mat[:, 1, 1] / div1
    inv_mat[:, 0, 1] = -mat[:, 0, 1] / div1
    inv_mat[:, 0, 2] = (
        -(mat[:, 0, 2] * mat[:, 1, 1] - mat[:, 0, 1] * mat[:, 1, 2]) / div1
    )
    div2 = mat[:, 0, 1] * mat[:, 1, 0] - mat[:, 0, 0] * mat[:, 1, 1]
    inv_mat[:, 1, 0] = mat[:, 1, 0] / div2
    inv_mat[:, 1, 1] = -mat[:, 0, 0] / div2
    inv_mat[:, 1, 2] = (
        -(mat[:, 0, 2] * mat[:, 1, 0] - mat[:, 0, 0] * mat[:, 1, 2]) / div2
    )
    return inv_mat


def align_face(
    img: np.ndarray, key_points: np.ndarray, crop_size: int, mode: str = "ffhq"
) -> Tuple[Iterable[np.ndarray], Iterable[np.ndarray]]:
    align_imgs = []
    transforms = []
    for i in range(key_points.shape[0]):
        kps = key_points[i]
        transform_matrix, _ = estimate_norm(kps, crop_size, mode=mode)
        align_img = cv2.warpAffine(
            img, transform_matrix, (crop_size, crop_size), borderValue=0.0
        )
        align_imgs.append(align_img)
        transforms.append(transform_matrix)

    return align_imgs, transforms