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)