Spaces:
Build error
Build error
| import face_alignment | |
| import os | |
| import cv2 | |
| import skimage.transform as trans | |
| import argparse | |
| import torch | |
| import numpy as np | |
| import tqdm | |
| device = 'cuda' if torch.cuda.is_available() else 'cpu' | |
| def get_affine(src): | |
| dst = np.array([[87, 59], | |
| [137, 59], | |
| [112, 120]], dtype=np.float32) | |
| tform = trans.SimilarityTransform() | |
| tform.estimate(src, dst) | |
| M = tform.params[0:2, :] | |
| return M | |
| def affine_align_img(img, M, crop_size=224): | |
| warped = cv2.warpAffine(img, M, (crop_size, crop_size), borderValue=0.0) | |
| return warped | |
| def affine_align_3landmarks(landmarks, M): | |
| new_landmarks = np.concatenate([landmarks, np.ones((3, 1))], 1) | |
| affined_landmarks = np.matmul(new_landmarks, M.transpose()) | |
| return affined_landmarks | |
| def get_eyes_mouths(landmark): | |
| three_points = np.zeros((3, 2)) | |
| three_points[0] = landmark[36:42].mean(0) | |
| three_points[1] = landmark[42:48].mean(0) | |
| three_points[2] = landmark[60:68].mean(0) | |
| return three_points | |
| def get_mouth_bias(three_points): | |
| bias = np.array([112, 120]) - three_points[2] | |
| return bias | |
| def align_folder(folder_path, folder_save_path): | |
| fa = face_alignment.FaceAlignment(face_alignment.LandmarksType._2D, device=device) | |
| preds = fa.get_landmarks_from_directory(folder_path) | |
| sumpoints = 0 | |
| three_points_list = [] | |
| for img in tqdm.tqdm(preds.keys(), desc='preprocessing..'): | |
| pred_points = np.array(preds[img]) | |
| if pred_points is None or len(pred_points.shape) != 3: | |
| print('preprocessing failed') | |
| return False | |
| else: | |
| num_faces, size, _ = pred_points.shape | |
| if num_faces == 1 and size == 68: | |
| three_points = get_eyes_mouths(pred_points[0]) | |
| sumpoints += three_points | |
| three_points_list.append(three_points) | |
| else: | |
| print('preprocessing failed') | |
| return False | |
| avg_points = sumpoints / len(preds) | |
| M = get_affine(avg_points) | |
| p_bias = None | |
| for i, img_pth in tqdm.tqdm(enumerate(preds.keys()), desc='affine and save'): | |
| three_points = three_points_list[i] | |
| affined_3landmarks = affine_align_3landmarks(three_points, M) | |
| bias = get_mouth_bias(affined_3landmarks) | |
| if p_bias is None: | |
| bias = bias | |
| else: | |
| bias = p_bias * 0.2 + bias * 0.8 | |
| p_bias = bias | |
| M_i = M.copy() | |
| M_i[:, 2] = M[:, 2] + bias | |
| img = cv2.imread(img_pth) | |
| wrapped = affine_align_img(img, M_i) | |
| img_save_path = os.path.join(folder_save_path, img_pth.split('/')[-1]) | |
| cv2.imwrite(img_save_path, wrapped) | |
| print('cropped files saved at {}'.format(folder_save_path)) | |
| def main(): | |
| parser = argparse.ArgumentParser() | |
| parser.add_argument('--folder_path', help='the folder which needs processing') | |
| args = parser.parse_args() | |
| if os.path.isdir(args.folder_path): | |
| home_path = '/'.join(args.folder_path.split('/')[:-1]) | |
| save_img_path = os.path.join(home_path, args.folder_path.split('/')[-1] + '_cropped') | |
| os.makedirs(save_img_path, exist_ok=True) | |
| align_folder(args.folder_path, save_img_path) | |
| if __name__ == '__main__': | |
| main() | |