paulo061/codeformer

1import cv22import numpy as np3 4from .matlab_cp2tform import get_similarity_transform_for_cv25 6# reference facial points, a list of coordinates (x,y)7REFERENCE_FACIAL_POINTS = [[30.29459953, 51.69630051], [65.53179932, 51.50139999], [48.02519989, 71.73660278],8                           [33.54930115, 92.3655014], [62.72990036, 92.20410156]]9 10DEFAULT_CROP_SIZE = (96, 112)11 12 13class FaceWarpException(Exception):14 15    def __str__(self):16        return 'In File {}:{}'.format(__file__, super.__str__(self))17 18 19def get_reference_facial_points(output_size=None, inner_padding_factor=0.0, outer_padding=(0, 0), default_square=False):20    """21    Function:22    ----------23        get reference 5 key points according to crop settings:24        0. Set default crop_size:25            if default_square:26                crop_size = (112, 112)27            else:28                crop_size = (96, 112)29        1. Pad the crop_size by inner_padding_factor in each side;30        2. Resize crop_size into (output_size - outer_padding*2),31            pad into output_size with outer_padding;32        3. Output reference_5point;33    Parameters:34    ----------35        @output_size: (w, h) or None36            size of aligned face image37        @inner_padding_factor: (w_factor, h_factor)38            padding factor for inner (w, h)39        @outer_padding: (w_pad, h_pad)40            each row is a pair of coordinates (x, y)41        @default_square: True or False42            if True:43                default crop_size = (112, 112)44            else:45                default crop_size = (96, 112);46        !!! make sure, if output_size is not None:47                (output_size - outer_padding)48                = some_scale * (default crop_size * (1.0 +49                inner_padding_factor))50    Returns:51    ----------52        @reference_5point: 5x2 np.array53            each row is a pair of transformed coordinates (x, y)54    """55 56    tmp_5pts = np.array(REFERENCE_FACIAL_POINTS)57    tmp_crop_size = np.array(DEFAULT_CROP_SIZE)58 59    # 0) make the inner region a square60    if default_square:61        size_diff = max(tmp_crop_size) - tmp_crop_size62        tmp_5pts += size_diff / 263        tmp_crop_size += size_diff64 65    if (output_size and output_size[0] == tmp_crop_size[0] and output_size[1] == tmp_crop_size[1]):66 67        return tmp_5pts68 69    if (inner_padding_factor == 0 and outer_padding == (0, 0)):70        if output_size is None:71            return tmp_5pts72        else:73            raise FaceWarpException('No paddings to do, output_size must be None or {}'.format(tmp_crop_size))74 75    # check output size76    if not (0 <= inner_padding_factor <= 1.0):77        raise FaceWarpException('Not (0 <= inner_padding_factor <= 1.0)')78 79    if ((inner_padding_factor > 0 or outer_padding[0] > 0 or outer_padding[1] > 0) and output_size is None):80        output_size = tmp_crop_size * \81            (1 + inner_padding_factor * 2).astype(np.int32)82        output_size += np.array(outer_padding)83    if not (outer_padding[0] < output_size[0] and outer_padding[1] < output_size[1]):84        raise FaceWarpException('Not (outer_padding[0] < output_size[0] and outer_padding[1] < output_size[1])')85 86    # 1) pad the inner region according inner_padding_factor87    if inner_padding_factor > 0:88        size_diff = tmp_crop_size * inner_padding_factor * 289        tmp_5pts += size_diff / 290        tmp_crop_size += np.round(size_diff).astype(np.int32)91 92    # 2) resize the padded inner region93    size_bf_outer_pad = np.array(output_size) - np.array(outer_padding) * 294 95    if size_bf_outer_pad[0] * tmp_crop_size[1] != size_bf_outer_pad[1] * tmp_crop_size[0]:96        raise FaceWarpException('Must have (output_size - outer_padding)'97                                '= some_scale * (crop_size * (1.0 + inner_padding_factor)')98 99    scale_factor = size_bf_outer_pad[0].astype(np.float32) / tmp_crop_size[0]100    tmp_5pts = tmp_5pts * scale_factor101    #    size_diff = tmp_crop_size * (scale_factor - min(scale_factor))102    #    tmp_5pts = tmp_5pts + size_diff / 2103    tmp_crop_size = size_bf_outer_pad104 105    # 3) add outer_padding to make output_size106    reference_5point = tmp_5pts + np.array(outer_padding)107    tmp_crop_size = output_size108 109    return reference_5point110 111 112def get_affine_transform_matrix(src_pts, dst_pts):113    """114    Function:115    ----------116        get affine transform matrix 'tfm' from src_pts to dst_pts117    Parameters:118    ----------119        @src_pts: Kx2 np.array120            source points matrix, each row is a pair of coordinates (x, y)121        @dst_pts: Kx2 np.array122            destination points matrix, each row is a pair of coordinates (x, y)123    Returns:124    ----------125        @tfm: 2x3 np.array126            transform matrix from src_pts to dst_pts127    """128 129    tfm = np.float32([[1, 0, 0], [0, 1, 0]])130    n_pts = src_pts.shape[0]131    ones = np.ones((n_pts, 1), src_pts.dtype)132    src_pts_ = np.hstack([src_pts, ones])133    dst_pts_ = np.hstack([dst_pts, ones])134 135    A, res, rank, s = np.linalg.lstsq(src_pts_, dst_pts_)136 137    if rank == 3:138        tfm = np.float32([[A[0, 0], A[1, 0], A[2, 0]], [A[0, 1], A[1, 1], A[2, 1]]])139    elif rank == 2:140        tfm = np.float32([[A[0, 0], A[1, 0], 0], [A[0, 1], A[1, 1], 0]])141 142    return tfm143 144 145def warp_and_crop_face(src_img, facial_pts, reference_pts=None, crop_size=(96, 112), align_type='smilarity'):146    """147    Function:148    ----------149        apply affine transform 'trans' to uv150    Parameters:151    ----------152        @src_img: 3x3 np.array153            input image154        @facial_pts: could be155            1)a list of K coordinates (x,y)156        or157            2) Kx2 or 2xK np.array158            each row or col is a pair of coordinates (x, y)159        @reference_pts: could be160            1) a list of K coordinates (x,y)161        or162            2) Kx2 or 2xK np.array163            each row or col is a pair of coordinates (x, y)164        or165            3) None166            if None, use default reference facial points167        @crop_size: (w, h)168            output face image size169        @align_type: transform type, could be one of170            1) 'similarity': use similarity transform171            2) 'cv2_affine': use the first 3 points to do affine transform,172                    by calling cv2.getAffineTransform()173            3) 'affine': use all points to do affine transform174    Returns:175    ----------176        @face_img: output face image with size (w, h) = @crop_size177    """178 179    if reference_pts is None:180        if crop_size[0] == 96 and crop_size[1] == 112:181            reference_pts = REFERENCE_FACIAL_POINTS182        else:183            default_square = False184            inner_padding_factor = 0185            outer_padding = (0, 0)186            output_size = crop_size187 188            reference_pts = get_reference_facial_points(output_size, inner_padding_factor, outer_padding,189                                                        default_square)190 191    ref_pts = np.float32(reference_pts)192    ref_pts_shp = ref_pts.shape193    if max(ref_pts_shp) < 3 or min(ref_pts_shp) != 2:194        raise FaceWarpException('reference_pts.shape must be (K,2) or (2,K) and K>2')195 196    if ref_pts_shp[0] == 2:197        ref_pts = ref_pts.T198 199    src_pts = np.float32(facial_pts)200    src_pts_shp = src_pts.shape201    if max(src_pts_shp) < 3 or min(src_pts_shp) != 2:202        raise FaceWarpException('facial_pts.shape must be (K,2) or (2,K) and K>2')203 204    if src_pts_shp[0] == 2:205        src_pts = src_pts.T206 207    if src_pts.shape != ref_pts.shape:208        raise FaceWarpException('facial_pts and reference_pts must have the same shape')209 210    if align_type == 'cv2_affine':211        tfm = cv2.getAffineTransform(src_pts[0:3], ref_pts[0:3])212    elif align_type == 'affine':213        tfm = get_affine_transform_matrix(src_pts, ref_pts)214    else:215        tfm = get_similarity_transform_for_cv2(src_pts, ref_pts)216 217    face_img = cv2.warpAffine(src_img, tfm, (crop_size[0], crop_size[1]))218 219    return face_img220