-
Notifications
You must be signed in to change notification settings - Fork 12
/
Copy pathgenerate_data_and_score.py
168 lines (138 loc) · 6.51 KB
/
generate_data_and_score.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
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
# python3.7
"""Generates a collection of images with specified model.
Commonly, this file is used for data preparation. More specifically, before
exploring the hidden semantics from the latent space, user need to prepare a
collection of images. These images can be used for further attribute prediction.
In this way, it is able to build a relationship between input latent codes and
the corresponding attribute scores.
"""
import os.path
import argparse
from collections import defaultdict
import cv2
import numpy as np
from tqdm import tqdm
from styleGAN2_model.stylegan2_generator import StyleGAN2Generator
from interface.utils.logger import setup_logger
from classifier.classify import get_model,check_double_chin
from utils import str2bool
def parse_args():
"""Parses arguments."""
parser = argparse.ArgumentParser(
description='Generate images with given model.')
parser.add_argument('-o', '--output_dir', type=str, required=True,
help='Directory to save the output results. (required)')
parser.add_argument('-i', '--latent_codes_path', type=str, default='',
help='If specified, will load latent codes from given '
'path instead of randomly sampling. (optional)')
parser.add_argument('-n', '--num', type=int, default=50000,
help='Number of images to generate. This field will be '
'ignored if `latent_codes_path` is specified. '
'(default: 1)')
parser.add_argument('-S', '--generate_style', action='store_true',
help='If specified, will generate layer-wise style codes '
'in Style GAN. (default: do not generate styles)')
parser.add_argument('-I', '--generate_image', action='store_false',
help='If specified, will skip generating images in '
'Style GAN. (default: generate images)')
parser.add_argument('-p', '--truncation_psi', type=float,default='0.8')
parser.add_argument("--double_chin_only", type=str2bool, nargs='?',
const=False, default=False,
help="Only generate double chin images.")
return parser.parse_args()
def main():
"""Main function."""
args = parse_args()
logger = setup_logger(args.output_dir, logger_name='generate_data')
double_chin_checker=get_model()
logger.info(f'Initializing generator.')
model_name = 'stylegan2_ffhq'
model = StyleGAN2Generator(model_name, logger, truncation_psi=args.truncation_psi)
kwargs = {'latent_space_type':'z'}
logger.info(f'Preparing latent codes.')
logger.info(f' Sample latent codes randomly.')
latent_codes = model.easy_sample(args.num, **kwargs)
total_num = latent_codes.shape[0]
logger.info(f'Generating {total_num} samples.')
results = defaultdict(list)
pbar = tqdm(total=total_num, leave=False)
scores = []
count_double_chin = 0
count=0
if args.double_chin_only:
print('Only generate images that have double chin!')
for latent_codes_batch in model.get_batch_inputs(latent_codes):
count+=1
outputs = model.easy_synthesize(latent_codes_batch,
**kwargs,
generate_style=args.generate_style,
generate_image=args.generate_image)
if args.double_chin_only:
choose = []
key='image'
val=outputs[key]
for image in val:
score = check_double_chin(img=image[:, :, ::-1], model=double_chin_checker)
pbar.update(1)
if (score == 1):
choose.append(True)
scores.append(score)
save_path = os.path.join(args.output_dir, f'{count_double_chin + 50186:06d}.jpg')
cv2.imwrite(save_path, image[:, :, ::-1])
count_double_chin += 1
else:
choose.append(False)
#os.remove(save_path)
for key, val in outputs.items():
if not key == 'image':
if choose[0]:
results[key].append(val)
else:
for key, val in outputs.items():
if key == 'image':
for image in val:
save_path = os.path.join(args.output_dir, f'{pbar.n:06d}.jpg')
cv2.imwrite(save_path, image[:, :, ::-1])
score = check_double_chin(img=save_path, model=double_chin_checker)
scores.append(score)
if (score == 1):
count_double_chin += 1
pbar.update(1)
else:
results[key].append(val)
if 'image' not in outputs:
pbar.update(latent_codes_batch.shape[0])
if pbar.n % 1000 == 0 or pbar.n == total_num:
logger.debug(f' Finish {pbar.n:6d} samples.')
pbar.close()
logger.info(f'Saving results.')
for key, val in results.items():
if key=='s':
for s_i in range(26):
save_path = os.path.join(args.output_dir, f'{key}_{s_i}.npy')
s_latent=np.concatenate([v[s_i] for v in val], axis=0)
print(s_latent.shape)
np.save(save_path,s_latent )
s_mean=s_latent.mean(axis=0)
s_std=s_latent.std(axis=0)
mean_save_path = os.path.join(args.output_dir, f'{key}_{s_i}_mean.npy')
std_save_path = os.path.join(args.output_dir, f'{key}_{s_i}_std.npy')
np.save(mean_save_path, s_mean)
np.save(std_save_path, s_std)
else:
save_path = os.path.join(args.output_dir, f'{key}.npy')
np.save(save_path, np.concatenate(val, axis=0))
print(np.concatenate(val, axis=0).shape)
score_save_path=os.path.join(args.output_dir,'double_chin_scores.npy')
scores_array=np.array(scores)[:,np.newaxis]
print(scores_array.shape)
np.save(score_save_path,scores_array)
print('%d double chin images'%count_double_chin)
double_chin_w=[]
for i in range(len(results['w'])):
if scores_array[i]==1:
double_chin_w.append(results['w'][i])
save_path = os.path.join(args.output_dir, 'double_chin_w.npy')
np.save(save_path, np.concatenate(double_chin_w, axis=0))
if __name__ == '__main__':
main()