| 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() |
|
|
