Skip to content

Commit d92d8bb

Browse files
committed
Initial commit
0 parents  commit d92d8bb

Some content is hidden

Large Commits have some content hidden by default. Use the searchbox below for content that may be hidden.

78 files changed

+3928
-0
lines changed

.idea/.gitignore

+8
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

.idea/deployment.xml

+77
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

.idea/inspectionProfiles/profiles_settings.xml

+6
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

.idea/misc.xml

+4
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

.idea/modules.xml

+8
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

.idea/pythonProject5.iml

+8
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.
+55
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,55 @@
1+
import os
2+
import bz2
3+
import argparse
4+
from Project.aligned_image.ffhq_dataset.face_alignment import image_align
5+
from Project.aligned_image.ffhq_dataset.landmarks_detector import LandmarksDetector
6+
7+
8+
def unpack_bz2(src_path):
9+
data = bz2.BZ2File(src_path).read()
10+
dst_path = src_path[:-4]
11+
with open(dst_path, 'wb') as fp:
12+
fp.write(data)
13+
return dst_path
14+
15+
16+
def align():
17+
"""
18+
Extracts and aligns all faces from images using DLib and a function from original FFHQ dataset preparation step
19+
python align_images.py /raw_images /aligned_images
20+
"""
21+
parser = argparse.ArgumentParser(description='Align faces from input images', formatter_class=argparse.ArgumentDefaultsHelpFormatter)
22+
parser.add_argument('--raw_dir',default='static/img_in', help='Directory with raw images for face alignment')
23+
parser.add_argument('--aligned_dir',default='static/img_aligned', help='Directory for storing aligned images')
24+
parser.add_argument('--output_size', default=1024, help='The dimension of images for input to the model', type=int)
25+
parser.add_argument('--x_scale', default=1, help='Scaling factor for x dimension', type=float)
26+
parser.add_argument('--y_scale', default=1, help='Scaling factor for y dimension', type=float)
27+
parser.add_argument('--em_scale', default=0.1, help='Scaling factor for eye-mouth distance', type=float)
28+
parser.add_argument('--use_alpha', default=False, help='Add an alpha channel for masking', type=bool)
29+
30+
args, other_args = parser.parse_known_args()
31+
32+
landmarks_model_path = unpack_bz2("Project/pretrained_models/shape_predictor_68_face_landmarks.dat.bz2")
33+
RAW_IMAGES_DIR = args.raw_dir
34+
ALIGNED_IMAGES_DIR = args.aligned_dir
35+
36+
landmarks_detector = LandmarksDetector(landmarks_model_path)
37+
for img_name in os.listdir(RAW_IMAGES_DIR):
38+
print('Aligning %s ...' % img_name)
39+
try:
40+
raw_img_path = os.path.join(RAW_IMAGES_DIR, img_name)
41+
fn = face_img_name = '%s_%02d.png' % (os.path.splitext(img_name)[0], 1)
42+
if os.path.isfile(fn):
43+
continue
44+
print('Getting landmarks...')
45+
for i, face_landmarks in enumerate(landmarks_detector.get_landmarks(raw_img_path), start=1):
46+
try:
47+
print('Starting face alignment...')
48+
face_img_name = '%s_%02d.png' % (os.path.splitext(img_name)[0], i)
49+
aligned_face_path = os.path.join(ALIGNED_IMAGES_DIR, face_img_name)
50+
image_align(raw_img_path, aligned_face_path, face_landmarks, output_size=args.output_size, x_scale=args.x_scale, y_scale=args.y_scale, em_scale=args.em_scale, alpha=args.use_alpha)
51+
print('Wrote result %s' % aligned_face_path)
52+
except:
53+
print("Exception in face alignment!")
54+
except:
55+
print("Exception in landmark detection!")

Project/aligned_image/ffhq_dataset/__init__.py

Whitespace-only changes.
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,92 @@
1+
import numpy as np
2+
import scipy.ndimage
3+
import os
4+
import PIL.Image
5+
6+
7+
def image_align(src_file, dst_file, face_landmarks, output_size=1024, transform_size=4096, enable_padding=True, x_scale=1, y_scale=1, em_scale=0.1, alpha=False):
8+
# Align function from FFHQ dataset pre-processing step
9+
# https://github.com/NVlabs/ffhq-dataset/blob/master/download_ffhq.py
10+
11+
lm = np.array(face_landmarks)
12+
lm_chin = lm[0 : 17] # left-right
13+
lm_eyebrow_left = lm[17 : 22] # left-right
14+
lm_eyebrow_right = lm[22 : 27] # left-right
15+
lm_nose = lm[27 : 31] # top-down
16+
lm_nostrils = lm[31 : 36] # top-down
17+
lm_eye_left = lm[36 : 42] # left-clockwise
18+
lm_eye_right = lm[42 : 48] # left-clockwise
19+
lm_mouth_outer = lm[48 : 60] # left-clockwise
20+
lm_mouth_inner = lm[60 : 68] # left-clockwise
21+
22+
# Calculate auxiliary vectors.
23+
eye_left = np.mean(lm_eye_left, axis=0)
24+
eye_right = np.mean(lm_eye_right, axis=0)
25+
eye_avg = (eye_left + eye_right) * 0.5
26+
eye_to_eye = eye_right - eye_left
27+
mouth_left = lm_mouth_outer[0]
28+
mouth_right = lm_mouth_outer[6]
29+
mouth_avg = (mouth_left + mouth_right) * 0.5
30+
eye_to_mouth = mouth_avg - eye_avg
31+
32+
# Choose oriented crop rectangle.
33+
x = eye_to_eye - np.flipud(eye_to_mouth) * [-1, 1]
34+
x /= np.hypot(*x)
35+
x *= max(np.hypot(*eye_to_eye) * 2.0, np.hypot(*eye_to_mouth) * 1.8)
36+
x *= x_scale
37+
y = np.flipud(x) * [-y_scale, y_scale]
38+
c = eye_avg + eye_to_mouth * em_scale
39+
quad = np.stack([c - x - y, c - x + y, c + x + y, c + x - y])
40+
qsize = np.hypot(*x) * 2
41+
42+
# Load in-the-wild image.
43+
if not os.path.isfile(src_file):
44+
print('\nCannot find source image. Please run "--wilds" before "--align".')
45+
return
46+
img = PIL.Image.open(src_file).convert('RGBA').convert('RGB')
47+
48+
# Shrink.
49+
shrink = int(np.floor(qsize / output_size * 0.5))
50+
if shrink > 1:
51+
rsize = (int(np.rint(float(img.size[0]) / shrink)), int(np.rint(float(img.size[1]) / shrink)))
52+
img = img.resize(rsize, PIL.Image.ANTIALIAS)
53+
quad /= shrink
54+
qsize /= shrink
55+
56+
# Crop.
57+
border = max(int(np.rint(qsize * 0.1)), 3)
58+
crop = (int(np.floor(min(quad[:,0]))), int(np.floor(min(quad[:,1]))), int(np.ceil(max(quad[:,0]))), int(np.ceil(max(quad[:,1]))))
59+
crop = (max(crop[0] - border, 0), max(crop[1] - border, 0), min(crop[2] + border, img.size[0]), min(crop[3] + border, img.size[1]))
60+
if crop[2] - crop[0] < img.size[0] or crop[3] - crop[1] < img.size[1]:
61+
img = img.crop(crop)
62+
quad -= crop[0:2]
63+
64+
# Pad.
65+
pad = (int(np.floor(min(quad[:,0]))), int(np.floor(min(quad[:,1]))), int(np.ceil(max(quad[:,0]))), int(np.ceil(max(quad[:,1]))))
66+
pad = (max(-pad[0] + border, 0), max(-pad[1] + border, 0), max(pad[2] - img.size[0] + border, 0), max(pad[3] - img.size[1] + border, 0))
67+
if enable_padding and max(pad) > border - 4:
68+
pad = np.maximum(pad, int(np.rint(qsize * 0.3)))
69+
img = np.pad(np.float32(img), ((pad[1], pad[3]), (pad[0], pad[2]), (0, 0)), 'reflect')
70+
h, w, _ = img.shape
71+
y, x, _ = np.ogrid[:h, :w, :1]
72+
mask = np.maximum(1.0 - np.minimum(np.float32(x) / pad[0], np.float32(w-1-x) / pad[2]), 1.0 - np.minimum(np.float32(y) / pad[1], np.float32(h-1-y) / pad[3]))
73+
blur = qsize * 0.02
74+
img += (scipy.ndimage.gaussian_filter(img, [blur, blur, 0]) - img) * np.clip(mask * 3.0 + 1.0, 0.0, 1.0)
75+
img += (np.median(img, axis=(0,1)) - img) * np.clip(mask, 0.0, 1.0)
76+
img = np.uint8(np.clip(np.rint(img), 0, 255))
77+
if alpha:
78+
mask = 1-np.clip(3.0 * mask, 0.0, 1.0)
79+
mask = np.uint8(np.clip(np.rint(mask*255), 0, 255))
80+
img = np.concatenate((img, mask), axis=2)
81+
img = PIL.Image.fromarray(img, 'RGBA')
82+
else:
83+
img = PIL.Image.fromarray(img, 'RGB')
84+
quad += pad[:2]
85+
86+
# Transform.
87+
img = img.transform((transform_size, transform_size), PIL.Image.QUAD, (quad + 0.5).flatten(), PIL.Image.BILINEAR)
88+
if output_size < transform_size:
89+
img = img.resize((output_size, output_size), PIL.Image.ANTIALIAS)
90+
91+
# Save aligned image.
92+
img.save(dst_file, 'PNG')
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,21 @@
1+
import dlib
2+
3+
4+
class LandmarksDetector:
5+
def __init__(self, predictor_model_path):
6+
"""
7+
:param predictor_model_path: path to shape_predictor_68_face_landmarks.dat file
8+
"""
9+
self.detector = dlib.get_frontal_face_detector() # cnn_face_detection_model_v1 also can be used
10+
self.shape_predictor = dlib.shape_predictor(predictor_model_path)
11+
12+
def get_landmarks(self, image):
13+
img = dlib.load_rgb_image(image)
14+
dets = self.detector(img, 1)
15+
16+
for detection in dets:
17+
try:
18+
face_landmarks = [(item.x, item.y) for item in self.shape_predictor(img, detection).parts()]
19+
yield face_landmarks
20+
except:
21+
print("Exception in get_landmarks()!")

Project/configs/__init__.py

Whitespace-only changes.

Project/configs/data_configs.py

+41
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,41 @@
1+
from Project.configs import transforms_config
2+
from Project.configs.paths_config import dataset_paths
3+
4+
5+
DATASETS = {
6+
'ffhq_encode': {
7+
'transforms': transforms_config.EncodeTransforms,
8+
'train_source_root': dataset_paths['ffhq'],
9+
'train_target_root': dataset_paths['ffhq'],
10+
'test_source_root': dataset_paths['celeba_test'],
11+
'test_target_root': dataset_paths['celeba_test'],
12+
},
13+
'cars_encode': {
14+
'transforms': transforms_config.CarsEncodeTransforms,
15+
'train_source_root': dataset_paths['cars_train'],
16+
'train_target_root': dataset_paths['cars_train'],
17+
'test_source_root': dataset_paths['cars_test'],
18+
'test_target_root': dataset_paths['cars_test'],
19+
},
20+
'horse_encode': {
21+
'transforms': transforms_config.EncodeTransforms,
22+
'train_source_root': dataset_paths['horse_train'],
23+
'train_target_root': dataset_paths['horse_train'],
24+
'test_source_root': dataset_paths['horse_test'],
25+
'test_target_root': dataset_paths['horse_test'],
26+
},
27+
'church_encode': {
28+
'transforms': transforms_config.EncodeTransforms,
29+
'train_source_root': dataset_paths['church_train'],
30+
'train_target_root': dataset_paths['church_train'],
31+
'test_source_root': dataset_paths['church_test'],
32+
'test_target_root': dataset_paths['church_test'],
33+
},
34+
'cats_encode': {
35+
'transforms': transforms_config.EncodeTransforms,
36+
'train_source_root': dataset_paths['cats_train'],
37+
'train_target_root': dataset_paths['cats_train'],
38+
'test_source_root': dataset_paths['cats_test'],
39+
'test_target_root': dataset_paths['cats_test'],
40+
}
41+
}

Project/configs/paths_config.py

+28
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,28 @@
1+
dataset_paths = {
2+
# Face Datasets (In the paper: FFHQ - train, CelebAHQ - test)
3+
'ffhq': '',
4+
'celeba_test': '',
5+
6+
# Cars Dataset (In the paper: Stanford cars)
7+
'cars_train': '',
8+
'cars_test': '',
9+
10+
# Horse Dataset (In the paper: LSUN Horse)
11+
'horse_train': '',
12+
'horse_test': '',
13+
14+
# Church Dataset (In the paper: LSUN Church)
15+
'church_train': '',
16+
'church_test': '',
17+
18+
# Cats Dataset (In the paper: LSUN Cat)
19+
'cats_train': '',
20+
'cats_test': ''
21+
}
22+
23+
model_paths = {
24+
'stylegan_ffhq': 'pretrained_models/stylegan2-ffhq-config-f.pt',
25+
'ir_se50': 'pretrained_models/model_ir_se50.pth',
26+
'shape_predictor': 'pretrained_models/shape_predictor_68_face_landmarks.dat',
27+
'moco': 'pretrained_models/moco_v2_800ep_pretrain.pth'
28+
}

0 commit comments

Comments
 (0)