def main(args, cfg, save_path='transferred_image.png'):
setup(args, cfg)
inference = Inference(cfg, args.model_path)
postprocess = PostProcess(cfg)
source = Image.open(args.source_path).convert("RGB")
reference_paths = list(Path(args.reference_dir).glob("*"))
np.random.shuffle(reference_paths)
for reference_path in reference_paths:
if not reference_path.is_file():
print(reference_path, "is not a valid file.")
continue
reference = Image.open(reference_path).convert("RGB")
# Transfer the psgan from reference to source.
image, face = inference.transfer(source, reference, with_face=True)
image.save('before.png')
source_crop = source.crop(
(face.left(), face.top(), face.right(), face.bottom()))
image = postprocess(source_crop, image)
image.save(save_path)
if __name__ == '__main__':
args = parse_args()
cfg = get_config(args.config_file)
main(args, cfg)
def main(args, cfg):
# init environment, include logger, dynamic graph, seed, device, train or test mode...
setup(args, cfg)
# build trainer
trainer = Trainer(cfg)
# continue train or evaluate, checkpoint need contain epoch and optimizer info
if args.resume:
trainer.resume(args.resume)
# evaluate or finute, only load generator weights
elif args.load:
trainer.load(args.load)
if args.evaluate_only:
trainer.test()
return
# training, when keyboard interrupt save weights
try:
trainer.train()
except KeyboardInterrupt as e:
trainer.save(trainer.current_epoch)
trainer.close()
if __name__ == '__main__':
args = parse_args()
cfg = get_config(args.config_file, args.opt)
main(args, cfg)