def get_inversions_on_batch(inputs, net, avg_image, opts):
result_batch, result_latents = run_on_batch(inputs, net, opts, avg_image)
# we'll take the final inversion as the inversion to edit
y_hat = [result_batch[idx][-1] for idx in range(len(result_batch))]
latents = [
torch.from_numpy(result_latents[idx][-1]).cuda()
for idx in range(len(result_batch))
]
return y_hat, torch.stack(latents)
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 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 process(input_path, output_path, verbose, preloaded_params):
time_before = time.time()
image_path = input_path
net = preloaded_params[0]
opts = preloaded_params[1]
# VISUALIZE INPUT
# image_path = EXPERIMENT_DATA_ARGS[experiment_type]["image_path"]
original_image = Image.open(image_path).convert("RGB")
# 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 = transforms.Compose([
transforms.Resize((256, 256)),
transforms.ToTensor(),
transforms.Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5])
])
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()
if verbose:
print('Inference took {:.4f} seconds.'.format(toc - tic))
# KEEP LAST W-VECTOR
w_last = result_latents[0][NUM_STEPS - 1]
# make sure output folder exists, otherwise saving won’t work
if not os.path.exists(output_path):
os.makedirs(output_path)
np.save(f'{output_path}/w_vector.npy', w_last, allow_pickle=True)
time_after = time.time()
if verbose:
print('Time to load img and get its w-vector took {:.4f} seconds.'.
format(time_after - time_before))
# Emit DONE signal
Path(f'{output_path}/w_vector.done').touch()
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