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) 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: Inverting %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, 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_estimate_pil = Tensor2PIL(y_estimate_list[img_id]) 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) 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): disc = nn.DataParallel(get_derivable_generator(args.disc_model, args.disc_type, args)).cuda() gen = nn.DataParallel(get_derivable_generator(args.gan_model, args.gan_type, args)).cuda() total_blocks = 18 if args.resolution == 1024 else 14 gen_block_ids = list(np.linspace(1, total_blocks - 4, total_blocks - 4).astype(np.int).tolist()) disc_block_ids = [gen.module.n_blocks - item for item in gen_block_ids] batch_size = args.batch_size num_data = args.num_data count = 0 for i in range(0, num_data, batch_size): print('(%d/%d)' % (count, num_data)) if i == 0: print('produce dataset of len %d. ' % num_data) with h5py.File(os.path.join(args.outputs, args.file_name), 'w') as file: z = torch.randn(size=(batch_size, args.dim_z), dtype=torch.float32).cuda() if num_data - count < batch_size: z = z[:(num_data - count)] noise_dset = file.create_dataset('noise', z.shape, dtype=np.float32, maxshape=(num_data, args.dim_z), chunks=(args.chunk_size, args.dim_z), compression=args.compression) noise_dset[...] = z.cpu().detach().numpy() h = z for layer_i in range(len(gen_block_ids)): layer_name = 'gen_layer_%d' % gen_block_ids[layer_i] if layer_i == 0: h = gen([z.reshape((batch_size, args.dim_z, 1, 1))], pre_model=True, which_block=gen_block_ids[layer_i]).detach() else: h = gen([h], free_model=True, start=gen_block_ids[layer_i - 1], end=gen_block_ids[layer_i]).detach() h_numpy = h.detach().cpu().numpy() layer_i_dset = file.create_dataset(layer_name, maxshape=(num_data, h_numpy.shape[1], h_numpy.shape[2], h_numpy.shape[3]), dtype=np.float32, chunks=(args.chunk_size, h_numpy.shape[1], h_numpy.shape[2], h_numpy.shape[3]), compression=args.compression) layer_i_dset[...] = h_numpy h = gen([h], post_model=True, which_block=gen_block_ids[-1]).detach() name = 'generation' h_numpy = h.detach().cpu().numpy() generation_dset = file.create_dataset(name, maxshape=(num_data, h_numpy.shape[1], h_numpy.shape[2], h_numpy.shape[3]), dtype=np.float32, chunks=(args.chunk_size, h_numpy.shape[1], h_numpy.shape[2], h_numpy.shape[3]), compression=args.compression) generation_dset[...] = h_numpy results = forward_discriminator(h, disc, disc_block_ids, disc_detach=True) for layer_i in range(len(results)): name = 'disc_layer_%d' % disc_block_ids[layer_i] h_numpy = results[layer_i].detach().cpu().numpy() layer_i_dset = file.create_dataset(name, maxshape=(num_data, h_numpy.shape[1], h_numpy.shape[2], h_numpy.shape[3]), dtype=np.float32, chunks=(args.chunk_size, h_numpy.shape[1], h_numpy.shape[2], h_numpy.shape[3]), compression=args.compression) layer_i_dset[...] = h_numpy else: with h5py.File(os.path.join(args.outputs, args.file_name), 'a') as file: h = torch.randn(size=(batch_size, args.dim_z), dtype=torch.float32).cuda() if num_data - count < batch_size: h = h[:(num_data - count)] h_numpy = h.detach().cpu().numpy() file['noise'].resize(file['noise'].shape[0] + h_numpy.shape[0], axis=0) file['noise'][-h_numpy.shape[0]:] = h_numpy for layer_i in range(len(gen_block_ids)): layer_name = 'gen_layer_%d' % gen_block_ids[layer_i] if layer_i == 0: h = gen([h.reshape((batch_size, args.dim_z, 1, 1))], pre_model=True, which_block=gen_block_ids[layer_i]).detach() else: h = gen([h], free_model=True, start=gen_block_ids[layer_i - 1], end=gen_block_ids[layer_i]).detach() h_numpy = h.detach().cpu().numpy() file[layer_name].resize(file[layer_name].shape[0] + h_numpy.shape[0], axis=0) file[layer_name][-h_numpy.shape[0]:] = h_numpy h = gen([h], post_model=True, which_block=gen_block_ids[-1]).detach() name = 'generation' h_numpy = h.detach().cpu().numpy() file[name].resize(file[name].shape[0] + h_numpy.shape[0], axis=0) file[name][-h_numpy.shape[0]:] = h_numpy results = forward_discriminator(h, disc, disc_block_ids, disc_detach=True) for layer_i in range(len(results)): name = 'disc_layer_%d' % disc_block_ids[layer_i] h_numpy = results[layer_i].detach().cpu().numpy() file[name].resize(file[name].shape[0] + h_numpy.shape[0], axis=0) file[name][-h_numpy.shape[0]:] = h_numpy count += batch_size
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 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() print('There are %d blocks in this generator.' % generator.n_blocks) # 19 for pggan latent_space = generator.PGGAN_LATENT[args.layer] print('The latent space is ', latent_space) Matrix = torch.ones([latent_space[0], latent_space[0]], dtype=torch.float32).cuda() * 0.0001 Matrix.requires_grad = True 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) if args.optimization == 'GD': optimizer = torch.optim.SGD([Matrix], lr=args.lr) elif args.optimization == 'Adam': optimizer = torch.optim.Adam([Matrix], lr=args.lr) else: raise NotImplemented('Not Implemented.') 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(test_zs.size()) losses_dict = {'total_loss': [], 'feature_loss': [], 'data_loss': [], 'sparse_loss': []} for iter in range(args.iterations): if args.beta0 > 0: Z = torch.randn([args.batch_size, 512, 1, 1], dtype=torch.float32).cuda() F = generator([Z], which_block=args.layer, pre_model=True) F_reshape = F.transpose(0, 1).reshape((latent_space[0], -1)) F_rec = torch.matmul(Matrix * (1. - torch.eye(n=latent_space[0], m=latent_space[0]).cuda()), F_reshape).\ reshape((latent_space[0], args.batch_size,) + tuple(latent_space[1:])).transpose(0, 1) X_rec = generator([F_rec], which_block=args.layer, post_model=True) X = generator([F], which_block=args.layer, post_model=True) optimizer.zero_grad() if args.sparse_type == 'L2': sparse_loss = torch.mean(torch.pow(Matrix * (1. - torch.eye(n=latent_space[0], m=latent_space[0]).cuda()), 2.0)) elif args.sparse_type == 'L1': sparse_loss = torch.mean(torch.abs(Matrix * (1. - torch.eye(n=latent_space[0], m=latent_space[0]).cuda()))) else: raise NotImplemented('Type not implemented.') feature_loss = torch.mean(torch.pow(F - F_rec, 2.0)) data_loss = torch.mean(torch.pow(X - X_rec, 2.0)) loss = feature_loss + args.beta0 * data_loss + args.beta1 * sparse_loss loss.backward() optimizer.step() elif args.beta0 == 0: Z = torch.randn([args.batch_size, 512, 1, 1], dtype=torch.float32).cuda() F = generator([Z], which_block=args.layer, pre_model=True) F_reshape = F.transpose(0, 1).reshape((latent_space[0], -1)) F_rec = torch.matmul(Matrix * (1. - torch.eye(n=latent_space[0], m=latent_space[0]).cuda()), F_reshape). \ reshape((latent_space[0], args.batch_size,) + tuple(latent_space[1:])).transpose(0, 1) optimizer.zero_grad() if args.sparse_type == 'L2': sparse_loss = torch.mean( torch.pow(Matrix * (1. - torch.eye(n=latent_space[0], m=latent_space[0]).cuda()), 2.0)) elif args.sparse_type == 'L1': sparse_loss = torch.mean( torch.abs(Matrix * (1. - torch.eye(n=latent_space[0], m=latent_space[0]).cuda()))) else: raise NotImplemented('Type not implemented.') feature_loss = torch.mean(torch.pow(F - F_rec, 2.0)) loss = feature_loss + args.beta1 * sparse_loss loss.backward() optimizer.step() else: raise NotImplemented() if iter % args.report_value == 0: if args.beta0 == 0: print('Iter %d, Layer %d, loss=%.6f, f_loss=%.6f, sparse_loss=%.6f' % (iter, args.layer, float(loss.item()), float(feature_loss.item()), float(sparse_loss.item()))) else: print('Iter %d, Layer %d, loss=%.6f, f_loss=%.6f, x_loss=%.6f, sparse_loss=%.6f' % (iter, args.layer, float(loss.item()), float(feature_loss.item()), float(data_loss.item()), float(sparse_loss.item()))) losses_dict['total_loss'].append(float(loss.item())) losses_dict['feature_loss'].append(float(feature_loss.item())) losses_dict['data_loss'].append(float(data_loss.item())) losses_dict['sparse_loss'].append(float(sparse_loss.item())) if iter % args.report_image == 0: # save reconstruction images. if args.beta0 == 0: X_rec = generator([F_rec], which_block=args.layer, post_model=True).detach() X = generator([F], which_block=args.layer, post_model=True).detach() test_feature = generator([test_zs[:4]], which_block=args.layer, pre_model=True).detach() test_feature_reshape = test_feature.transpose(0, 1).reshape((latent_space[0], -1)) test_feature_rec = torch.matmul(Matrix * (1. - torch.eye(n=latent_space[0], m=latent_space[0]).cuda()), test_feature_reshape). \ reshape((latent_space[0], args.batch_size,) + tuple(latent_space[1:])).transpose(0, 1) test_rec = generator([test_feature_rec], which_block=args.layer, post_model=True).detach() test_images = generator([test_feature], which_block=args.layer, post_model=True).detach() image_number = min(64, args.batch_size) Xs = torch.clamp(_tanh_to_sigmoid(torch.cat([test_rec, test_images], dim=0).detach().cpu()), min=0.0, max=1.0) torchvision.utils.save_image(Xs, os.path.join(out_dir, 'ReconstructionImages_%d.png' % iter), nrow=4) # visualize the affinity matrix. Matrix_abs = torch.relu(Matrix) S_val = Matrix_abs.detach().cpu().numpy() S_val = thrC(S_val.copy().T, args.alpha).T #S_val_norm = S_val / np.sum(S_val, axis=1, keepdims=True) #S_val_norm = S_val_norm / np.sum(S_val_norm, axis=0, keepdims=True) #S_val_visual = np.clip(S_val_norm, 0, args.times * 1/512) #S_val_visual = S_val_visual / np.max(S_val_visual, axis=1, keepdims=True) #plt.figure() #plt.imshow(S_val_visual, cmap='Oranges') #plt.savefig(os.path.join(out_dir, 'thrd_matrix_iter_%d.png' % iter)) predict, L_val = post_proC(S_val, args.n_subs, args.d_subs, args.power) p_sum = [sum(predict == k) for k in range(1, args.n_subs+1, 1)] p_sum = np.array(p_sum) print(p_sum) p_sort = np.argsort(p_sum)[::-1] predict_new = predict.copy() for i in range(1, args.n_subs+1, 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.n_subs+1, 1)] print(predict) print(p_sum) S_val_blockized = switch_matrix(S_val.copy(), predict.copy()) # S_val_blockized = np.clip(S_val_blockized, 0, S_val_blockized.mean() * args.times) torchvision.utils.save_image(torch.from_numpy(S_val_blockized + S_val_blockized.T), os.path.join(out_dir, 'SwitchedMatrix%d.png' % iter), nrow=1, normalize=True) 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).detach() features = F.detach().cpu().numpy() for class_i in range(1, args.n_subs+1, 1): ex_images = [] for ii in range(image_number): ex_rows = [] for jj in range(image_number): f_a = features[ii].copy() f_b = features[jj].copy() f_a[predict == class_i] = f_b[predict == class_i] f_a = f_a.reshape((1, ) + latent_space) ex_rows.append(f_a) ex_rows = np.concatenate(ex_rows, axis=0).astype(np.float32) ex_rows = torch.from_numpy(ex_rows).cuda() ex_ys = generator([ex_rows], which_block=args.layer, post_model=True).detach() ex_ys = torch.clamp(_tanh_to_sigmoid(ex_ys), min=0.0, max=1.0) ex_images.append(ex_ys) ex_images = torch.cat(ex_images, dim=0) torchvision.utils.save_image(ex_images.detach().cpu(), os.path.join(out_dir, 'layer_%d_class_%d_iter_%d.png' % (args.layer, class_i, iter)), nrow=image_number) # the subspace mask of class_i, save it. subspace_i = predict == class_i with open(os.path.join(out_dir, 'subspace_mask_layer%d_iter%d_class%d.pkl' % (args.layer, iter, class_i)), 'wb') as file_out: pkl.dump(subspace_i, file_out) print('Save layer %d iter %d class %d, out_dir=%s.' % (args.layer, iter, class_i, out_dir)) if iter % args.report_model == 0: with open(os.path.join(out_dir, 'value%d_layer%d.pkl'%(iter, args.layer)), 'wb') as file_out: pkl.dump(Matrix.detach().cpu().numpy(), file_out) print('Save S.') with open(os.path.join(out_dir, 'loss_dict%d_layer%d.pkl'%(iter, args.layer)), 'wb') as file_out: pkl.dump(losses_dict, file_out) print('Save dict.')
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): generator = get_derivable_generator(args.gan_model, args.inversion_type, args) loss = get_loss(args.loss_type, args) loss.cuda() generator.cuda() inversion = get_inversion(args.optimization, args) os.makedirs(args.outputs, exist_ok=True) save_manipulate_dir = os.path.join(args.outputs, args.attribute_name) images_list = image_files(args.target_images) for i, images in enumerate(images_list): print('%d: Processing images ' % (i + 1), end='') image_name = os.path.split(images)[1] print(image_name) image_name_list = [] image_tensor_list = [] image_name_list.append(os.path.split(images)[1]) image_tensor_list.append(_add_batch_one(load_as_tensor(images))) y_gt = _sigmoid_to_tanh(torch.cat(image_tensor_list, dim=0)).cuda() # Invert latent_estimates, history = inversion.invert(generator, y_gt, loss, batch_size=1) image_manipulate_dir = os.path.join(save_manipulate_dir, image_name[:-4]) os.makedirs(image_manipulate_dir, exist_ok=True) wp = latent_estimates[0].cpu().detach().numpy() mask = latent_estimates[1].cpu().detach().numpy() # Visualize results with given w+ latent vector. if args.original: print('Save inversion.') image = generator(latent_estimates) image_cv2 = convert_array_to_images(image.detach().cpu().numpy()) cv2.imwrite( os.path.join(image_manipulate_dir, 'original_inversion.png'), image_cv2[0][:, :, ::-1]) boundary, bias = get_boundary( os.path.join( BOUNDARY_DIR, 'pggan_celebahq_%s_boundary.npy' % args.attribute_name)) wp_list = get_interpolated_wp(wp, boundary, max_step=args.max_step, num_frames=args.fps * args.duration) # Create video for attribute manipulation with given w+ latent_vector. if args.video: print('Create attribute manipulation video.') video = cv2.VideoWriter(filename=os.path.join( image_manipulate_dir, '%s_manipulate.avi' % args.attribute_name), fourcc=cv2.VideoWriter_fourcc(*'MJPG'), fps=args.fps, frameSize=(1024, 1024)) print('Save frames.') for i, sample in enumerate(wp_list): image = generator([ torch.from_numpy(sample).view((1, ) + sample.shape).cuda(), torch.from_numpy(mask).cuda() ]) image_cv2 = convert_array_to_images( image.detach().cpu().numpy())[0][:, :, ::-1] cv2.imwrite(os.path.join(image_manipulate_dir, '%d.png' % i), image_cv2) 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.')