Shellbrady/LivePortrait5

1import cv22import numpy as np3from skimage import transform as trans4 5 6arcface_dst = np.array(7    [[38.2946, 51.6963], [73.5318, 51.5014], [56.0252, 71.7366],8     [41.5493, 92.3655], [70.7299, 92.2041]],9    dtype=np.float32)10 11def estimate_norm(lmk, image_size=112,mode='arcface'):12    assert lmk.shape == (5, 2)13    assert image_size%112==0 or image_size%128==014    if image_size%112==0:15        ratio = float(image_size)/112.016        diff_x = 017    else:18        ratio = float(image_size)/128.019        diff_x = 8.0*ratio20    dst = arcface_dst * ratio21    dst[:,0] += diff_x22    tform = trans.SimilarityTransform()23    tform.estimate(lmk, dst)24    M = tform.params[0:2, :]25    return M26 27def norm_crop(img, landmark, image_size=112, mode='arcface'):28    M = estimate_norm(landmark, image_size, mode)29    warped = cv2.warpAffine(img, M, (image_size, image_size), borderValue=0.0)30    return warped31 32def norm_crop2(img, landmark, image_size=112, mode='arcface'):33    M = estimate_norm(landmark, image_size, mode)34    warped = cv2.warpAffine(img, M, (image_size, image_size), borderValue=0.0)35    return warped, M36 37def square_crop(im, S):38    if im.shape[0] > im.shape[1]:39        height = S40        width = int(float(im.shape[1]) / im.shape[0] * S)41        scale = float(S) / im.shape[0]42    else:43        width = S44        height = int(float(im.shape[0]) / im.shape[1] * S)45        scale = float(S) / im.shape[1]46    resized_im = cv2.resize(im, (width, height))47    det_im = np.zeros((S, S, 3), dtype=np.uint8)48    det_im[:resized_im.shape[0], :resized_im.shape[1], :] = resized_im49    return det_im, scale50 51 52def transform(data, center, output_size, scale, rotation):53    scale_ratio = scale54    rot = float(rotation) * np.pi / 180.055    #translation = (output_size/2-center[0]*scale_ratio, output_size/2-center[1]*scale_ratio)56    t1 = trans.SimilarityTransform(scale=scale_ratio)57    cx = center[0] * scale_ratio58    cy = center[1] * scale_ratio59    t2 = trans.SimilarityTransform(translation=(-1 * cx, -1 * cy))60    t3 = trans.SimilarityTransform(rotation=rot)61    t4 = trans.SimilarityTransform(translation=(output_size / 2,62                                                output_size / 2))63    t = t1 + t2 + t3 + t464    M = t.params[0:2]65    cropped = cv2.warpAffine(data,66                             M, (output_size, output_size),67                             borderValue=0.0)68    return cropped, M69 70 71def trans_points2d(pts, M):72    new_pts = np.zeros(shape=pts.shape, dtype=np.float32)73    for i in range(pts.shape[0]):74        pt = pts[i]75        new_pt = np.array([pt[0], pt[1], 1.], dtype=np.float32)76        new_pt = np.dot(M, new_pt)77        #print('new_pt', new_pt.shape, new_pt)78        new_pts[i] = new_pt[0:2]79 80    return new_pts81 82 83def trans_points3d(pts, M):84    scale = np.sqrt(M[0][0] * M[0][0] + M[0][1] * M[0][1])85    #print(scale)86    new_pts = np.zeros(shape=pts.shape, dtype=np.float32)87    for i in range(pts.shape[0]):88        pt = pts[i]89        new_pt = np.array([pt[0], pt[1], 1.], dtype=np.float32)90        new_pt = np.dot(M, new_pt)91        #print('new_pt', new_pt.shape, new_pt)92        new_pts[i][0:2] = new_pt[0:2]93        new_pts[i][2] = pts[i][2] * scale94 95    return new_pts96 97 98def trans_points(pts, M):99    if pts.shape[1] == 2:100        return trans_points2d(pts, M)101    else:102        return trans_points3d(pts, M)103 104