Shellbrady/LivePortrait5

1import cv22import math3import numpy as np4from skimage import transform as trans5 6 7def transform(data, center, output_size, scale, rotation):8    scale_ratio = scale9    rot = float(rotation) * np.pi / 180.010    #translation = (output_size/2-center[0]*scale_ratio, output_size/2-center[1]*scale_ratio)11    t1 = trans.SimilarityTransform(scale=scale_ratio)12    cx = center[0] * scale_ratio13    cy = center[1] * scale_ratio14    t2 = trans.SimilarityTransform(translation=(-1 * cx, -1 * cy))15    t3 = trans.SimilarityTransform(rotation=rot)16    t4 = trans.SimilarityTransform(translation=(output_size / 2,17                                                output_size / 2))18    t = t1 + t2 + t3 + t419    M = t.params[0:2]20    cropped = cv2.warpAffine(data,21                             M, (output_size, output_size),22                             borderValue=0.0)23    return cropped, M24 25 26def trans_points2d(pts, M):27    new_pts = np.zeros(shape=pts.shape, dtype=np.float32)28    for i in range(pts.shape[0]):29        pt = pts[i]30        new_pt = np.array([pt[0], pt[1], 1.], dtype=np.float32)31        new_pt = np.dot(M, new_pt)32        #print('new_pt', new_pt.shape, new_pt)33        new_pts[i] = new_pt[0:2]34 35    return new_pts36 37 38def trans_points3d(pts, M):39    scale = np.sqrt(M[0][0] * M[0][0] + M[0][1] * M[0][1])40    #print(scale)41    new_pts = np.zeros(shape=pts.shape, dtype=np.float32)42    for i in range(pts.shape[0]):43        pt = pts[i]44        new_pt = np.array([pt[0], pt[1], 1.], dtype=np.float32)45        new_pt = np.dot(M, new_pt)46        #print('new_pt', new_pt.shape, new_pt)47        new_pts[i][0:2] = new_pt[0:2]48        new_pts[i][2] = pts[i][2] * scale49 50    return new_pts51 52 53def trans_points(pts, M):54    if pts.shape[1] == 2:55        return trans_points2d(pts, M)56    else:57        return trans_points3d(pts, M)58 59def estimate_affine_matrix_3d23d(X, Y):60    ''' Using least-squares solution 61    Args:62        X: [n, 3]. 3d points(fixed)63        Y: [n, 3]. corresponding 3d points(moving). Y = PX64    Returns:65        P_Affine: (3, 4). Affine camera matrix (the third row is [0, 0, 0, 1]).66    '''67    X_homo = np.hstack((X, np.ones([X.shape[0],1]))) #n x 468    P = np.linalg.lstsq(X_homo, Y)[0].T # Affine matrix. 3 x 469    return P70 71def P2sRt(P):72    ''' decompositing camera matrix P73    Args: 74        P: (3, 4). Affine Camera Matrix.75    Returns:76        s: scale factor.77        R: (3, 3). rotation matrix.78        t: (3,). translation. 79    '''80    t = P[:, 3]81    R1 = P[0:1, :3]82    R2 = P[1:2, :3]83    s = (np.linalg.norm(R1) + np.linalg.norm(R2))/2.084    r1 = R1/np.linalg.norm(R1)85    r2 = R2/np.linalg.norm(R2)86    r3 = np.cross(r1, r2)87 88    R = np.concatenate((r1, r2, r3), 0)89    return s, R, t90 91def matrix2angle(R):92    ''' get three Euler angles from Rotation Matrix93    Args:94        R: (3,3). rotation matrix95    Returns:96        x: pitch97        y: yaw98        z: roll99    '''100    sy = math.sqrt(R[0,0] * R[0,0] +  R[1,0] * R[1,0])101     102    singular = sy < 1e-6103 104    if  not singular :105        x = math.atan2(R[2,1] , R[2,2])106        y = math.atan2(-R[2,0], sy)107        z = math.atan2(R[1,0], R[0,0])108    else :109        x = math.atan2(-R[1,2], R[1,1])110        y = math.atan2(-R[2,0], sy)111        z = 0112 113    # rx, ry, rz = np.rad2deg(x), np.rad2deg(y), np.rad2deg(z)114    rx, ry, rz = x*180/np.pi, y*180/np.pi, z*180/np.pi115    return rx, ry, rz116 117