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")