def main(args): os.makedirs(args.outputs, exist_ok=True) generator = get_derivable_generator(args.gan_model, args.inversion_type, args) loss = get_loss(args.loss_type, args) mask = parsing_mask(args.mask).cuda() mask_cpu = parsing_mask(args.mask) crop_loss = masked_loss(loss, mask) # to cuda generator.cuda() loss.cuda() inversion = get_inversion(args.optimization, args) image_list = image_files(args.target_images) frameSize = MODEL_POOL[args.gan_model]['resolution'] for i, images in enumerate(split_to_batches(image_list, 1)): print('%d: Processing %d images :' % (i + 1, 1), end='') pt_image_str = '%s\n' print(pt_image_str % tuple(images)) image_name_list = [] image_tensor_list = [] for image in images: image_name_list.append(os.path.split(image)[1]) image_tensor_list.append(_add_batch_one(load_as_tensor(image))) y_gt = _sigmoid_to_tanh(torch.cat(image_tensor_list, dim=0)).cuda() # Invert latent_estimates, history = inversion.invert(generator, y_gt, crop_loss, batch_size=1, video=args.video) # Get Images y_estimate_list = torch.split(torch.clamp(_tanh_to_sigmoid(generator(latent_estimates)), min=0., max=1.).cpu(), 1, dim=0) # Save for img_id, image in enumerate(images): y_gt_pil = Tensor2PIL(mask_images(image_tensor_list[img_id], mask_cpu)) y_estimate_pil = Tensor2PIL(y_estimate_list[img_id]) y_estimate_pil.save(os.path.join(args.outputs, '%s-%s-%d.png' % (image_name_list[img_id], os.path.split(args.mask[:-4])[1], math.floor(time.time())))) y_gt_pil.save(os.path.join(args.outputs, '%s-%s.png' % (image_name_list[img_id], os.path.split(args.mask[:-4])[1]))) # Create video if args.video: print('Create GAN-Inversion video.') video = cv2.VideoWriter( filename=os.path.join(args.outputs, '%s_inpainting_%s.avi' % (image_name_list[img_id][:-4], os.path.split(args.mask[:-4])[1])), fourcc=cv2.VideoWriter_fourcc(*'MJPG'), fps=args.fps, frameSize=(frameSize, frameSize)) print('Save frames.') for i, sample in enumerate(history): image = generator(sample) image_cv2 = convert_array_to_images(image.detach().cpu().numpy())[0][:, :, ::-1] video.write(image_cv2) video.release()
def saveIMG(): image_list = image_files(default_load_img_pth) for i, images in enumerate(split_to_batches(image_list, 1)): image_name_list = [] image_tensor_list = [] for image in images: image_name_list.append(os.path.split(image)[1]) image_tensor_list.append(_add_batch_one(load_as_tensor(image))) # 保存结果 for img_id, image in enumerate(images): down = convert_test(image_tensor_list[img_id], factor=16, mode='nearest') y_nn_pil = Tensor2PIL(down) y_nn_pil.save( os.path.join( os.path.join('img_test_output', '%s-nn.png' % image_name_list[img_id][:-4])))
def main(args): os.makedirs(args.outputs, exist_ok=True) generator = get_derivable_generator(args.gan_model, args.inversion_type, args) # 生成器 loss = get_loss(args.loss_type, args) sr_loss = SR_loss(loss, args.down, args.factor) # SR计算loss的方式 # to cuda generator.cuda() loss.cuda() inversion = get_inversion(args.optimization, args) image_list = image_files(args.target_images) frameSize = MODEL_POOL[args.gan_model]['resolution'] for i, images in enumerate(split_to_batches(image_list, 1)): print('%d: Super-resolving %d images ' % (i + 1, 1), end='') pt_image_str = '%s\n' print(pt_image_str % tuple(images)) image_name_list = [] image_tensor_list = [] for image in images: image_name_list.append(os.path.split(image)[1]) # image = _add_batch_one(load_as_tensor(image)) image = convert2target(_add_batch_one(load_as_tensor(image)), 'nearest') # 更改size使得适用于更多类型的分辨率图像 image_tensor_list.append(image) # print("add..: ", _add_batch_one(load_as_tensor(image)).size()) # torch.Size([1, 3, 64, 64]) y_gt = _sigmoid_to_tanh(torch.cat(image_tensor_list, dim=0)).cuda() # Invert latent_estimates, history = inversion.invert(generator, y_gt, sr_loss, batch_size=BATCH_SIZE, video=args.video) # Get Images # 将optimizer优化好的latent_estimates再放入generator中生成图像 # 并且将batch_size那一列的数据去除 y_estimate_list = torch.split(torch.clamp(_tanh_to_sigmoid( generator(latent_estimates)), min=0., max=1.).cpu(), 1, dim=0) # 保存结果 for img_id, image in enumerate(images): # up_nn, up_bic, down = downsample_images(image_tensor_list[img_id], factor=args.factor, mode=args.down) # y_nn_pil = Tensor2PIL(up_nn) # 低分辨率化后的图像 y_estimate_pil = Tensor2PIL(y_estimate_list[img_id]) y_estimate_pil.save( os.path.join( os.path.join(args.outputs, '%s.png' % image_name_list[img_id][:-4]))) #y_nn_pil.save(os.path.join(os.path.join(args.outputs, '%s-nn.png' % image_name_list[img_id][:-4]))) # Create video if args.video: print('Create GAN-Inversion video.') video = cv2.VideoWriter(filename=os.path.join( args.outputs, '%s_sr.avi' % image_name_list[img_id][:-4]), fourcc=cv2.VideoWriter_fourcc(*'MJPG'), fps=args.fps, frameSize=(frameSize, frameSize)) print('Save frames.') for i, sample in enumerate(history): image = generator( sample ) # 用generator从history(保存的训练中的estimate_latent的值)中生成图像 image_cv2 = convert_array_to_images( image.detach().cpu().numpy())[0][:, :, ::-1] video.write(image_cv2) video.release()
def main(args): os.makedirs(args.outputs, exist_ok=True) out_dir, exp_name = create_experiments_directory(args, args.exp_id) print(out_dir) print(exp_name) generator = get_derivable_generator(args.gan_model, args.inversion_type, args) generator.cuda() if args.target_images.endswith('.png') or args.target_images.endswith( '.jpg'): image_list = os.path.abspath(args.target_images) image_list = [image_list] else: image_list = image_files(args.target_images) frameSize = MODEL_POOL[args.gan_model]['resolution'] n_blocks = generator.n_blocks print('There are %d blocks in this generator.' % n_blocks) # 19 for pggan latent_space = generator.PGGAN_LATENT[args.layer] print('The latent space is ', latent_space) with open(args.matrix_dir, 'rb') as file_in: matrix = pkl.load(file_in) print('Load matrix successfully.') print('Matrix shape ', matrix.shape) matrix2 = thrC(matrix, args.alpha) predict, _ = post_proC(matrix2, args.n_subs, args.d_subs, args.power) print(predict) p_sum = [sum(predict == k) for k in range(1, args.cluster_numbers, 1)] p_sum = np.array(p_sum) p_sort = np.argsort(p_sum)[::-1] print(p_sum) predict_new = predict.copy() for i in range(1, args.cluster_numbers, 1): predict_new[predict == (p_sort[i - 1] + 1)] = i predict = predict_new.copy() p_sum = [sum(predict == k) for k in range(1, args.cluster_numbers, 1)] print(predict) print(p_sum) # pre_see the images gan_type, image_type = args.gan_model.split("_") print('The gan type is %s, and the image type is %s' % (gan_type, image_type)) test_image_dir = os.path.join('./bin/', gan_type, image_type) print(test_image_dir) files = os.listdir(test_image_dir) test_zs = [] for i in range(len(files)): if files[i].endswith('.pkl'): with open(os.path.join(test_image_dir, files[i]), 'rb') as file_in: test_zs.append(pkl.load(file_in)) test_zs = torch.from_numpy( np.concatenate(test_zs, axis=0).astype(np.float32)).cuda() print('Load all testing zs, shape is ', test_zs.size()) image_number = 3 sel_idx = np.random.choice(test_zs.shape[0], size=[image_number], replace=False) F = generator([test_zs[sel_idx]], which_block=args.layer, pre_model=True) features = F.detach().cpu().numpy() predict_masks = [] for i in range(1, args.cluster_numbers + 1, 1): mask = torch.from_numpy((predict == i).astype(np.float32)).cuda() predict_masks += [mask.reshape((1, -1, 1, 1))] for i, images in enumerate(split_to_batches(image_list, 1)): # input_size = generator.PGGAN_LATENT[args.layer + 1] # print("We are making ", input_size, ". ") print('%d: Inverting %d images :' % (i + 1, 1), end='') pt_image_str = '%s\n' print(pt_image_str % tuple(images)) image_name_only = images[0].split(".")[0] image_name_only = image_name_only.split("/")[-1] print(image_name_only) image_name_list = [] image_tensor_list = [] for image in images: image_name_list.append(os.path.split(image)[1]) image_tensor_list.append(_add_batch_one(load_as_tensor(image))) print("image_name_list", image_name_list) print("image_tensor_list, [", image_tensor_list[0].size(), "]") y_image = _sigmoid_to_tanh(torch.cat(image_tensor_list, dim=0)).cuda() print('image size is ', y_image.size()) z_estimate = generator.init_value(batch_size=1, which_layer=0, init=args.init_type, z_numbers=args.cluster_numbers * args.code_per_cluster) base_estimate = generator.init_value(batch_size=1, which_layer=0, init=args.init_type, z_numbers=1) if args.optimization == 'GD': z_optimizer = torch.optim.SGD(z_estimate + base_estimate, lr=args.lr) elif args.optimization == 'Adam': z_optimizer = torch.optim.Adam(z_estimate + base_estimate, lr=args.lr) else: raise NotImplemented( 'We don\'t support this type of optimization.') for iter in range(args.iterations): for estimate in z_estimate: estimate.requires_grad = True for estimate in base_estimate: estimate.requires_grad = True features = generator([ z_estimate[0].reshape( [args.cluster_numbers * args.code_per_cluster, 512, 1, 1]) ], which_block=args.layer, pre_model=True) base_feature = generator( [base_estimate[0].reshape([1, 512, 1, 1])], which_block=args.layer, pre_model=True) for t in range(args.cluster_numbers * args.code_per_cluster): if t == 0: f_mix = features[t].view(*(1, ) + latent_space) * predict_masks[int( t / args.code_per_cluster)] else: f_mix = f_mix + features[t].view( *(1, ) + latent_space) * predict_masks[int( t / args.code_per_cluster)] f_mix = f_mix + base_feature y_estimate = generator([f_mix], which_block=args.layer, post_model=True) y_raw_estimate = generator([base_feature], which_block=args.layer, post_model=True) z_optimizer.zero_grad() loss = 0.01 * torch.mean(torch.pow( y_estimate - y_image, 2.0)) + torch.mean( torch.pow(y_raw_estimate - y_image, 2.0)) loss.backward() z_optimizer.step() if iter % args.report_value == 0: print('Iter %d, layer %d, loss = %.4f.' % (iter, args.layer, float(loss.item()))) if iter % args.report_image == 0: print('Saving the images.') y_estimate_pil = Tensor2PIL( torch.clamp(_tanh_to_sigmoid(y_estimate.detach().cpu()), min=0.0, max=1.0)) y_estimate_pil.save( os.path.join( out_dir, image_name_only + "_estimate_iter%d.png" % iter)) y_estimate_pil = Tensor2PIL( torch.clamp(_tanh_to_sigmoid( y_raw_estimate.detach().cpu()), min=0.0, max=1.0)) y_estimate_pil.save( os.path.join( out_dir, image_name_only + "_raw_estimate_iter%d.png" % iter)) # add bias added output picture. # save all the codes codes = [] for code_idx in range(args.cluster_numbers * args.code_per_cluster): code_f = generator([z_estimate[0][code_idx]], which_block=args.layer + 1, pre_model=True) code_y = generator([code_f], which_block=args.layer + 1, post_model=True).detach().cpu() codes.append( torch.clamp(_tanh_to_sigmoid(code_y), min=0, max=1)) codes = torch.cat(codes, dim=0).detach().cpu() torchvision.utils.save_image( codes, os.path.join( out_dir, image_name_only + '_codes_iter%d.png' % (iter)), nrow=(args.cluster_numbers * args.code_per_cluster) // 2) if iter % args.report_model == 0: print('Save the models') save_dict = { "z": z_estimate[0].detach().cpu().numpy(), 'matrix': matrix, 'layer': args.layer, 'predict': predict } with open( os.path.join( out_dir, 'save_dict_iter_%d_layer_%d.pkl' % (iter, args.layer)), 'wb') as file_out: pkl.dump(save_dict, file_out) print('Save the models OK!')
def main(args): os.makedirs(args.outputs, exist_ok=True) generator = get_derivable_generator( args.gan_model, args.inversion_type, args) # pggan_celebahq/stylegan2_ffhq256 PGGAN-Multi-Z/StyleGAN-w+ loss = get_loss(args.loss_type, args) sr_loss = SR_loss(loss, args.down, args.factor) # to cuda generator.cuda() loss.cuda() inversion = get_inversion(args.optimization, args) # SGD image_list = image_files(args.target_images) frameSize = MODEL_POOL[args.gan_model][ 'resolution'] # gan_model=pggan_celebahq for i, images in enumerate(split_to_batches(image_list, 1)): print('%d: Super-resolving %d images ' % (i + 1, 1), end='') pt_image_str = '%s\n' print(pt_image_str % tuple(images)) image_name_list = [] image_tensor_list = [] for image in images: image_name_list.append(os.path.split(image)[1]) image_tensor_list.append(_add_batch_one(load_as_tensor(image))) y_gt = _sigmoid_to_tanh(torch.cat(image_tensor_list, dim=0)).cuda() # Invert latent_estimates, history = inversion.invert(generator, y_gt, sr_loss, batch_size=1, video=args.video) # Get Images y_estimate_list = torch.split(torch.clamp(_tanh_to_sigmoid( generator(latent_estimates)), min=0., max=1.).cpu(), 1, dim=0) # Save for img_id, image in enumerate(images): up_nn, up_bic = downsample_images(image_tensor_list[img_id], factor=args.factor, mode=args.down) y_nn_pil = Tensor2PIL(up_nn) y_estimate_pil = Tensor2PIL(y_estimate_list[img_id]) y_estimate_pil.save( os.path.join( os.path.join(args.outputs, '%s.png' % image_name_list[img_id][:-4]))) y_nn_pil.save( os.path.join( os.path.join(args.outputs, '%s-nn.png' % image_name_list[img_id][:-4]))) # Create video if args.video: print('Create GAN-Inversion video.') video = cv2.VideoWriter(filename=os.path.join( args.outputs, '%s_sr.avi' % image_name_list[img_id][:-4]), fourcc=cv2.VideoWriter_fourcc(*'MJPG'), fps=args.fps, frameSize=(frameSize, frameSize)) print('Save frames.') for i, sample in enumerate(history): image = generator(sample) image_cv2 = convert_array_to_images( image.detach().cpu().numpy())[0][:, :, ::-1] video.write(image_cv2) video.release()
def invert(self, generator, gt_image, loss_function, batch_size=1, video=True, init=(), z_iterations=3000, args=None, out_dir=None, image_name=None): input_size_list = generator.input_size() if len(init) == 0: if generator.init is False: latent_estimate = [] for input_size in input_size_list: if self.init_type == 'Zero': print("Zero") latent_estimate.append( torch.zeros((batch_size, ) + input_size).cuda()) elif self.init_type == 'Normal': latent_estimate.append( torch.randn((batch_size, ) + input_size).cuda()) else: latent_estimate = list(generator.init_value(batch_size)) else: assert len(init) == len( input_size_list), 'Please check the number of init value' latent_estimate = init for latent in latent_estimate: latent.requires_grad = True # 19 PGGAN - 0: 1 4 - 18: 1e-8 x 10 = 1e-6 print(latent_estimate[0].requires_grad) optimizer = self.optimizer(latent_estimate, lr=self.lr) history = [] # Opt for i in tqdm(range(z_iterations)): y_estimate = generator(latent_estimate) if i % args.report_image == 0: if out_dir is not None: y_estimate_pil = Tensor2PIL( torch.clamp(_tanh_to_sigmoid( y_estimate.detach().cpu()), min=0.0, max=1.0)) y_estimate_pil.save( os.path.join( out_dir, image_name.split('.')[0] + "_" + str(i) + "_layer_" + str(args.composing_layer) + ".png")) zs = latent_estimate[0].detach().cpu().numpy() ps = latent_estimate[1].detach().cpu().numpy() save_dict = {'zs': zs, 'ps': ps} with open( os.path.join( out_dir, image_name.split('.')[0] + "_" + str(i) + "_layer_" + str(args.composing_layer) + ".pkl"), 'wb') as file_out: pkl.dump(save_dict, file_out) optimizer.zero_grad() loss = loss_function(y_estimate, gt_image) loss.backward() optimizer.step() if video: history.append(copy.deepcopy(latent_estimate)) return latent_estimate, history
def main(args): os.makedirs(args.outputs + '/input', exist_ok=True) os.makedirs(args.outputs + '/GT', exist_ok=True) os.makedirs(args.outputs + '/mGANoutput', exist_ok=True) with open(args.outputs + '/mGANargs.txt', 'w') as f: json.dump(args.__dict__, f, indent=2) generator = get_derivable_generator(args.gan_model, args.inversion_type, args) loss = get_loss(args.loss_type, args) mask = parsing_mask('code/mganprior/masks/' + args.mask).cuda() mask_cpu = parsing_mask('code/mganprior/masks/' + args.mask) crop_loss = masked_loss(loss, mask) generator.cuda() loss.cuda() inversion = get_inversion(args.optimization, args) image_list = image_files(args.target_images) if len(image_list) > 300: print('Limiting the image set to 300.') image_list = image_list[:300] frameSize = MODEL_POOL[args.gan_model]['resolution'] start_time = time.time() for i, images in enumerate(split_to_batches(image_list, 1)): print('%d: Processing %d images :' % (i, 1), end='') pt_image_str = '%s\n' print(pt_image_str % tuple(images)) image_name_list = [] image_tensor_list = [] for image in images: image_name_list.append(os.path.split(image)[1]) image_tensor_list.append(_add_batch_one(load_as_tensor(image))) y_gt = _sigmoid_to_tanh(torch.cat(image_tensor_list, dim=0)).cuda() # Invert if args.varmask: os.makedirs(args.outputs + '/mask', exist_ok=True) mask_cpu = get_var_mask(y_gt.shape[-2:], args.min_p, args.max_p, args.width_mean, args.width_var) mask = mask_cpu.cuda() save_image(mask, os.path.join(args.outputs + '/mask/%d%s' % (i, '.png'))) crop_loss = masked_loss(loss, mask) latent_estimates, history = inversion.invert(generator, y_gt, crop_loss, batch_size=1, video=args.video) # Get Images y_estimate_list = torch.split(torch.clamp(_tanh_to_sigmoid( generator(latent_estimates)), min=0., max=1.).cpu(), 1, dim=0) # Save for img_id, image in enumerate(images): y_RGB = Tensor2PIL(image_tensor_list[img_id]) y_RGB.save(args.outputs + '/GT/%d%s' % (i, image_name_list[img_id][-4:])) y_gt_pil = Tensor2PIL( mask_images(image_tensor_list[img_id], mask_cpu)) y_estimate_pil = Tensor2PIL(y_estimate_list[img_id]) y_estimate_pil.save( os.path.join(args.outputs + '/mGANoutput/%d%s' % (i, image_name_list[img_id][-4:]))) y_gt_pil.save( os.path.join(args.outputs + '/input/%d%s' % (i, image_name_list[img_id][-4:]))) # Create video if args.video: print('Create GAN-Inversion video.') video = cv2.VideoWriter(filename=os.path.join( args.outputs, '%s_inpainting_%s.avi' % (image_name_list[img_id][:-4], os.path.split( args.mask[:-4])[1])), fourcc=cv2.VideoWriter_fourcc(*'MJPG'), fps=args.fps, frameSize=(frameSize, frameSize)) print('Save frames.') for i, sample in enumerate(history): image = generator(sample) image_cv2 = convert_array_to_images( image.detach().cpu().numpy())[0][:, :, ::-1] video.write(image_cv2) video.release() print(f'{(time.time()-start_time)/60:.2f}', 'minutes taken in total;', f'{(time.time()-start_time)/60/len(image_list):.2f}', 'per image.')
def run(args): os.makedirs(args.outputs, exist_ok=True) # 生成输出路径文件夹,存在则跳过 # 生成器 generator = get_derivable_generator(args.gan_model, args.inversion_type, args) loss = get_loss(args.loss_type, args) # 损失函数 generator.cuda() # pytorch需要手动放入GPU进行运算 loss.cuda() inversion = get_inversion(args.optimization, args) image_list = image_files(args.target_images) # 获取输入图片路径 frameSize = MODEL_POOL[args.gan_model]['resolution'] # 获取图像分辨率 # 按照batch大小分批处理图像 for i, images in enumerate(split_to_batches(image_list, 1)): print('%d: Inverting %d images :' % (i + 1, 1), end='') # pt_image_str = '%s\n' print('%s\n' % tuple(images)) image_name_list = [] image_tensor_list = [] for image in images: image_name_list.append(os.path.split(image)[1]) image_tensor_list.append(_add_batch_one(load_as_tensor(image))) # torch.cat(tensors, dim=0, out=None) → Tensor # tensors (sequence of Tensors) – any python sequence of tensors of the same type. Non-empty tensors provided must have the same shape, except in the cat dimension. # dim (int, optional) – the dimension over which the tensors are concatenated # out (Tensor, optional) – the output tensor. y_gt = _sigmoid_to_tanh( torch.cat(image_tensor_list, dim=0)).cuda() # 在维度0上连接所有的tensor并且将值域映射到[-1, 1] # 逆映射, 生成图像tensor latent_estimates, history = inversion.invert(generator, y_gt, loss, batch_size=BATCH_SIZE, video=args.video) # 将值域从[-1,1]映射到[0,1], 使用torch.clamp()进一步保证值域在[0,1] y_estimate_list = torch.split(torch.clamp(_tanh_to_sigmoid( generator(latent_estimates)), min=0., max=1.).cpu(), 1, dim=0) # Save for img_id, image in enumerate(images): y_estimate_pil = Tensor2PIL( y_estimate_list[img_id]) # 从tensor转化为PIL image并保存 y_estimate_pil.save( os.path.join(args.outputs, image_name_list[img_id])) # Create video if args.video: print('Create GAN-Inversion video.') video = cv2.VideoWriter(filename=os.path.join( args.outputs, '%s_inversion.avi' % image_name_list[img_id]), fourcc=cv2.VideoWriter_fourcc(*'MJPG'), fps=args.fps, frameSize=(frameSize, frameSize)) print('Save frames.') for i, sample in enumerate(history): image = generator(sample) image_cv2 = convert_array_to_images( image.detach().cpu().numpy())[0][:, :, ::-1] video.write(image_cv2) video.release()
def main(args): os.makedirs(args.outputs, exist_ok=True) generator = get_derivable_generator(args.gan_model, args.inversion_type, args) loss = get_loss(args.loss_type, args) cor_loss = Color_loss(loss) generator.cuda() loss.cuda() inversion = get_inversion(args.optimization, args) image_list = image_files(args.target_images) frameSize = MODEL_POOL[args.gan_model]['resolution'] for i, images in enumerate(split_to_batches(image_list, 1)): print('%d: Processing %d images :' % (i + 1, 1), end='') pt_image_str = '%s\n' print(pt_image_str % tuple(images)) image_name_list = [] image_tensor_list = [] for image in images: image_name_list.append(os.path.split(image)[1]) image_tensor_list.append(_add_batch_one(load_as_tensor(image))) y_gt = _sigmoid_to_tanh(torch.cat(image_tensor_list, dim=0)).cuda() # Invert latent_estimates, history = inversion.invert(generator, y_gt, cor_loss, batch_size=1, video=args.video) # Get Images y_estimate_list = torch.split(torch.clamp(_tanh_to_sigmoid( generator(latent_estimates)), min=0., max=1.).cpu(), 1, dim=0) # Save for img_id, image in enumerate(images): up_gray = colorization_images(image_tensor_list[img_id]) y_gray_pil = Tensor2PIL(up_gray, mode='L') y_gray_pil.save( os.path.join(args.outputs, '%s-%s.png' % (image_name_list[img_id], 'gray'))) Y_gt = Tensor2PIL(image_tensor_list[img_id], mode='RGB').convert('YCbCr') y_estimate_pil = Tensor2PIL(y_estimate_list[img_id], mode='RGB').convert('YCbCr') _, Cb, Cr = y_estimate_pil.split() Y, _, _ = Y_gt.split() y_colorization = Image.merge('YCbCr', (Y, Cb, Cr)) y_colorization.convert('RGB').save( os.path.join( args.outputs, '%s-%d.png' % (image_name_list[img_id], math.floor(time.time())))) # Create video if args.video: print('Create GAN-Inversion video.') video = cv2.VideoWriter(filename=os.path.join( args.outputs, '%s_inversion.avi' % image_name_list[img_id]), fourcc=cv2.VideoWriter_fourcc(*'MJPG'), fps=args.fps, frameSize=(frameSize, frameSize)) print('Save frames.') for i, sample in enumerate(history): image = torch.clamp(_tanh_to_sigmoid(generator(sample)), min=0., max=1.).cpu() image_pil = Tensor2PIL(image, mode='RGB').convert('YCbCr') _, Cb, Cr = image_pil.split() y_colorization = Image.merge('YCbCr', (Y, Cb, Cr)).convert('RGB') image_cv2 = cv2.cvtColor(np.asarray(y_colorization), cv2.COLOR_RGB2BGR) # image_cv2 = cv2.cvtColor(np.asarray(image_pil.convert('RGB')), cv2.COLOR_RGB2BGR) video.write(image_cv2) video.release()
def main(args): os.makedirs(args.outputs+'/input', exist_ok=True) os.makedirs(args.outputs+'/GT', exist_ok=True) os.makedirs(args.outputs+'/mGANoutput', exist_ok=True) with open(args.outputs+'/mGANargs.txt', 'w') as f: json.dump(args.__dict__, f, indent=2) generator = get_derivable_generator(args.gan_model, args.inversion_type, args) loss = get_loss(args.loss_type, args) cor_loss = Color_loss(loss) generator.cuda() loss.cuda() inversion = get_inversion(args.optimization, args) image_list = image_files(args.target_images) if len(image_list)>300: print('Limiting the image set to 300.') image_list = image_list[:300] frameSize = MODEL_POOL[args.gan_model]['resolution'] start_time = time.time() for i, images in enumerate(split_to_batches(image_list, 1)): print('%d: Processing %d images :' % (i, 1), end='') pt_image_str = '%s\n' print(pt_image_str % tuple(images)) image_name_list = [] image_tensor_list = [] for image in images: image_name_list.append(os.path.split(image)[1]) image_tensor_list.append(_add_batch_one(load_as_tensor(image))) y_gt = _sigmoid_to_tanh(torch.cat(image_tensor_list, dim=0)).cuda() # Invert latent_estimates, history = inversion.invert(generator, y_gt, cor_loss, batch_size=1, video=args.video) # Get Images y_estimate_list = torch.split(torch.clamp(_tanh_to_sigmoid(generator(latent_estimates)), min=0., max=1.).cpu(), 1, dim=0) # Save for img_id, image in enumerate(images): up_gray = colorization_images(image_tensor_list[img_id]) y_gray_pil = Tensor2PIL(up_gray, mode='L') y_gray_pil.save(args.outputs + '/input/%d%s' % (i, image_name_list[img_id][-4:])) y_RGB = Tensor2PIL(image_tensor_list[img_id]) y_RGB.save(args.outputs + '/GT/%d%s' % (i, image_name_list[img_id][-4:])) Y_gt = Tensor2PIL(image_tensor_list[img_id], mode='RGB').convert('YCbCr') y_estimate_pil = Tensor2PIL(y_estimate_list[img_id], mode='RGB').convert('YCbCr') _, Cb, Cr = y_estimate_pil.split() Y, _, _ = Y_gt.split() y_colorization = Image.merge('YCbCr', (Y, Cb, Cr)) y_colorization.convert('RGB').save(args.outputs + '/mGANoutput/%d%s' % (i, image_name_list[img_id][-4:])) # Create video if args.video: print('Create GAN-Inversion video.') video = cv2.VideoWriter( filename=os.path.join(args.outputs, '%s_inversion.avi' % image_name_list[img_id]), fourcc=cv2.VideoWriter_fourcc(*'MJPG'), fps=args.fps, frameSize=(frameSize, frameSize)) print('Save frames.') for i, sample in enumerate(history): image = torch.clamp(_tanh_to_sigmoid(generator(sample)), min=0., max=1.).cpu() image_pil = Tensor2PIL(image, mode='RGB').convert('YCbCr') _, Cb, Cr = image_pil.split() y_colorization = Image.merge('YCbCr', (Y, Cb, Cr)).convert('RGB') image_cv2 = cv2.cvtColor(np.asarray(y_colorization), cv2.COLOR_RGB2BGR) # image_cv2 = cv2.cvtColor(np.asarray(image_pil.convert('RGB')), cv2.COLOR_RGB2BGR) video.write(image_cv2) video.release() print(f'{(time.time()-start_time)/60:.2f}','minutes taken in total;', f'{(time.time()-start_time)/60/len(image_list):.2f}', 'per image.')