-
Notifications
You must be signed in to change notification settings - Fork 3
/
gen_data.py
100 lines (81 loc) · 3.7 KB
/
gen_data.py
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
# -*- coding: utf-8 -*-
"""
Created on Fri Apr 19 10:32:37 2019
@author: lyfeng
Make train/test datasets from raw datasets
input: raw original datasets
output: aligned datasets(customized size of images)
"""
import cv2
import os
import numpy as np
import torchvision.transforms as transforms
from matlab_cp2tform import get_similarity_transform_for_cv2
from src import detect_faces
from PIL import Image
from tqdm import tqdm
import matplotlib.pyplot as plt
def alignment(src_img,src_pts):
of = 0
ref_pts = [ [30.2946+of, 51.6963+of],[65.5318+of, 51.5014+of],
[48.0252+of, 71.7366+of],[33.5493+of, 92.3655+of],[62.7299+of, 92.2041+of] ]
crop_size = (96+of*2, 112+of*2) # 96,112
s = np.array(src_pts).astype(np.float32)
r = np.array(ref_pts).astype(np.float32)
tfm = get_similarity_transform_for_cv2(s, r) #
face_img = cv2.warpAffine(src_img, tfm, crop_size) # affine transformation
return face_img
# have no landmarks at initial
def produce_aligned_images(raw_data_root, new_data_root):
all_sequences = os.listdir(raw_data_root)
num_id = len(all_sequences)
for each_id in tqdm(all_sequences):
imgs_path = raw_data_root + '/' + each_id
for img_name in os.listdir(imgs_path):
img = Image.open(imgs_path +'/'+ img_name)
bounding_boxes, landmarks = detect_faces(img)
try:
landmarks = landmarks[0]
src_pts = [[landmarks[0],landmarks[5]],
[landmarks[1],landmarks[6]],
[landmarks[2],landmarks[7]],
[landmarks[3],landmarks[8]],
[landmarks[4],landmarks[9]]]
except: pass
img = alignment(np.array(img),src_pts)
try:
img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB) # cv2 default:BGR
except:
pass
save_path = new_data_root + '/' + each_id
if not os.path.exists(save_path):
os.makedirs(save_path)
cv2.imwrite(save_path+'/'+img_name, img)
break
# have landmarks aleardy
def produce_aligned_images_from_txt(raw_data_root, new_data_root, landmark_file):
with open(landmark_file) as f:
for index, line in enumerate(f):
split = line.split('\t')
nameinzip = split[0] # eg: 0000159/266.jpg
classid = int(split[1])
src_pts = [] # store the landmarks
for i in range(5):
src_pts.append([int(split[2*i+2]),int(split[2*i+3])])
imgs_path = raw_data_root + '/' + nameinzip
img = Image.open(imgs_path)
img = alignment(np.array(img),src_pts)
try:
img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB) # cv2 default:BGR
except:
pass
save_path = new_data_root + '/' + nameinzip.split('/')[0]
if not os.path.exists(save_path):
os.makedirs(save_path)
cv2.imwrite(save_path + '/' + nameinzip.split('/')[1], img)
#if __name__ =='__main__':
# produce_aligned_images('G:/数据资料/人脸数据集/CASIA-WebFace/CASIA-WebFace',
# 'D:/Face Recognition/CASIA-WebFace-aligned')
# produce_aligned_images_from_txt('G:/数据资料/人脸数据集/CASIA-WebFace/CASIA-WebFace',
# 'D:/Face Recognition/CASIA-WebFace-aligned',
# 'D:/Face Recognition/sphereface_pytorch-master/data/casia_landmark.txt')