def main():
"""Main function."""
args = parse_args()
work_dir = args.output_dir or f'{args.model_name}_rescore'
logger_name = f'{args.model_name}_rescore_logger'
logger = setup_logger(work_dir, args.logfile_name, logger_name)
logger.info(f'Initializing generator.')
model = build_generator(args.model_name, logger=logger)
logger.info(f'Preparing latent codes.')
if args.num <= 0:
raise ValueError(f'Argument `num` should be specified as a positive '
f'number, but `{args.num}` received!')
latent_codes = model.easy_sample(num=args.num, latent_space_type='z')
latent_codes = model.easy_synthesize(latent_codes=latent_codes,
latent_space_type='z',
generate_style=False,
generate_image=False)
for key, val in latent_codes.items():
np.save(os.path.join(work_dir, f'{key}.npy'), val)
logger.info(f'Initializing predictor.')
predictor = build_predictor(args.predictor_name)
boundaries = parse_boundary_list(args.boundary_list_path)
logger.info(f'========================================')
logger.info(f'Rescoring.')
score_changing = []
for boundary_info, boundary_path in boundaries.items():
logger.info(f'----------------------------------------')
boundary_name, space_type = boundary_info
logger.info(
f'Boundary `{boundary_name}` from {space_type.upper()} space.')
prefix = f'{boundary_name}_{space_type}'
attr_idx = predictor.attribute_name_to_idx[boundary_name]
try:
boundary_file = np.load(boundary_path, allow_pickle=True).item()
boundary = boundary_file['boundary']
except ValueError:
boundary = np.load(boundary_path)
np.save(os.path.join(work_dir, f'{prefix}_boundary.npy'), boundary)
if space_type == 'z':
layerwise_manipulation = False
is_code_layerwise = False
is_boundary_layerwise = False
num_layers = 0
strength = 1.0
else:
layerwise_manipulation = True
is_code_layerwise = True
is_boundary_layerwise = (space_type == 'wp')
num_layers = model.num_layers if args.layerwise_rescoring else 0
if space_type == 'w':
strength = get_layerwise_manipulation_strength(
model.num_layers, model.truncation_psi,
model.truncation_layers)
else:
strength = 1.0
space_type = 'wp'
codes = []
codes.append(latent_codes[space_type][:, np.newaxis])
for l in range(-1, num_layers):
codes.append(
manipulate(latent_codes[space_type],
boundary,
start_distance=2.0,
end_distance=2.0,
step=1,
layerwise_manipulation=layerwise_manipulation,
num_layers=model.num_layers,
manipulate_layers=None if l < 0 else l,
is_code_layerwise=is_code_layerwise,
is_boundary_layerwise=is_boundary_layerwise,
layerwise_manipulation_strength=strength))
codes = np.concatenate(codes, axis=1)
scores = []
for i in tqdm(range(args.num), leave=False):
images = model.easy_synthesize(latent_codes=codes[i],
latent_space_type=space_type,
generate_style=False,
generate_image=True)['image']
scores.append(
predictor.easy_predict(images)['attribute'][:, attr_idx])
scores = np.stack(scores, axis=0)
np.save(os.path.join(work_dir, f'{prefix}_scores.npy'), scores)
delta = scores[:, 1] - scores[:, 0]
delta[delta < 0] = 0
score_changing.append((boundary_name, np.mean(delta)))
if num_layers:
layerwise_score_changing = []
for l in range(num_layers):
delta = scores[:, l + 2] - scores[:, 0]
delta[delta < 0] = 0
layerwise_score_changing.append(
(f'Layer {l:02d}', np.mean(delta)))
layerwise_score_changing.sort(key=lambda x: x[1], reverse=True)
for layer_name, delta_score in layerwise_score_changing:
logger.info(f' {layer_name}: {delta_score:7.4f}')
logger.info(f'----------------------------------------')
logger.info(f'Most relevant semantics:')
score_changing.sort(key=lambda x: x[1], reverse=True)
for boundary_name, delta_score in score_changing:
logger.info(f' {boundary_name.ljust(15)}: {delta_score:7.4f}')
args = parser.parse_args()
set_cuda_devices(args.gpu_id)
name_list = open(args.name_list, "r").readlines()
z, image_stroke, image_mask, label_stroke, label_mask = \
read_data(args.data_dir, name_list)
param_dict = dict(latent_strategy="mixwp",
optimizer='adam',
n_iter=50,
base_lr=0.002)
print(z.shape, image_stroke.shape, image_mask.shape, label_stroke.shape,
label_mask.shape)
G_name = "stylegan2_ffhq"
DIR = "predictors/pretrain/"
G = build_generator(G_name).net
P = build_predictor("face_seg")
SE_full = load_semantic_extractor(f"{DIR}/{G_name}_LSE.pth").cuda()
SE_fewshot = load_semantic_extractor(
f"{DIR}/{G_name}_8shot_LSE.pth").cuda()
proc = lambda x: ((bu(x.detach().cpu(), 256) + 1) / 2)
vizproc = lambda x: (bu(segviz_torch(x.detach().cpu()).unsqueeze(0), 256))
origins, labels, baselines, fewshots, fulls = [], [], [], [], []
for i in tqdm(range(z.shape[0])):
zs = z[i].cuda()
with torch.no_grad():
wp = EditStrategy.z_to_wp(G,
zs,
in_type="zs",
out_type="notrunc-wp")
def main():
"""Main function."""
args = parse_args()
file_name = args.output_file
work_dir = args.output_dir or f'{args.model_name}_synthesis'
logger_name = f'{args.model_name}_synthesis_logger'
logger = setup_logger(work_dir, args.logfile_name, logger_name)
logger.info(f'Initializing generator.')
model = build_generator(args.model_name, logger=logger)
logger.info(f'Preparing latent codes.')
if os.path.isfile(args.latent_codes_path):
logger.info(f' Load latent codes from `{args.latent_codes_path}`.')
latent_codes = np.load(args.latent_codes_path)
latent_codes = model.preprocess(
latent_codes=latent_codes,
latent_space_type=args.latent_space_type)
else:
if args.num <= 0:
raise ValueError(
f'Argument `num` should be specified as a positive '
f'number since the latent code path '
f'`{args.latent_codes_path}` does not exist!')
logger.info(f' Sample latent codes randomly.')
latent_codes = model.easy_sample(
num=args.num, latent_space_type=args.latent_space_type)
total_num = latent_codes.shape[0]
if args.generate_prediction:
logger.info(f'Initializing predictor.')
predictor = build_predictor(args.predictor_name)
if args.generate_html:
viz_size = None if args.viz_size == 0 else args.viz_size
visualizer = HtmlPageVisualizer(num_rows=args.html_row,
num_cols=args.html_col,
grid_size=total_num,
viz_size=viz_size)
logger.info(f'Generating {total_num} samples.')
results = defaultdict(list)
predictions = defaultdict(list)
pbar = tqdm(total=total_num, leave=False)
for inputs in model.get_batch_inputs(latent_codes):
outputs = model.easy_synthesize(
latent_codes=inputs,
latent_space_type=args.latent_space_type,
generate_style=args.generate_style,
generate_image=not args.skip_image)
for key, val in outputs.items():
if key == 'image':
if args.generate_prediction:
pred_outputs = predictor.easy_predict(val)
for pred_key, pred_val in pred_outputs.items():
predictions[pred_key].append(pred_val)
for image in val:
if args.save_raw_synthesis:
dest = os.path.join(work_dir, f'{pbar.n:06d}.jpg')
if file_name != "":
dest = os.path.join(work_dir, file_name)
print('saving image to ', dest)
save_image(dest, image)
if args.generate_html:
row_idx = pbar.n // visualizer.num_cols
col_idx = pbar.n % visualizer.num_cols
visualizer.set_cell(row_idx, col_idx, image=image)
pbar.update(1)
else:
results[key].append(val)
if 'image' not in outputs:
pbar.update(inputs.shape[0])
pbar.close()
logger.info(f'Saving results.')
if args.generate_html:
visualizer.save(os.path.join(work_dir, args.html_name))
for key, val in results.items():
np.save(os.path.join(work_dir, f'{key}.npy'),
np.concatenate(val, axis=0))
if predictions:
if args.predictor_name == 'scene':
# Categories
categories = np.concatenate(predictions['category'], axis=0)
detailed_categories = {
'score': categories,
'name_to_idx': predictor.category_name_to_idx,
'idx_to_name': predictor.category_idx_to_name,
}
np.save(os.path.join(work_dir, 'category.npy'),
detailed_categories)
# Attributes
attributes = np.concatenate(predictions['attribute'], axis=0)
detailed_attributes = {
'score': attributes,
'name_to_idx': predictor.attribute_name_to_idx,
'idx_to_name': predictor.attribute_idx_to_name,
}
np.save(os.path.join(work_dir, 'attribute.npy'),
detailed_attributes)
else:
for key, val in predictions.items():
np.save(os.path.join(work_dir, f'{key}.npy'),
np.concatenate(val, axis=0))
def main():
"""Main function."""
args = parse_args()
set_cuda_devices(args.gpu_id)
work_dir = args.output_dir
os.system("mkdir " + work_dir)
model = build_generator(
args.model_name,
truncation_psi=None if args.truncation < 0 else args.truncation)
if os.path.isfile(args.latent_codes_path):
latent_codes = np.load(args.latent_codes_path)
latent_codes = model.preprocess(
latent_codes=latent_codes,
latent_space_type=args.latent_space_type)
else:
if args.num <= 0:
raise ValueError(
f'Argument `num` should be specified as a positive '
f'number since the latent code path '
f'`{args.latent_codes_path}` does not exist!')
latent_codes = model.easy_sample(
num=args.num, latent_space_type=args.latent_space_type)
total_num = latent_codes.shape[0]
if args.generate_prediction:
predictor = build_predictor(args.predictor_name)
if args.generate_html:
visualizer = HtmlPageVisualizer(num_rows=args.html_row,
num_cols=args.html_col,
grid_size=total_num,
viz_size=args.viz_size)
results = defaultdict(list)
predictions = defaultdict(list)
pbar = tqdm(total=total_num, leave=False)
for inputs in model.get_batch_inputs(latent_codes):
outputs = model.easy_synthesize(
latent_codes=inputs,
latent_space_type=args.latent_space_type,
generate_style=args.generate_style,
generate_image=not args.skip_image)
for key, val in outputs.items():
if key == 'image':
if args.generate_prediction:
pred_outputs = predictor.easy_predict(val)
for pred_key, pred_val in pred_outputs.items():
predictions[pred_key].append(pred_val)
for image in val:
if args.save_raw_synthesis:
save_image(os.path.join(work_dir, f'{pbar.n:06d}.jpg'),
image)
if args.generate_html:
row_idx = pbar.n // visualizer.num_cols
col_idx = pbar.n % visualizer.num_cols
visualizer.set_cell(row_idx,
col_idx,
text=f'Sample {pbar.n:06d}',
image=image)
pbar.update(1)
else:
results[key].append(val)
if 'image' not in outputs:
pbar.update(inputs.shape[0])
pbar.close()
if args.generate_html:
visualizer.save(os.path.join(work_dir, args.html_name))
for key, val in results.items():
np.save(os.path.join(work_dir, f'{key}.npy'),
np.concatenate(val, axis=0))
if predictions:
print(len(predictions))
predictor.save(predictions, work_dir)