def run_toonify_bootstrapping(self, input, num_iterations,
display_intermediate_results):
# load ffhq model
print("Loading faces model...")
ckpt = self.checkpoints["faces"]
opts = ckpt['opts']
opts['checkpoint_path'] = self.model_paths["faces"]
opts = Namespace(**opts)
net_ffhq = pSp(opts)
net_ffhq.eval()
net_ffhq.cuda()
print("Done!")
# load toonify model
print("Loading toonify model...")
ckpt = self.checkpoints["toonify"]
opts = ckpt['opts']
opts['checkpoint_path'] = self.model_paths["toonify"]
opts = Namespace(**opts)
net_toonify = pSp(opts)
net_toonify.eval()
net_toonify.cuda()
print("Done!")
# define some arguments
opts.n_iters_per_batch = num_iterations
opts.resize_outputs = False
# load, align, and preprocess image
print('Aligning image...')
input_image = self.run_alignment(str(input))
print('Done!')
transformed_image = self.default_transforms(input_image)
# run inference
print("Running inference...")
with torch.no_grad():
start = time.time()
avg_image = self.get_avg_image(net_ffhq, encoding_type="faces")
result_batch = encoder_bootstrapping_inference.run_on_batch(
transformed_image.unsqueeze(0).cuda(), net_ffhq, net_toonify,
opts, avg_image)
total_time = time.time() - start
print(f"Finished inference in {total_time} seconds!")
# post-processing
print("Preparing result...")
resize_amount = (1024, 1024)
res = self.get_final_output(result_batch, resize_amount,
display_intermediate_results, opts)
# display output
out_path = Path(tempfile.mkdtemp()) / "output.png"
imageio.imwrite(str(out_path), res)
return out_path
def __init__(self, opts):
self.opts = opts
self.global_step = 0
self.device = 'cuda:0' # TODO: Allow multiple GPU? currently using CUDA_VISIBLE_DEVICES
self.opts.device = self.device
if self.opts.use_wandb:
from utils.wandb_utils import WBLogger
self.wb_logger = WBLogger(self.opts)
# Initialize network
self.net = pSp(self.opts).to(self.device)
# Estimate latent_avg via dense sampling if latent_avg is not available
if self.net.latent_avg is None:
self.net.latent_avg = self.net.decoder.mean_latent(int(1e5))[0].detach()
# Initialize loss
if self.opts.id_lambda > 0 and self.opts.moco_lambda > 0:
raise ValueError('Both ID and MoCo loss have lambdas > 0! Please select only one to have non-zero lambda!')
self.mse_loss = nn.MSELoss().to(self.device).eval()
if self.opts.lpips_lambda > 0:
self.lpips_loss = LPIPS(net_type='alex').to(self.device).eval()
if self.opts.id_lambda > 0:
self.id_loss = id_loss.IDLoss().to(self.device).eval()
if self.opts.w_norm_lambda > 0:
self.w_norm_loss = w_norm.WNormLoss(start_from_latent_avg=self.opts.start_from_latent_avg)
if self.opts.moco_lambda > 0:
self.moco_loss = moco_loss.MocoLoss().to(self.device).eval()
# Initialize optimizer
self.optimizer = self.configure_optimizers()
# Initialize dataset
self.train_dataset, self.test_dataset = self.configure_datasets()
self.train_dataloader = DataLoader(self.train_dataset,
batch_size=self.opts.batch_size,
shuffle=True,
num_workers=int(self.opts.workers),
drop_last=True)
self.test_dataloader = DataLoader(self.test_dataset,
batch_size=self.opts.test_batch_size,
shuffle=False,
num_workers=int(self.opts.test_workers),
drop_last=True)
# Initialize logger
log_dir = os.path.join(opts.exp_dir, 'logs')
os.makedirs(log_dir, exist_ok=True)
self.logger = SummaryWriter(log_dir=log_dir)
# Initialize checkpoint dir
self.checkpoint_dir = os.path.join(opts.exp_dir, 'checkpoints')
os.makedirs(self.checkpoint_dir, exist_ok=True)
self.best_val_loss = None
if self.opts.save_interval is None:
self.opts.save_interval = self.opts.max_steps
def load_model(model_path='pretrained_models/e4e_ffhq_encode.pt'):
ckpt = torch.load(model_path, map_location='cpu')
opts = ckpt['opts']
opts['checkpoint_path'] = model_path
opts = Namespace(**opts)
net = pSp(opts)
net.eval()
net.cuda()
return net
def predict(self, model, image):
opts = self.opts[model]
opts = Namespace(**opts)
pprint.pprint(opts)
net = pSp(opts)
net.eval()
net.cuda()
print('Model successfully loaded!')
original_image = Image.open(str(image))
if opts.label_nc == 0:
original_image = original_image.convert("RGB")
else:
original_image = original_image.convert("L")
original_image.resize(
(self.opts[model]['output_size'], self.opts[model]['output_size']))
# Align Image
if model not in ["celebs_sketch_to_face", "celebs_seg_to_face"]:
input_image = self.run_alignment(str(image))
else:
input_image = original_image
img_transforms = self.transforms[model]
transformed_image = img_transforms(input_image)
if model in ["celebs_sketch_to_face", "celebs_seg_to_face"]:
latent_mask = [8, 9, 10, 11, 12, 13, 14, 15, 16, 17]
else:
latent_mask = None
with torch.no_grad():
result_image = run_on_batch(transformed_image.unsqueeze(0), net,
latent_mask)[0]
input_vis_image = log_input_image(transformed_image, opts)
output_image = tensor2im(result_image)
if model == "celebs_super_resolution":
res = np.concatenate([
np.array(
input_vis_image.resize((self.opts[model]['output_size'],
self.opts[model]['output_size']))),
np.array(
output_image.resize((self.opts[model]['output_size'],
self.opts[model]['output_size'])))
],
axis=1)
else:
res = np.array(
output_image.resize((self.opts[model]['output_size'],
self.opts[model]['output_size'])))
out_path = Path(tempfile.mkdtemp()) / "out.png"
Image.fromarray(np.array(res)).save(str(out_path))
return out_path
def run():
# Load configs
test_opts = TestOptions().parse()
ckpt = torch.load(test_opts.checkpoint_path, map_location='cpu')
opts = ckpt['opts']
opts.update(vars(test_opts))
if 'learn_in_w' not in opts:
opts['learn_in_w'] = False
# opts['device'] = 'cpu'
opts = Namespace(**opts)
# Init model
model = pSp(opts)
model.eval()
model.cuda()
with torch.no_grad():
# Use this an input trace to serialize the model
# input_shape = (1, 256, 256)
# x = torch.randn(1, 1, 256, 256, requires_grad=True)
x = torch.randn(1, 1, 256, 256, requires_grad=True).cuda()
# TensorRT
# Export the model to an ONNX file
# orward(self, x, resize=True, latent_mask=None, input_code=False, randomize_noise=True,
# inject_latent=None, return_latents=False, alpha=None):
# dummy_input = (x, True, None, False, False, None, False, None)
dummy_input = x
model_onnx_path = "psp_sketch_model_default.onnx"
# tensorRT
# model_rt_path = "./psp_sketch_model_tensorRT.pth"
# convert to TensorRT feeding sample data as input
# model_trt = torch2trt(model, [x])
# torch.save(model_trt.state_dict(), model_rt_path)
# print("Export of torch_model.tensorRT complete!")
output = torch_onnx.export(
model,
dummy_input,
model_onnx_path,
input_names=["x"],
output_names=["images"],
opset_version=13, # the ONNX version to export the model to
export_params=True,
verbose=True)
# output = torch_onnx.export(model,
# dummy_input,
# model_onnx_path,
# input_names=["x","resize","latent_mask","input_code","randomize_noise","inject_latent","return_latents","alpha"],
# output_names=["images"],
# verbose=True)
print("Export of torch_model.onnx complete!")
def run_default_encoding(self, input, encoding_type, num_iterations,
display_intermediate_results):
# load model
print(f'Loading {encoding_type} model...')
ckpt = self.checkpoints[encoding_type]
opts = ckpt['opts']
opts['checkpoint_path'] = self.model_paths[encoding_type]
opts = Namespace(**opts)
net = e4e(opts) if encoding_type == "horses" else pSp(opts)
net.eval()
net.cuda()
print('Done!')
# define some arguments
opts.n_iters_per_batch = num_iterations
opts.resize_outputs = False
# define transforms
image_transforms = self.cars_transforms if encoding_type == "cars" else self.default_transforms
# if working on faces load and align the image
if encoding_type == "faces":
print('Aligning image...')
input_image = self.run_alignment(str(input))
print('Done!')
# otherwise simply load the image
else:
input_image = Image.open(str(input)).convert("RGB")
# preprocess image
transformed_image = image_transforms(input_image)
# run inference
print("Running inference...")
with torch.no_grad():
start = time.time()
avg_image = self.get_avg_image(net, encoding_type)
result_batch, result_latents = run_on_batch(
transformed_image.unsqueeze(0).cuda(), net, opts, avg_image)
total_time = time.time() - start
print(f"Finished inference in {total_time} seconds!")
# post-processing
print("Preparing result...")
resize_amount = (512, 384) if encoding_type == "cars_encode" else (
opts.output_size, opts.output_size)
res = self.get_final_output(result_batch, resize_amount,
display_intermediate_results, opts)
# display output
out_path = Path(tempfile.mkdtemp()) / "output.png"
imageio.imwrite(str(out_path), res)
return out_path
def __init__(self, opts):
self.opts = opts
self.global_step = 0
self.device = "cuda:0" # TODO: Allow multiple GPU? currently using CUDA_VISIBLE_DEVICES
self.opts.device = self.device
# Initialize network
self.net = pSp(self.opts).to(self.device)
# Initialize loss
if self.opts.lpips_lambda > 0:
self.lpips_loss = LPIPS(net_type="alex").to(self.device).eval()
if self.opts.id_lambda > 0 or self.opts.id_lambda_input > 0:
self.id_loss = id_loss.IDLoss().to(self.device).eval()
if self.opts.w_norm_lambda > 0:
self.w_norm_loss = w_norm.WNormLoss(
start_from_latent_avg=self.opts.start_from_latent_avg
)
self.mse_loss = nn.MSELoss().to(self.device).eval()
# Initialize optimizer
self.optimizer = self.configure_optimizers()
# Initialize dataset
self.train_dataset, self.test_dataset = self.configure_datasets()
self.train_dataloader = DataLoader(
self.train_dataset,
batch_size=self.opts.batch_size,
shuffle=True,
num_workers=int(self.opts.workers),
drop_last=True,
)
self.test_dataloader = DataLoader(
self.test_dataset,
batch_size=self.opts.test_batch_size,
shuffle=False,
num_workers=int(self.opts.test_workers),
drop_last=False,
)
# Initialize logger
log_dir = os.path.join(opts.exp_dir, "logs")
os.makedirs(log_dir, exist_ok=True)
self.logger = SummaryWriter(log_dir=log_dir)
# Initialize checkpoint dir
self.checkpoint_dir = os.path.join(opts.exp_dir, "checkpoints")
os.makedirs(self.checkpoint_dir, exist_ok=True)
self.best_val_loss = None
if self.opts.save_interval is None:
self.opts.save_interval = self.opts.max_steps
def setup_model(checkpoint_path, device='cuda'):
ckpt = torch.load(checkpoint_path, map_location='cpu')
opts = ckpt['opts']
opts['checkpoint_path'] = checkpoint_path
opts['device'] = device
opts = argparse.Namespace(**opts)
net = pSp(opts)
net.eval()
net = net.to(device)
return net, opts
def __init__(self, opts):
self.opts = opts
self.global_step = 0
self.device = 'cuda'
self.opts.device = self.device
# Initialize network
self.net = pSp(self.opts).to(self.device)
# Initialize loss
self.mse_loss = nn.MSELoss().to(self.device).eval()
if self.opts.lpips_lambda > 0:
self.lpips_loss = LPIPS(net_type='alex').to(self.device).eval()
if self.opts.id_lambda > 0:
self.id_loss = id_loss.IDLoss().to(self.device).eval()
if self.opts.w_norm_lambda > 0:
self.w_norm_loss = w_norm.WNormLoss(opts=self.opts)
if self.opts.aging_lambda > 0:
self.aging_loss = AgingLoss(self.opts)
# Initialize optimizer
self.optimizer = self.configure_optimizers()
# Initialize dataset
self.train_dataset, self.test_dataset = self.configure_datasets()
self.train_dataloader = DataLoader(self.train_dataset,
batch_size=self.opts.batch_size,
shuffle=True,
num_workers=int(self.opts.workers),
drop_last=True)
self.test_dataloader = DataLoader(self.test_dataset,
batch_size=self.opts.test_batch_size,
shuffle=False,
num_workers=int(
self.opts.test_workers),
drop_last=True)
self.age_transformer = AgeTransformer(target_age=self.opts.target_age)
# Initialize logger
log_dir = os.path.join(opts.exp_dir, 'logs')
os.makedirs(log_dir, exist_ok=True)
self.logger = SummaryWriter(log_dir=log_dir)
# Initialize checkpoint dir
self.checkpoint_dir = os.path.join(opts.exp_dir, 'checkpoints')
os.makedirs(self.checkpoint_dir, exist_ok=True)
self.best_val_loss = None
if self.opts.save_interval is None:
self.opts.save_interval = self.opts.max_steps
def create_model_from_path(model_path):
ckpt = torch.load(model_path, map_location='cpu')
opts = ckpt['opts']
#pprint.pprint(opts)
# update the training options
opts['checkpoint_path'] = model_path
if 'learn_in_w' not in opts:
opts['learn_in_w'] = False
opts = Namespace(**opts)
net = pSp(opts)
net.eval()
net.cuda()
print('Model successfully loaded!')
return net, opts
def model_init(opts, seed=42):
print(f"Model init time = {time.time()}")
# CUDNN SETTING
# random.seed(seed)
# np.random.seed(seed)
# torch.manual_seed(seed)
# torch.cuda.manual_seed(seed)
torch.backends.cudnn.benchmark = True
# update test options with options used during training
ckpt = torch.load(opts.checkpoint_path, map_location='cpu')
net = pSp(opts)
net.eval()
net.cuda()
return net
def build_model():
# Load configs
test_opts = TestOptions().parse()
ckpt = torch.load(test_opts.checkpoint_path, map_location='cpu')
opts = ckpt['opts']
opts.update(vars(test_opts))
if 'learn_in_w' not in opts:
opts['learn_in_w'] = False
# opts['device'] = 'cpu'
opts = Namespace(**opts)
# Init model
model = pSp(opts)
# model.eval()
return model
def setup_model(checkpoint_path, device='cuda'):
ckpt = torch.load(checkpoint_path, map_location='cpu')
opts = ckpt['opts']
is_cars = 'car' in opts['dataset_type']
is_faces = 'ffhq' in opts['dataset_type']
if is_faces:
opts['stylegan_size'] = 1024
elif is_cars:
opts['stylegan_size'] = 512
else:
opts['stylegan_size'] = 256
opts['checkpoint_path'] = checkpoint_path
opts = argparse.Namespace(**opts)
net = pSp(opts)
net.eval()
net = net.to(device)
return net, opts
def load_model(self,
model_path=REPO_PATH +
"/pretrained_models/psp_ffhq_encode.pt"):
ckpt = torch.load(model_path, map_location='cpu')
opts = ckpt['opts']
# pprint.pprint(opts)
# update the training options
opts['checkpoint_path'] = model_path
if 'learn_in_w' not in opts:
opts['learn_in_w'] = False
opts = Namespace(**opts)
net = pSp(opts)
net.eval()
net.cuda()
print('Model successfully loaded!')
self.net = net
self.opts = opts
def predict(self, image, target_age="default"):
net = pSp(self.opts)
net.eval()
if torch.cuda.is_available():
net.cuda()
# align image
aligned_image = run_alignment(str(image))
aligned_image.resize((256, 256))
input_image = self.transform(aligned_image)
if target_age == "default":
target_ages = [0, 10, 20, 30, 40, 50, 60, 70, 80, 90, 100]
age_transformers = [
AgeTransformer(target_age=age) for age in target_ages
]
else:
age_transformers = [AgeTransformer(target_age=target_age)]
results = np.array(aligned_image.resize((1024, 1024)))
all_imgs = []
for age_transformer in age_transformers:
print(f"Running on target age: {age_transformer.target_age}")
with torch.no_grad():
input_image_age = [
age_transformer(input_image.cpu()).to("cuda")
]
input_image_age = torch.stack(input_image_age)
result_tensor = run_on_batch(input_image_age, net)[0]
result_image = tensor2im(result_tensor)
all_imgs.append(result_image)
results = np.concatenate([results, result_image], axis=1)
if target_age == "default":
out_path = Path(tempfile.mkdtemp()) / "output.gif"
imageio.mimwrite(str(out_path), all_imgs, duration=0.3)
else:
out_path = Path(tempfile.mkdtemp()) / "output.png"
imageio.imwrite(str(out_path), all_imgs[0])
return out_path
def run():
test_opts = TestOptions().parse()
out_path_coupled = os.path.join(test_opts.exp_dir, 'inference_coupled')
os.makedirs(out_path_coupled, exist_ok=True)
# update test options with options used during training
ckpt = torch.load(test_opts.checkpoint_path, map_location='cpu')
opts = ckpt['opts']
opts.update(vars(test_opts))
opts = Namespace(**opts)
net = pSp(opts)
net.eval()
net.cuda()
print('Loading dataset for {}'.format(opts.dataset_type))
dataset_args = data_configs.DATASETS[opts.dataset_type]
transforms_dict = dataset_args['transforms'](opts).get_transforms()
dataset = InferenceDataset(root=opts.data_path,
transform=transforms_dict['transform_inference'],
opts=opts)
dataloader = DataLoader(dataset,
batch_size=opts.test_batch_size,
shuffle=False,
num_workers=int(opts.test_workers),
drop_last=False)
if opts.n_images is None:
opts.n_images = len(dataset)
# get the image corresponding to the latent average
avg_image = net(net.latent_avg.unsqueeze(0),
input_code=True,
randomize_noise=False,
return_latents=False,
average_code=True)[0]
avg_image = avg_image.to('cuda').float().detach()
if opts.dataset_type == "cars_encode":
avg_image = avg_image[:, 32:224, :]
tensor2im(avg_image).save(os.path.join(opts.exp_dir, 'avg_image.jpg'))
if opts.dataset_type == "cars_encode":
resize_amount = (256, 192) if opts.resize_outputs else (512, 384)
else:
resize_amount = (256, 256) if opts.resize_outputs else (opts.output_size, opts.output_size)
global_i = 0
global_time = []
for input_batch in tqdm(dataloader):
if global_i >= opts.n_images:
break
with torch.no_grad():
input_cuda = input_batch.cuda().float()
tic = time.time()
result_batch, result_latents = run_on_batch(input_cuda, net, opts, avg_image)
toc = time.time()
global_time.append(toc - tic)
for i in range(input_batch.shape[0]):
results = [tensor2im(result_batch[i][iter_idx]) for iter_idx in range(opts.n_iters_per_batch)]
im_path = dataset.paths[global_i]
# save step-by-step results side-by-side
input_im = tensor2im(input_batch[i])
res = np.array(results[0].resize(resize_amount))
for idx, result in enumerate(results[1:]):
res = np.concatenate([res, np.array(result.resize(resize_amount))], axis=1)
res = np.concatenate([res, input_im.resize(resize_amount)], axis=1)
Image.fromarray(res).save(os.path.join(out_path_coupled, os.path.basename(im_path)))
global_i += 1
stats_path = os.path.join(opts.exp_dir, 'stats.txt')
result_str = 'Runtime {:.4f}+-{:.4f}'.format(np.mean(global_time), np.std(global_time))
print(result_str)
with open(stats_path, 'w') as f:
f.write(result_str)
def __init__(self, opts):
self.opts = opts
self.global_step = 0
self.device = 'cuda:0' # TODO: Allow multiple GPU? currently using CUDA_VISIBLE_DEVICES
self.opts.device = self.device
# Initialize network
self.net = pSp(self.opts).to(self.device)
# get the image corresponding to the latent average
self.avg_image = self.net(self.net.latent_avg.unsqueeze(0),
input_code=True,
randomize_noise=False,
return_latents=False,
average_code=True)[0]
self.avg_image = self.avg_image.to(self.device).float().detach()
if self.opts.dataset_type == "cars_encode":
self.avg_image = self.avg_image[:, 32:224, :]
common.tensor2im(self.avg_image).save(
os.path.join(self.opts.exp_dir, 'avg_image.jpg'))
# Initialize loss
if self.opts.id_lambda > 0 and self.opts.moco_lambda > 0:
raise ValueError(
'Both ID and MoCo loss have lambdas > 0! Please select only one to have non-zero lambda!'
)
self.mse_loss = nn.MSELoss().to(self.device).eval()
if self.opts.lpips_lambda > 0:
self.lpips_loss = LPIPS(net_type='alex').to(self.device).eval()
if self.opts.id_lambda > 0:
self.id_loss = id_loss.IDLoss().to(self.device).eval()
if self.opts.w_norm_lambda > 0:
self.w_norm_loss = w_norm.WNormLoss(
start_from_latent_avg=self.opts.start_from_latent_avg)
if self.opts.moco_lambda > 0:
self.moco_loss = moco_loss.MocoLoss()
# Initialize optimizer
self.optimizer = self.configure_optimizers()
# Initialize dataset
self.train_dataset, self.test_dataset = self.configure_datasets()
self.train_dataloader = DataLoader(self.train_dataset,
batch_size=self.opts.batch_size,
shuffle=True,
num_workers=int(self.opts.workers),
drop_last=True)
self.test_dataloader = DataLoader(self.test_dataset,
batch_size=self.opts.test_batch_size,
shuffle=False,
num_workers=int(
self.opts.test_workers),
drop_last=True)
# Initialize logger
log_dir = os.path.join(opts.exp_dir, 'logs')
os.makedirs(log_dir, exist_ok=True)
self.logger = SummaryWriter(log_dir=log_dir)
# Initialize checkpoint dir
self.checkpoint_dir = os.path.join(opts.exp_dir, 'checkpoints')
os.makedirs(self.checkpoint_dir, exist_ok=True)
self.best_val_loss = None
if self.opts.save_interval is None:
self.opts.save_interval = self.opts.max_steps
def run():
test_opts = TestOptions().parse()
if test_opts.resize_factors is not None:
factors = test_opts.resize_factors.split(',')
assert len(
factors
) == 1, "When running inference, please provide a single downsampling factor!"
mixed_path_results = os.path.join(
test_opts.exp_dir, 'style_mixing',
'downsampling_{}'.format(test_opts.resize_factors))
else:
mixed_path_results = os.path.join(test_opts.exp_dir, 'style_mixing')
os.makedirs(mixed_path_results, exist_ok=True)
# update test options with options used during training
ckpt = torch.load(test_opts.checkpoint_path, map_location='cpu')
opts = ckpt['opts']
opts.update(vars(test_opts))
if 'learn_in_w' not in opts:
opts['learn_in_w'] = False
opts = Namespace(**opts)
net = pSp(opts)
net.eval()
net.cuda()
print('Loading dataset for {}'.format(opts.dataset_type))
dataset_args = data_configs.DATASETS[opts.dataset_type]
transforms_dict = dataset_args['transforms'](opts).get_transforms()
dataset = InferenceDataset(
root=opts.data_path,
transform=transforms_dict['transform_inference'],
opts=opts)
dataloader = DataLoader(dataset,
batch_size=opts.test_batch_size,
shuffle=False,
num_workers=int(opts.test_workers),
drop_last=True)
latent_mask = [int(l) for l in opts.latent_mask.split(",")]
if opts.n_images is None:
opts.n_images = len(dataset)
global_i = 0
for input_batch in tqdm(dataloader):
if global_i > opts.n_images:
break
with torch.no_grad():
input_batch = input_batch.cuda()
for image_idx, input_image in enumerate(input_batch):
# generate random vectors to inject into input image
vecs_to_inject = np.random.randn(opts.n_outputs_to_generate,
512).astype('float32')
multi_modal_outputs = []
for vec_to_inject in vecs_to_inject:
cur_vec = torch.from_numpy(vec_to_inject).unsqueeze(0).to(
"cuda")
# get latent vector to inject into our input image
_, latent_to_inject = net(cur_vec,
input_code=True,
return_latents=True)
# get output image with injected style vector
res = net(input_image.unsqueeze(0).to("cuda").float(),
latent_mask=latent_mask,
inject_latent=latent_to_inject,
alpha=opts.mix_alpha)
multi_modal_outputs.append(res[0])
# visualize multi modal outputs
input_im_path = dataset.paths[global_i]
image = input_batch[image_idx]
input_image = log_input_image(image, opts)
res = np.array(input_image.resize((256, 256)))
for output in multi_modal_outputs:
output = tensor2im(output)
res = np.concatenate(
[res, np.array(output.resize((256, 256)))], axis=1)
Image.fromarray(res).save(
os.path.join(mixed_path_results,
os.path.basename(input_im_path)))
global_i += 1
def run():
test_opts = TestOptions().parse()
if test_opts.resize_factors is not None:
assert len(
test_opts.resize_factors.split(',')
) == 1, "When running inference, provide a single downsampling factor!"
out_path_results = os.path.join(
test_opts.exp_dir, 'inference_results',
'downsampling_{}'.format(test_opts.resize_factors))
out_path_coupled = os.path.join(
test_opts.exp_dir, 'inference_coupled',
'downsampling_{}'.format(test_opts.resize_factors))
else:
out_path_results = os.path.join(test_opts.exp_dir, 'inference_results')
out_path_coupled = os.path.join(test_opts.exp_dir, 'inference_coupled')
os.makedirs(out_path_results, exist_ok=True)
os.makedirs(out_path_coupled, exist_ok=True)
# update test options with options used during training
ckpt = torch.load(test_opts.checkpoint_path, map_location='cpu')
opts = ckpt['opts']
opts.update(vars(test_opts))
if 'learn_in_w' not in opts:
opts['learn_in_w'] = False
opts = Namespace(**opts)
net = pSp(opts)
net.eval()
net.cuda()
print('Loading dataset for {}'.format(opts.dataset_type))
dataset_args = data_configs.DATASETS[opts.dataset_type]
transforms_dict = dataset_args['transforms'](opts).get_transforms()
dataset = InferenceDataset(
root=opts.data_path,
transform=transforms_dict['transform_inference'],
opts=opts)
dataloader = DataLoader(dataset,
batch_size=opts.test_batch_size,
shuffle=False,
num_workers=int(opts.test_workers),
drop_last=True)
if opts.n_images is None:
opts.n_images = len(dataset)
global_i = 0
global_time = []
for input_batch in tqdm(dataloader):
if global_i >= opts.n_images:
break
with torch.no_grad():
input_cuda = input_batch.cuda().float()
tic = time.time()
result_batch = run_on_batch(input_cuda, net, opts)
toc = time.time()
global_time.append(toc - tic)
for i in range(opts.test_batch_size):
result = tensor2im(result_batch[i])
im_path = dataset.paths[global_i]
if opts.couple_outputs or global_i % 100 == 0:
input_im = log_input_image(input_batch[i], opts)
resize_amount = (256, 256) if opts.resize_outputs else (1024,
1024)
if opts.resize_factors is not None:
# for super resolution, save the original, down-sampled, and output
source = Image.open(im_path)
res = np.concatenate([
np.array(source.resize(resize_amount)),
np.array(
input_im.resize(resize_amount,
resample=Image.NEAREST)),
np.array(result.resize(resize_amount))
],
axis=1)
else:
# otherwise, save the original and output
res = np.concatenate([
np.array(input_im.resize(resize_amount)),
np.array(result.resize(resize_amount))
],
axis=1)
Image.fromarray(res).save(
os.path.join(out_path_coupled, os.path.basename(im_path)))
im_save_path = os.path.join(out_path_results,
os.path.basename(im_path))
Image.fromarray(np.array(result)).save(im_save_path)
global_i += 1
stats_path = os.path.join(opts.exp_dir, 'stats.txt')
result_str = 'Runtime {:.4f}+-{:.4f}'.format(np.mean(global_time),
np.std(global_time))
print(result_str)
with open(stats_path, 'w') as f:
f.write(result_str)
def run():
test_opts = TestOptions().parse()
out_path_results = os.path.join(test_opts.exp_dir,
'reference_guided_inference')
os.makedirs(out_path_results, exist_ok=True)
# update test options with options used during training
ckpt = torch.load(test_opts.checkpoint_path, map_location='cpu')
opts = ckpt['opts']
opts.update(vars(test_opts))
opts = Namespace(**opts)
net = pSp(opts)
net.eval()
net.cuda()
age_transformers = [
AgeTransformer(target_age=age) for age in opts.target_age.split(',')
]
print(f'Loading dataset for {opts.dataset_type}')
dataset_args = data_configs.DATASETS[opts.dataset_type]
transforms_dict = dataset_args['transforms'](opts).get_transforms()
source_dataset = InferenceDataset(
root=opts.data_path,
transform=transforms_dict['transform_inference'],
opts=opts)
source_dataloader = DataLoader(source_dataset,
batch_size=opts.test_batch_size,
shuffle=False,
num_workers=int(opts.test_workers),
drop_last=False)
ref_dataset = InferenceDataset(
paths_list=opts.ref_images_paths_file,
transform=transforms_dict['transform_inference'],
opts=opts)
ref_dataloader = DataLoader(ref_dataset,
batch_size=1,
shuffle=False,
num_workers=1,
drop_last=False)
if opts.n_images is None:
opts.n_images = len(source_dataset)
for age_transformer in age_transformers:
target_age = age_transformer.target_age
print(f"Running on target age: {target_age}")
age_save_path = os.path.join(out_path_results, str(target_age))
os.makedirs(age_save_path, exist_ok=True)
global_i = 0
for i, source_batch in enumerate(tqdm(source_dataloader)):
if global_i >= opts.n_images:
break
results_per_source = {idx: [] for idx in range(len(source_batch))}
with torch.no_grad():
for ref_batch in ref_dataloader:
source_batch = source_batch.cuda().float()
ref_batch = ref_batch.cuda().float()
source_input_age_batch = [
age_transformer(img.cpu()).to('cuda')
for img in source_batch
]
source_input_age_batch = torch.stack(
source_input_age_batch)
# compute w+ of ref images to be injected for style-mixing
ref_latents = net.pretrained_encoder(
ref_batch) + net.latent_avg
# run age transformation on source images with style-mixing
res_batch_mixed = run_on_batch(
source_input_age_batch,
net,
opts,
latent_to_inject=ref_latents)
# store results
for idx in range(len(source_batch)):
results_per_source[idx].append(
[ref_batch[0], res_batch_mixed[idx]])
# save results
resize_amount = (256, 256) if opts.resize_outputs else (1024,
1024)
for image_idx, image_results in results_per_source.items():
input_im_path = source_dataset.paths[global_i]
image = source_batch[image_idx]
input_image = log_image(image, opts)
# initialize results image
ref_inputs = np.zeros_like(
input_image.resize(resize_amount))
mixing_results = np.array(
input_image.resize(resize_amount))
for ref_idx in range(len(image_results)):
ref_input, mixing_result = image_results[ref_idx]
ref_input = log_image(ref_input, opts)
mixing_result = log_image(mixing_result, opts)
# append current results
ref_inputs = np.concatenate([
ref_inputs,
np.array(ref_input.resize(resize_amount))
],
axis=1)
mixing_results = np.concatenate([
mixing_results,
np.array(mixing_result.resize(resize_amount))
],
axis=1)
res = np.concatenate([ref_inputs, mixing_results], axis=0)
save_path = os.path.join(age_save_path,
os.path.basename(input_im_path))
Image.fromarray(res).save(save_path)
global_i += 1
def __init__(self, opts, prev_train_checkpoint=None):
self.opts = opts
self.global_step = 0
self.device = 'cuda:0'
self.opts.device = self.device
# Initialize network
self.net = pSp(self.opts).to(self.device)
# Initialize loss
if self.opts.lpips_lambda > 0:
self.lpips_loss = LPIPS(net_type=self.opts.lpips_type).to(
self.device).eval()
if self.opts.id_lambda > 0:
if 'ffhq' in self.opts.dataset_type or 'celeb' in self.opts.dataset_type:
self.id_loss = id_loss.IDLoss().to(self.device).eval()
else:
self.id_loss = moco_loss.MocoLoss(opts).to(self.device).eval()
self.mse_loss = nn.MSELoss().to(self.device).eval()
# Initialize optimizer
self.optimizer = self.configure_optimizers()
# Initialize discriminator
if self.opts.w_discriminator_lambda > 0:
self.discriminator = LatentCodesDiscriminator(512,
4).to(self.device)
self.discriminator_optimizer = torch.optim.Adam(
list(self.discriminator.parameters()),
lr=opts.w_discriminator_lr)
self.real_w_pool = LatentCodesPool(self.opts.w_pool_size)
self.fake_w_pool = LatentCodesPool(self.opts.w_pool_size)
# Initialize dataset
self.train_dataset, self.test_dataset = self.configure_datasets()
self.train_dataloader = DataLoader(self.train_dataset,
batch_size=self.opts.batch_size,
shuffle=True,
num_workers=int(self.opts.workers),
drop_last=True)
self.test_dataloader = DataLoader(self.test_dataset,
batch_size=self.opts.test_batch_size,
shuffle=False,
num_workers=int(
self.opts.test_workers),
drop_last=True)
# Initialize logger
log_dir = os.path.join(opts.exp_dir, 'logs')
os.makedirs(log_dir, exist_ok=True)
self.logger = SummaryWriter(log_dir=log_dir)
# Initialize checkpoint dir
self.checkpoint_dir = os.path.join(opts.exp_dir, 'checkpoints')
os.makedirs(self.checkpoint_dir, exist_ok=True)
self.best_val_loss = None
if self.opts.save_interval is None:
self.opts.save_interval = self.opts.max_steps
if prev_train_checkpoint is not None:
self.load_from_train_checkpoint(prev_train_checkpoint)
prev_train_checkpoint = None
def run():
test_opts = TestOptions().parse()
out_path_results = os.path.join(test_opts.exp_dir, 'inference_side_by_side')
os.makedirs(out_path_results, exist_ok=True)
# update test options with options used during training
ckpt = torch.load(test_opts.checkpoint_path, map_location='cpu')
opts = ckpt['opts']
opts.update(vars(test_opts))
opts = Namespace(**opts)
net = pSp(opts)
net.eval()
net.cuda()
age_transformers = [AgeTransformer(target_age=age) for age in opts.target_age.split(',')]
print(f'Loading dataset for {opts.dataset_type}')
dataset_args = data_configs.DATASETS[opts.dataset_type]
transforms_dict = dataset_args['transforms'](opts).get_transforms()
dataset = InferenceDataset(root=opts.data_path,
transform=transforms_dict['transform_inference'],
opts=opts,
return_path=True)
dataloader = DataLoader(dataset,
batch_size=opts.test_batch_size,
shuffle=False,
num_workers=int(opts.test_workers),
drop_last=False)
if opts.n_images is None:
opts.n_images = len(dataset)
global_time = []
global_i = 0
for input_batch, image_paths in tqdm(dataloader):
if global_i >= opts.n_images:
break
batch_results = {}
for idx, age_transformer in enumerate(age_transformers):
with torch.no_grad():
input_age_batch = [age_transformer(img.cpu()).to('cuda') for img in input_batch]
input_age_batch = torch.stack(input_age_batch)
input_cuda = input_age_batch.cuda().float()
tic = time.time()
result_batch = run_on_batch(input_cuda, net, opts)
toc = time.time()
global_time.append(toc - tic)
resize_amount = (256, 256) if opts.resize_outputs else (1024, 1024)
for i in range(len(input_batch)):
result = tensor2im(result_batch[i])
im_path = image_paths[i]
input_im = log_image(input_batch[i], opts)
if im_path not in batch_results.keys():
batch_results[im_path] = np.array(input_im.resize(resize_amount))
batch_results[im_path] = np.concatenate([batch_results[im_path],
np.array(result.resize(resize_amount))],
axis=1)
for im_path, res in batch_results.items():
image_name = os.path.basename(im_path)
im_save_path = os.path.join(out_path_results, image_name)
Image.fromarray(np.array(res)).save(im_save_path)
global_i += 1
raise ValueError("Pretrained model was unable to be downlaoded correctly!")
## Step 4: Load Pretrained Model
model_path = EXPERIMENT_ARGS['model_path']
ckpt = torch.load(model_path, map_location='cpu')
opts = ckpt['opts']
pprint.pprint(opts)
# update the training options
opts['checkpoint_path'] = model_path
if 'learn_in_w' not in opts:
opts['learn_in_w'] = False
opts = Namespace(**opts)
net = pSp(opts)
net.eval()
net.cuda()
print('Model successfully loaded!')
## Step 5: Visualize Input
# Alignment face
def run_alignment(image_path):
predictor = dlib.shape_predictor(
"./pretrained_models/shape_predictor_68_face_landmarks.dat")
aligned_image = align_face(filepath=image_path, predictor=predictor)
print("Aligned image has shape: {}".format(aligned_image.size))
return aligned_image
def psp(image_path=None, net=None, opts=None):
from argparse import Namespace
import time
import os
import sys
import pprint
import numpy as np
from PIL import Image
import torch
import torchvision.transforms as transforms
sys.path.append(".")
sys.path.append("..")
from datasets import augmentations
from utils.common import tensor2im, log_input_image
from models.psp import pSp
import numpy as np
from sklearn.manifold import TSNE
import os
CODE_DIR = 'pixel2style2pixel'
experiment_type = 'ffhq_encode'
def get_download_model_command(file_id, file_name):
""" Get wget download command for downloading the desired model and save to directory ../pretrained_models. """
current_directory = os.getcwd()
save_path = os.path.join(os.path.dirname(current_directory), CODE_DIR,
"pretrained_models")
if not os.path.exists(save_path):
os.makedirs(save_path)
url = r"""wget --load-cookies /tmp/cookies.txt "https://docs.google.com/uc?export=download&confirm=$(wget --quiet --save-cookies /tmp/cookies.txt --keep-session-cookies --no-check-certificate 'https://docs.google.com/uc?export=download&id={FILE_ID}' -O- | sed -rn 's/.*confirm=([0-9A-Za-z_]+).*/\1\n/p')&id={FILE_ID}" -O {SAVE_PATH}/{FILE_NAME} && rm -rf /tmp/cookies.txt""".format(
FILE_ID=file_id, FILE_NAME=file_name, SAVE_PATH=save_path)
return url
MODEL_PATHS = {
"ffhq_encode": {
"id": "1bMTNWkh5LArlaWSc_wa8VKyq2V42T2z0",
"name": "psp_ffhq_encode.pt"
},
"ffhq_frontalize": {
"id": "1_S4THAzXb-97DbpXmanjHtXRyKxqjARv",
"name": "psp_ffhq_frontalization.pt"
},
"celebs_sketch_to_face": {
"id": "1lB7wk7MwtdxL-LL4Z_T76DuCfk00aSXA",
"name": "psp_celebs_sketch_to_face.pt"
},
"celebs_seg_to_face": {
"id": "1VpEKc6E6yG3xhYuZ0cq8D2_1CbT0Dstz",
"name": "psp_celebs_seg_to_face.pt"
},
"celebs_super_resolution": {
"id": "1ZpmSXBpJ9pFEov6-jjQstAlfYbkebECu",
"name": "psp_celebs_super_resolution.pt"
},
"toonify": {
"id": "1YKoiVuFaqdvzDP5CZaqa3k5phL-VDmyz",
"name": "psp_ffhq_toonify.pt"
}
}
path = MODEL_PATHS[experiment_type]
download_command = get_download_model_command(file_id=path["id"],
file_name=path["name"])
EXPERIMENT_DATA_ARGS = {
"ffhq_encode": {
"model_path":
"pretrained_models/psp_ffhq_encode.pt",
"image_path":
"/home/dibabdal/Desktop/MySpace/Devs/SfSNet-Pytorch-master/Images/REF_ID00353_Cam11_0063.png",
"transform":
transforms.Compose([
transforms.Resize((256, 256)),
transforms.ToTensor(),
transforms.Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5])
])
},
"ffhq_frontalize": {
"model_path":
"pretrained_models/psp_ffhq_frontalization.pt",
"image_path":
"notebooks/images/input_img.jpg",
"transform":
transforms.Compose([
transforms.Resize((256, 256)),
transforms.ToTensor(),
transforms.Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5])
])
},
"celebs_sketch_to_face": {
"model_path":
"pretrained_models/psp_celebs_sketch_to_face.pt",
"image_path":
"notebooks/images/input_sketch.jpg",
"transform":
transforms.Compose(
[transforms.Resize((256, 256)),
transforms.ToTensor()])
},
"celebs_seg_to_face": {
"model_path":
"pretrained_models/psp_celebs_seg_to_face.pt",
"image_path":
"notebooks/images/input_mask.png",
"transform":
transforms.Compose([
transforms.Resize((256, 256)),
augmentations.ToOneHot(n_classes=19),
transforms.ToTensor()
])
},
"celebs_super_resolution": {
"model_path":
"pretrained_models/psp_celebs_super_resolution.pt",
"image_path":
"notebooks/images/input_img.jpg",
"transform":
transforms.Compose([
transforms.Resize((256, 256)),
augmentations.BilinearResize(factors=[16]),
transforms.Resize((256, 256)),
transforms.ToTensor(),
transforms.Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5])
])
},
"toonify": {
"model_path":
"pretrained_models/psp_ffhq_toonify.pt",
"image_path":
"notebooks/images/input_img.jpg",
"transform":
transforms.Compose([
transforms.Resize((256, 256)),
transforms.ToTensor(),
transforms.Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5])
])
},
}
EXPERIMENT_ARGS = EXPERIMENT_DATA_ARGS[experiment_type]
if net is None:
model_path = EXPERIMENT_ARGS['model_path']
ckpt = torch.load(model_path, map_location='cpu')
opts = ckpt['opts']
# update the training options
opts['checkpoint_path'] = model_path
if 'learn_in_w' not in opts:
opts['learn_in_w'] = False
opts = Namespace(**opts)
net = pSp(opts)
net.eval()
net.cuda()
print('Model successfully loaded!')
if image_path is None:
image_path = EXPERIMENT_DATA_ARGS[experiment_type]["image_path"]
original_image = Image.open(image_path)
if opts.label_nc == 0:
original_image = original_image.convert("RGB")
else:
original_image = original_image.convert("L")
original_image.resize((256, 256))
def run_alignment(image_path):
import dlib
from scripts.align_all_parallel import align_face
predictor = dlib.shape_predictor(
"shape_predictor_68_face_landmarks.dat")
aligned_image = align_face(filepath=image_path, predictor=predictor)
print("Aligned image has shape: {}".format(aligned_image.size))
return aligned_image
if experiment_type not in ["celebs_sketch_to_face", "celebs_seg_to_face"]:
input_image = run_alignment(image_path)
else:
input_image = original_image
input_image.resize((256, 256))
img_transforms = EXPERIMENT_ARGS['transform']
transformed_image = img_transforms(input_image)
def run_on_batch(inputs, net, latent_mask=None):
latent = None
if latent_mask is None:
result_batch, latent = net(inputs.to("cuda").float(),
randomize_noise=False,
return_latents=True)
else:
result_batch = []
for image_idx, input_image in enumerate(inputs):
# get latent vector to inject into our input image
vec_to_inject = np.random.randn(1, 512).astype('float32')
_, latent_to_inject = net(
torch.from_numpy(vec_to_inject).to("cuda"),
input_code=True,
return_latents=True)
# get output image with injected style vector
res = net(input_image.unsqueeze(0).to("cuda").float(),
latent_mask=latent_mask,
inject_latent=latent_to_inject)
result_batch.append(res)
result_batch = torch.cat(result_batch, dim=0)
return result_batch, latent
if experiment_type in ["celebs_sketch_to_face", "celebs_seg_to_face"]:
latent_mask = [8, 9, 10, 11, 12, 13, 14, 15, 16, 17]
else:
latent_mask = None
with torch.no_grad():
tic = time.time()
result_image, latent = run_on_batch(transformed_image.unsqueeze(0),
net, latent_mask)
result_image = result_image[0]
toc = time.time()
print('Inference took {:.4f} seconds.'.format(toc - tic))
input_vis_image = log_input_image(transformed_image, opts)
output_image = tensor2im(result_image)
if experiment_type == "celebs_super_resolution":
res = np.concatenate([
np.array(input_image.resize((256, 256))),
np.array(input_vis_image.resize((256, 256))),
np.array(output_image.resize((256, 256)))
],
axis=1)
else:
res = np.concatenate([
np.array(input_vis_image.resize((256, 256))),
np.array(output_image.resize((256, 256)))
],
axis=1)
res_image = Image.fromarray(res)
import gc
gc.collect()
return res_image, latent, net, opts
def initialize(network_path, verbose):
def get_download_model_command(file_id, file_name):
""" Get wget download command for downloading the desired model and save to directory ../pretrained_models. """
current_directory = os.getcwd()
save_path = os.path.join(os.path.dirname(current_directory), CODE_DIR,
"pretrained_models")
if not os.path.exists(save_path):
os.makedirs(save_path)
url = r"""wget --load-cookies /tmp/cookies.txt "https://docs.google.com/uc?export=download&confirm=$(wget --quiet --save-cookies /tmp/cookies.txt --keep-session-cookies --no-check-certificate 'https://docs.google.com/uc?export=download&id={FILE_ID}' -O- | sed -rn 's/.*confirm=([0-9A-Za-z_]+).*/\1\n/p')&id={FILE_ID}" -O {SAVE_PATH}/{FILE_NAME} && rm -rf /tmp/cookies.txt""".format(
FILE_ID=file_id, FILE_NAME=file_name, SAVE_PATH=save_path)
return url
MODEL_PATHS = {
"ffhq_encode": {
"id": "1sw6I2lRIB0MpuJkpc8F5BJiSZrc0hjfE",
"name": "restyle_psp_ffhq_encode.pt"
},
# "toonify": {"id": "1GtudVDig59d4HJ_8bGEniz5huaTSGO_0", "name": "restyle_psp_toonify.pt"}
}
path = MODEL_PATHS[experiment_type]
download_command = get_download_model_command(file_id=path["id"],
file_name=path["name"])
EXPERIMENT_DATA_ARGS = {
"ffhq_encode": {
"model_path":
network_path, #"pretrained_models/restyle_psp_ffhq_encode.pt",
# "image_path": "notebooks/images/face_img.jpg",
"transform":
transforms.Compose([
transforms.Resize((256, 256)),
transforms.ToTensor(),
transforms.Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5])
])
},
# "toonify": {
# "model_path": "pretrained_models/restyle_psp_toonify.pt",
# "image_path": "notebooks/images/toonify_img.jpg",
# "transform": transforms.Compose([
# transforms.Resize((256, 256)),
# transforms.ToTensor(),
# transforms.Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5])])
# },
}
EXPERIMENT_ARGS = EXPERIMENT_DATA_ARGS[experiment_type]
if not os.path.exists(EXPERIMENT_ARGS['model_path']) or os.path.getsize(
EXPERIMENT_ARGS['model_path']) < 1000000:
print(f'Downloading ReStyle model for {experiment_type}...')
os.system(f"wget {download_command}")
# if google drive receives too many requests, we'll reach the quota limit and be unable to download the model
if os.path.getsize(EXPERIMENT_ARGS['model_path']) < 1000000:
raise ValueError(
"Pretrained model was unable to be downloaded correctly!")
else:
print('Done.')
else:
print(f'ReStyle model for {experiment_type} already exists!')
time_before = time.time()
# LOAD PRETRAINED MODEL
model_path = EXPERIMENT_ARGS['model_path']
ckpt = torch.load(model_path, map_location='cpu')
opts = ckpt['opts']
# resize_amount = (256, 256) if opts.resize_outputs else (opts.output_size, opts.output_size)
# number of output images
opts['n_iters_per_batch'] = NUM_STEPS
opts['resize_outputs'] = False # generate outputs at full resolution
# update the training options
opts['checkpoint_path'] = model_path
opts = Namespace(**opts)
net = pSp(opts)
net.eval()
net.cuda()
print('Model successfully loaded!')
time_after_loading = time.time()
print('Time to load model took {:.4f} seconds.'.format(time_after_loading -
time_before))
if verbose:
print(f'Loaded network from {network_path}.')
return net, opts
def run():
test_opts = TestOptions().parse()
out_path_results = os.path.join(test_opts.exp_dir, 'inference_results')
os.makedirs(out_path_results, exist_ok=True)
# load model used for initializing encoder bootstrapping
ckpt = torch.load(test_opts.model_1_checkpoint_path, map_location='cpu')
opts = ckpt['opts']
opts.update(vars(test_opts))
opts['checkpoint_path'] = test_opts.model_1_checkpoint_path
opts = Namespace(**opts)
if opts.encoder_type in ENCODER_TYPES['pSp']:
net1 = pSp(opts)
else:
net1 = e4e(opts)
net1.eval()
net1.cuda()
# load model used for translating input image after initialization
ckpt = torch.load(test_opts.model_2_checkpoint_path, map_location='cpu')
opts = ckpt['opts']
opts.update(vars(test_opts))
opts['checkpoint_path'] = test_opts.model_2_checkpoint_path
opts = Namespace(**opts)
if opts.encoder_type in ENCODER_TYPES['pSp']:
net2 = pSp(opts)
else:
net2 = e4e(opts)
net2.eval()
net2.cuda()
print('Loading dataset for {}'.format(opts.dataset_type))
dataset_args = data_configs.DATASETS[opts.dataset_type]
transforms_dict = dataset_args['transforms'](opts).get_transforms()
dataset = InferenceDataset(
root=opts.data_path,
transform=transforms_dict['transform_inference'],
opts=opts)
dataloader = DataLoader(dataset,
batch_size=opts.test_batch_size,
shuffle=False,
num_workers=int(opts.test_workers),
drop_last=False)
if opts.n_images is None:
opts.n_images = len(dataset)
# get the image corresponding to the latent average
avg_image = get_average_image(net1, opts)
resize_amount = (256, 256) if opts.resize_outputs else (opts.output_size,
opts.output_size)
global_i = 0
global_time = []
for input_batch in tqdm(dataloader):
if global_i >= opts.n_images:
break
with torch.no_grad():
input_cuda = input_batch.cuda().float()
tic = time.time()
result_batch = run_on_batch(input_cuda, net1, net2, opts,
avg_image)
toc = time.time()
global_time.append(toc - tic)
for i in range(input_batch.shape[0]):
results = [
tensor2im(result_batch[i][iter_idx])
for iter_idx in range(opts.n_iters_per_batch + 1)
]
im_path = dataset.paths[global_i]
input_im = tensor2im(input_batch[i])
# save step-by-step results side-by-side
res = np.array(results[0].resize(resize_amount))
for idx, result in enumerate(results[1:]):
res = np.concatenate(
[res, np.array(result.resize(resize_amount))], axis=1)
res = np.concatenate([res, input_im.resize(resize_amount)], axis=1)
Image.fromarray(res).save(
os.path.join(out_path_results, os.path.basename(im_path)))
global_i += 1
stats_path = os.path.join(opts.exp_dir, 'stats.txt')
result_str = 'Runtime {:.4f}+-{:.4f}'.format(np.mean(global_time),
np.std(global_time))
print(result_str)
with open(stats_path, 'w') as f:
f.write(result_str)
def project(image_path: str, output_path: str, network: str,
NUM_OUTPUT_IMAGES: int):
experiment_type = 'ffhq_encode' #['ffhq_encode', 'cars_encode', 'church_encode', 'horse_encode', 'afhq_wild_encode', 'toonify']
CODE_DIR = 'restyle-encoder'
def get_download_model_command(file_id, file_name):
""" Get wget download command for downloading the desired model and save to directory ../pretrained_models. """
current_directory = os.getcwd()
save_path = os.path.join(os.path.dirname(current_directory), CODE_DIR,
"pretrained_models")
if not os.path.exists(save_path):
os.makedirs(save_path)
url = r"""wget --load-cookies /tmp/cookies.txt "https://docs.google.com/uc?export=download&confirm=$(wget --quiet --save-cookies /tmp/cookies.txt --keep-session-cookies --no-check-certificate 'https://docs.google.com/uc?export=download&id={FILE_ID}' -O- | sed -rn 's/.*confirm=([0-9A-Za-z_]+).*/\1\n/p')&id={FILE_ID}" -O {SAVE_PATH}/{FILE_NAME} && rm -rf /tmp/cookies.txt""".format(
FILE_ID=file_id, FILE_NAME=file_name, SAVE_PATH=save_path)
return url
MODEL_PATHS = {
"ffhq_encode": {
"id": "1sw6I2lRIB0MpuJkpc8F5BJiSZrc0hjfE",
"name": "restyle_psp_ffhq_encode.pt"
},
# "church_encode": {"id": "1bcxx7mw-1z7dzbJI_z7oGpWG1oQAvMaD", "name": "restyle_psp_church_encode.pt"},
# "horse_encode": {"id": "19_sUpTYtJmhSAolKLm3VgI-ptYqd-hgY", "name": "restyle_e4e_horse_encode.pt"},
# "afhq_wild_encode": {"id": "1GyFXVTNDUw3IIGHmGS71ChhJ1Rmslhk7", "name": "restyle_psp_afhq_wild_encode.pt"},
# "toonify": {"id": "1GtudVDig59d4HJ_8bGEniz5huaTSGO_0", "name": "restyle_psp_toonify.pt"}
}
path = MODEL_PATHS[experiment_type]
download_command = get_download_model_command(file_id=path["id"],
file_name=path["name"])
EXPERIMENT_DATA_ARGS = {
"ffhq_encode": {
"model_path":
network, #"pretrained_models/restyle_psp_ffhq_encode.pt",
# "image_path": "notebooks/images/face_img.jpg",
"transform":
transforms.Compose([
transforms.Resize((256, 256)),
transforms.ToTensor(),
transforms.Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5])
])
},
# "church_encode": {
# "model_path": "pretrained_models/restyle_psp_church_encode.pt",
# "image_path": "notebooks/images/church_img.jpg",
# "transform": transforms.Compose([
# transforms.Resize((256, 256)),
# transforms.ToTensor(),
# transforms.Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5])])
# },
# "horse_encode": {
# "model_path": "pretrained_models/restyle_e4e_horse_encode.pt",
# "image_path": "notebooks/images/horse_img.jpg",
# "transform": transforms.Compose([
# transforms.Resize((256, 256)),
# transforms.ToTensor(),
# transforms.Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5])])
# },
# "afhq_wild_encode": {
# "model_path": "pretrained_models/restyle_psp_afhq_wild_encode.pt",
# "image_path": "notebooks/images/afhq_wild_img.jpg",
# "transform": transforms.Compose([
# transforms.Resize((256, 256)),
# transforms.ToTensor(),
# transforms.Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5])])
# },
# "toonify": {
# "model_path": "pretrained_models/restyle_psp_toonify.pt",
# "image_path": "notebooks/images/toonify_img.jpg",
# "transform": transforms.Compose([
# transforms.Resize((256, 256)),
# transforms.ToTensor(),
# transforms.Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5])])
# },
}
EXPERIMENT_ARGS = EXPERIMENT_DATA_ARGS[experiment_type]
if not os.path.exists(EXPERIMENT_ARGS['model_path']) or os.path.getsize(
EXPERIMENT_ARGS['model_path']) < 1000000:
print(f'Downloading ReStyle model for {experiment_type}...')
os.system(f"wget {download_command}")
# if google drive receives too many requests, we'll reach the quota limit and be unable to download the model
if os.path.getsize(EXPERIMENT_ARGS['model_path']) < 1000000:
raise ValueError(
"Pretrained model was unable to be downloaded correctly!")
else:
print('Done.')
else:
print(f'ReStyle model for {experiment_type} already exists!')
# MOVE INFERENCE FUNCTIONS AND PRE-DECLARATIONS OUTSIDE OF INFINITE LOOP
def get_avg_image(net):
avg_image = net(net.latent_avg.unsqueeze(0),
input_code=True,
randomize_noise=False,
return_latents=False,
average_code=True)[0]
avg_image = avg_image.to('cuda').float().detach()
if experiment_type == "cars_encode":
avg_image = avg_image[:, 32:224, :]
return avg_image
from utils.inference_utils import run_on_batch
# MOVE FUNCTION TO DISPLAY RESULTS OUTSIDE OF LOOP
def get_coupled_results(result_batch, transformed_image):
"""
Visualize output images from left to right (the input image is on the right)
"""
result_tensors = result_batch[0] # there's one image in our batch
result_images = [
tensor2im(result_tensors[iter_idx])
for iter_idx in range(opts.n_iters_per_batch)
]
input_im = tensor2im(transformed_image)
res = np.array(result_images[0].resize(resize_amount))
for idx, result in enumerate(result_images[1:]):
res = np.concatenate(
[res, np.array(result.resize(resize_amount))], axis=1)
res = np.concatenate([res, input_im.resize(resize_amount)], axis=1)
res = Image.fromarray(res)
return res, result_images
##### ADD INFINITE LOOP ################################################################
time_before = time.time()
# LOAD PRETRAINED MODEL
model_path = EXPERIMENT_ARGS['model_path']
ckpt = torch.load(model_path, map_location='cpu')
opts = ckpt['opts']
# resize_amount = (256, 256) if opts.resize_outputs else (opts.output_size, opts.output_size)
# number of output images
opts['n_iters_per_batch'] = NUM_OUTPUT_IMAGES
opts['resize_outputs'] = False # generate outputs at full resolution
# update the training options
opts['checkpoint_path'] = model_path
opts = Namespace(**opts)
if experiment_type == 'horse_encode':
net = e4e(opts)
else:
net = pSp(opts)
net.eval()
net.cuda()
print('Model successfully loaded!')
time_after_loading = time.time()
print('Time to load model took {:.4f} seconds.'.format(time_after_loading -
time_before))
# VISUALIZE INPUT
# image_path = EXPERIMENT_DATA_ARGS[experiment_type]["image_path"]
original_image = Image.open(image_path).convert("RGB")
# if experiment_type == 'cars_encode':
# original_image = original_image.resize((192, 256))
# else:
# original_image = original_image.resize((256, 256))
# ALIGN IMAGE
# def run_alignment(image_path):
# import dlib
# from scripts.align_faces_parallel import align_face
# if not os.path.exists("shape_predictor_68_face_landmarks.dat"):
# print('Downloading files for aligning face image...')
# os.system('wget http://dlib.net/files/shape_predictor_68_face_landmarks.dat.bz2')
# os.system('bzip2 -dk shape_predictor_68_face_landmarks.dat.bz2')
# print('Done.')
# predictor = dlib.shape_predictor("shape_predictor_68_face_landmarks.dat")
# aligned_image = align_face(filepath=image_path, predictor=predictor)
# print("Aligned image has shape: {}".format(aligned_image.size))
# return aligned_image
# if experiment_type in ['ffhq_encode', 'toonify']:
# input_image = run_alignment(image_path)
# else:
# input_image = original_image
input_image = original_image
# PERFORM INFERENCE
img_transforms = EXPERIMENT_ARGS['transform']
transformed_image = img_transforms(input_image)
with torch.no_grad():
avg_image = get_avg_image(net)
tic = time.time()
result_batch, result_latents = run_on_batch(
transformed_image.unsqueeze(0).cuda(), net, opts, avg_image)
toc = time.time()
print('Inference took {:.4f} seconds.'.format(toc - tic))
# VISUALIZE RESULT
# get results & save
res, result_images = get_coupled_results(result_batch, transformed_image)
time_after = time.time()
print(
'Time to load model, perform projection, and save out the results took {:.4f} seconds.'
.format(time_after - time_before))
return res, result_images
def run():
test_opts = TestOptions().parse()
out_path_results = os.path.join(test_opts.exp_dir, 'inference_results')
out_path_coupled = os.path.join(test_opts.exp_dir, 'inference_coupled')
os.makedirs(out_path_results, exist_ok=True)
os.makedirs(out_path_coupled, exist_ok=True)
# update test options with options used during training
ckpt = torch.load(test_opts.checkpoint_path, map_location='cpu')
opts = ckpt['opts']
opts.update(vars(test_opts))
opts = Namespace(**opts)
net = pSp(opts)
net.eval()
net.cuda()
age_transformers = [
AgeTransformer(target_age=age) for age in opts.target_age.split(',')
]
print(f'Loading dataset for {opts.dataset_type}')
dataset_args = data_configs.DATASETS[opts.dataset_type]
transforms_dict = dataset_args['transforms'](opts).get_transforms()
dataset = InferenceDataset(
root=opts.data_path,
transform=transforms_dict['transform_inference'],
opts=opts)
dataloader = DataLoader(dataset,
batch_size=opts.test_batch_size,
shuffle=False,
num_workers=int(opts.test_workers),
drop_last=False)
if opts.n_images is None:
opts.n_images = len(dataset)
global_time = []
for age_transformer in age_transformers:
print(f"Running on target age: {age_transformer.target_age}")
global_i = 0
for input_batch in tqdm(dataloader):
if global_i >= opts.n_images:
break
with torch.no_grad():
input_age_batch = [
age_transformer(img.cpu()).to('cuda')
for img in input_batch
]
input_age_batch = torch.stack(input_age_batch)
input_cuda = input_age_batch.cuda().float()
tic = time.time()
result_batch = run_on_batch(input_cuda, net, opts)
toc = time.time()
global_time.append(toc - tic)
for i in range(len(input_batch)):
result = tensor2im(result_batch[i])
im_path = dataset.paths[global_i]
if opts.couple_outputs or global_i % 100 == 0:
input_im = log_image(input_batch[i], opts)
resize_amount = (
256, 256) if opts.resize_outputs else (1024, 1024)
res = np.concatenate([
np.array(input_im.resize(resize_amount)),
np.array(result.resize(resize_amount))
],
axis=1)
age_out_path_coupled = os.path.join(
out_path_coupled, age_transformer.target_age)
os.makedirs(age_out_path_coupled, exist_ok=True)
Image.fromarray(res).save(
os.path.join(age_out_path_coupled,
os.path.basename(im_path)))
age_out_path_results = os.path.join(
out_path_results, age_transformer.target_age)
os.makedirs(age_out_path_results, exist_ok=True)
image_name = os.path.basename(im_path)
im_save_path = os.path.join(age_out_path_results,
image_name)
Image.fromarray(np.array(
result.resize(resize_amount))).save(im_save_path)
global_i += 1
stats_path = os.path.join(opts.exp_dir, 'stats.txt')
result_str = 'Runtime {:.4f}+-{:.4f}'.format(np.mean(global_time),
np.std(global_time))
print(result_str)
with open(stats_path, 'w') as f:
f.write(result_str)
