Spaces:
Running
Running
| import cv2 | |
| import numpy as np | |
| def get_similarity_transform(src_pts, dst_pts): | |
| """Get similarity transform matrix from src_pts to dst_pts | |
| Args: | |
| src_pts: Kx2 np.array | |
| source points matrix, each row is a pair of coordinates (x, y) | |
| dst_pts: Kx2 np.array | |
| destination points matrix, each row is a pair of coordinates (x, y) | |
| Returns: | |
| xform_matrix: 3x3 np.array | |
| transform matrix from src_pts to dst_pts | |
| """ | |
| src_pts = np.asarray(src_pts) | |
| dst_pts = np.asarray(dst_pts) | |
| assert src_pts.shape == dst_pts.shape | |
| assert (src_pts.ndim == 2) and (src_pts.shape[-1] == 2) | |
| npts = src_pts.shape[0] | |
| src_x = src_pts[:, 0].reshape((-1, 1)) | |
| src_y = src_pts[:, 1].reshape((-1, 1)) | |
| tmp1 = np.hstack((src_x, -src_y, np.ones((npts, 1)), np.zeros((npts, 1)))) | |
| tmp2 = np.hstack((src_y, src_x, np.zeros((npts, 1)), np.ones((npts, 1)))) | |
| A = np.vstack((tmp1, tmp2)) | |
| dst_x = dst_pts[:, 0].reshape((-1, 1)) | |
| dst_y = dst_pts[:, 1].reshape((-1, 1)) | |
| b = np.vstack((dst_x, dst_y)) | |
| x = np.linalg.lstsq(A, b, rcond=-1)[0] | |
| x = np.squeeze(x) | |
| sc, ss, tx, ty = x[0], x[1], x[2], x[3] | |
| xform_matrix = np.array([ | |
| [sc, -ss, tx], | |
| [ss, sc, ty], | |
| [ 0, 0, 1] | |
| ]) | |
| return xform_matrix | |
| def align_and_crop(image, landmarks, std_landmarks, align_size, | |
| border_value=0, return_transform_matrix=False): | |
| landmarks = np.asarray(landmarks) | |
| std_landmarks = np.asarray(std_landmarks) | |
| xform_matrix = get_similarity_transform(landmarks, std_landmarks) | |
| landmarks_ex = np.pad(landmarks, ((0,0),(0,1)), mode='constant', constant_values=1) | |
| dst_landmarks = np.dot(landmarks_ex, xform_matrix[:2,:].T) | |
| dst_image = cv2.warpAffine(image, xform_matrix[:2,:], dsize=align_size, | |
| borderValue=border_value) | |
| if return_transform_matrix: | |
| return dst_image, dst_landmarks, xform_matrix | |
| else: | |
| return dst_image, dst_landmarks | |