File size: 3,270 Bytes
e34aada |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 |
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()
|