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 parse_and_log_images(self, id_logs, x, y, y_hat, title, subscript=None, display_count=2):
im_data = []
for i in range(display_count):
cur_im_data = {
'input_face': common.log_input_image(x[i], self.opts),
'target_face': common.tensor2im(y[i]),
'output_face': common.tensor2im(y_hat[i]),
}
if id_logs is not None:
for key in id_logs[i]:
cur_im_data[key] = id_logs[i][key]
im_data.append(cur_im_data)
self.log_images(title, im_data=im_data, subscript=subscript)
def log_images_to_wandb(x, y, y_hat, id_logs, prefix, step, opts):
im_data = []
column_names = ["Source", "Target", "Output"]
if id_logs is not None:
column_names.append("ID Diff Output to Target")
for i in range(len(x)):
cur_im_data = [
wandb.Image(common.log_input_image(x[i], opts)),
wandb.Image(common.tensor2im(y[i])),
wandb.Image(common.tensor2im(y_hat[i])),
]
if id_logs is not None:
cur_im_data.append(id_logs[i]["diff_target"])
im_data.append(cur_im_data)
outputs_table = wandb.Table(data=im_data, columns=column_names)
wandb.log(
{f"{prefix.title()} Step {step} Output Samples": outputs_table})
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 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 online_inference_fn(image_path, output_path):
# Step1: 步骤1，加载图片，图片大小调整。load image and resize
# resize image first
# Visiual input
# image_path = EXPERIMENT_DATA_ARGS[experiment_type]["image_path"]
file_name = image_path.split("/")[-1]
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)) # resize
original_image = original_image.resize((256, 256), resample=Image.BILINEAR)
# Step2: 步骤2： 裁剪出人脸。run_alignment
if experiment_type not in ["celebs_sketch_to_face", "celebs_seg_to_face"]:
try:
input_image = run_alignment(image_path)
print("run alignment error!...")
except Exception as e:
input_image = original_image
else:
input_image = original_image
# Step 3: 运行生成程序 Perform Inference
img_transforms = EXPERIMENT_ARGS['transform']
transformed_image = img_transforms(input_image)
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 = run_on_batch(transformed_image.unsqueeze(0), net,
latent_mask)[0]
toc = time.time()
print('Inference took {:.4f} seconds.'.format(toc - tic))
# Step 6.2: Visualize Result
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) # 实现array到image的转换, return an image object
# res_image
# save image to output path
res_image.save(output_path + "_" + file_name + "_toon.jpg")
