def eval(): model.eval() avg_psnr_predicted = 0.0 for batch in testing_data_loader: input, flow_f, flow_b, filename, d_dir = batch[0], batch[1], batch[2], batch[3], batch[4] with torch.no_grad(): t_im1 = Variable(input[0]).cuda(gpus_list[0]) t_im2 = Variable(input[1]).cuda(gpus_list[0]) t_flow_f = Variable(flow_f).cuda(gpus_list[0]).float() t_flow_b = Variable(flow_b).cuda(gpus_list[0]).float() t0 = time.time() if opt.chop_forward: with torch.no_grad(): pred_l = chop_forward(t_im1, t_im2, t_flow_f, t_flow_b, model) else: with torch.no_grad(): _, _, _, pred_l = model(t_im1, t_im2, t_flow_f, t_flow_b, train=False) t1 = time.time() print("===> Processing: %s || Timer: %.4f sec." % (d_dir[0]+'/frame10i11.png', (t1 - t0))) pred_l = utils.denorm(pred_l[0].cpu().data, vgg=True) pred_1 = utils.denorm(t_im1[0].cpu().data, vgg=True) pred_2 = utils.denorm(t_im2[0].cpu().data, vgg=True) if opt.data_dir == 'ucf101_interp_ours': save_img(pred_1, d_dir[0],'frame_00.png', False) save_img(pred_l, d_dir[0],'frame_01_gt.png', False) save_img(pred_2, d_dir[0],'frame_02.png', False) else: save_img(pred_1, d_dir[0],'im1.png', False) save_img(pred_l, d_dir[0],'im2.png', False) save_img(pred_2, d_dir[0],'im3.png', False)
def print_info(self, step, total_steps, pbar): current_epoch = (step + 1) / self.train_steps_per_epoch if (step + 1) % self.args.info_step == 0: elapsed_num = time.time() - self.start_time elapsed = str(datetime.timedelta(seconds=elapsed_num)) pbar.write( "Elapse:{:>.12s}, D_Step:{:>6d}/{}, G_Step:{:>6d}/{}, D_loss:{:>.4f}, G_loss:{:>.4f}, G_percep_loss:{:>.4f}, G_adv_loss:{:>.4f}, G_idt_loss:{:>.4f}" .format(elapsed, step + 1, total_steps, (step + 1), total_steps, self.d_loss, self.g_loss, self.g_percep_loss, self.g_adv_loss, self.g_idt_loss)) # sample images if (step + 1) % self.args.sample_step == 0: for i in range(0, self.real_raw.size(0)): save_imgs = torch.cat([ denorm(self.real_raw.data)[i:i + 1, :, :, :], denorm(self.fake_exp.data)[i:i + 1, :, :, :], denorm(self.real_exp.data)[i:i + 1, :, :, :] ], 3) save_image( save_imgs, os.path.join( self.sample_path, '{:s}_{:0>3.2f}_{:0>2d}_realRaw_fakeExp_realExp.png'. format(self.real_raw_name[i], current_epoch, i))) # save models if (step + 1) % self.model_save_step == 0: if self.args.parallel: if torch.cuda.device_count() > 1: checkpoint = { "G_net": self.G.module.state_dict(), "D_net": self.D.module.state_dict(), "epoch": current_epoch, "g_optimizer": self.g_optimizer.state_dict(), "d_optimizer": self.d_optimizer.state_dict(), "lr_scheduler_g": self.lr_scheduler_g.state_dict(), "lr_scheduler_d": self.lr_scheduler_d.state_dict() } else: checkpoint = { "G_net": self.G.state_dict(), "D_net": self.D.state_dict(), "epoch": current_epoch, "g_optimizer": self.g_optimizer.state_dict(), "d_optimizer": self.d_optimizer.state_dict(), "lr_scheduler_g": self.lr_scheduler_g.state_dict(), "lr_scheduler_d": self.lr_scheduler_d.state_dict() } torch.save( checkpoint, os.path.join( self.model_save_path, '{}_{}_{}.pth'.format(self.args.version, self.args.adv_loss_type, current_epoch))) pbar.write("======= Save model checkpoints into {} ======".format( self.model_save_path))
def generate_all(data_loader, device, args): """Generate Images using DIV2K Validation set""" # Single Results Path # if not os.path.exists(args.single_results_path): os.makedirs(args.single_results_path) # Prepare Networks # edsr = EDSR(channels=args.channels, features=args.dim, num_residuals=args.num_residuals, scale_factor=args.upscale_factor).to(device) # Weight Paths # edsr_weight_path = os.path.join( args.weights_path, '{}_Epoch_{}.pkl'.format(edsr.__class__.__name__, args.num_epochs)) #rdn_weight_path = os.path.join(args.weights_path, '{}_Epoch_{}.pkl'.format(rdn.__class__.__name__, args.num_epochs)) #srgan_weight_path = os.path.join(args.weights_path, '{}_Epoch_{}.pkl'.format(srgan.__class__.__name__, args.num_epochs)) #esrgan_weight_path = os.path.join(args.weights_path, '{}_Epoch_{}.pkl'.format(esrgan.__class__.__name__, args.num_epochs)) # Load State Dict # edsr.load_state_dict(torch.load(edsr_weight_path)) # rdn.load_state_dict(torch.load(rdn_weight_path)) # srgan.load_state_dict(torch.load(srgan_weight_path)) # esrgan.load_state_dict(torch.load(esrgan_weight_path)) # Up-sampling Network # up_sampler = torch.nn.Upsample(scale_factor=args.upscale_factor, mode='bicubic').to(device) for i, (high, low) in enumerate(data_loader): # Prepare Data # high = high.to(device) low = low.to(device) # Forward Data to Networks # with torch.no_grad(): bicubic = up_sampler(low.detach()) generated_edsr = edsr(low.detach()) # generated_rdn = rdn(low.detach()) # generated_srgan = srgan(low.detach()) # generated_esrgan = esrgan(low.detach()) # Normalize and Save Images # save_image( denorm(high.data), os.path.join(args.single_results_path, 'Inference_Samples_%03d_TARGET.png' % (i + 1))) save_image( denorm(bicubic.data), os.path.join(args.single_results_path, 'Inference_Samples_%03d_BICUBIC.png' % (i + 1))) save_image( denorm(generated_edsr.data), os.path.join( args.single_results_path, 'Inference_Samples_%03d_%s.png' % (i + 1, edsr.__class__.__name__)))
def train(epoch): epoch_loss = 0 model.train() feature_extractor.eval() for iteration, batch in enumerate(training_data_loader, 1): target, flow_f, flow_b, gt_flow_f, gt_flow_b = batch[0], batch[ 1], batch[2], batch[3], batch[4] if cuda: t_flow_f = Variable(flow_f).cuda(gpus_list[0]).float() t_flow_b = Variable(flow_b).cuda(gpus_list[0]).float() gt_flow_f = Variable(gt_flow_f).cuda(gpus_list[0]).float() gt_flow_b = Variable(gt_flow_b).cuda(gpus_list[0]).float() t_hr1 = Variable(target[0]).cuda(gpus_list[0]) t_hr = Variable(target[1]).cuda(gpus_list[0]) t_hr2 = Variable(target[2]).cuda(gpus_list[0]) optimizer.zero_grad() t0 = time.time() _, _, _, pred_l, pred_flow_f, pred_flow_b = model( t_hr1, t_hr2, t_flow_f, t_flow_b) l_mse_lr = mse_loss_calc(pred_l, t_hr) l_mse = l_mse_lr l_feat_lr = feat_loss_calc(pred_l, t_hr) l_feat = l_feat_lr l_flow = criterion(pred_flow_f, gt_flow_f) + criterion( pred_flow_b, gt_flow_b) #loss = l_flow loss = l_mse + 0.1 * l_feat + 0.1 * l_flow t1 = time.time() ###show sample pred_l = utils.denorm(pred_l[0][0].cpu().data, vgg=True) pred_l1 = utils.denorm(t_hr1[0].cpu().data, vgg=True) pred_l2 = utils.denorm(t_hr2[0].cpu().data, vgg=True) save_img(pred_l, "lr2") save_img(pred_l1, "lr1") save_img(pred_l2, "lr3") epoch_loss += loss.data loss.backward() optimizer.step() #print(loss_mse_hr.data , loss_mse_lr.data, loss_feat_hr.data, loss_feat_l.data) print( "===> Epoch[{}]({}/{}): Loss: {:.4f} MSE_LR: {:.4f} FEAT_LR: {:.4f} Flow: {:.8f} || Timer: {:.4f} sec." .format(epoch, iteration, len(training_data_loader), loss.data, l_mse_lr, l_feat_lr, l_flow, (t1 - t0))) print("===> Epoch {} Complete: Avg. Loss: {:.4f}".format( epoch, epoch_loss / len(training_data_loader)))
def loss_function(batch_x, batch_y): logits = model(batch_x, training=True) denorm_x = denorm(logits, _min, _max) denorm_y = denorm(batch_y, _min, _max) lossL2 = tf.add_n(model.losses) return rmse(denorm_x, denorm_y) + lossL2
def train(epoch): epoch_loss = 0 model.train() feature_extractor.eval() for iteration, batch in enumerate(training_data_loader, 1): input, target, target_l, flow_f, flow_b = batch[0], batch[1], batch[2], batch[3], batch[4] if cuda: t_im1 = Variable(input[0]).cuda(gpus_list[0]) t_im2 = Variable(input[1]).cuda(gpus_list[0]) t_flow_f = Variable(flow_f).cuda(gpus_list[0]).float() t_flow_b = Variable(flow_b).cuda(gpus_list[0]).float() t_target1 = Variable(target[0]).cuda(gpus_list[0]) t_target = Variable(target[1]).cuda(gpus_list[0]) t_target2 = Variable(target[2]).cuda(gpus_list[0]) t_target_l = Variable(target_l).cuda(gpus_list[0]) optimizer.zero_grad() t0 = time.time() pred_ht, pred_h1, pred_h2, pred_l = model(t_im1, t_im2, t_flow_f, t_flow_b) l_mse_ht = mse_loss_calc(pred_ht, t_target) l_mse_h1 = mse_loss_calc(pred_h1, t_target1) l_mse_h2 = mse_loss_calc(pred_h2, t_target2) l_mse_lr = mse_loss_calc(pred_l, t_target_l) l_mse = l_mse_ht + 0.5*l_mse_h1 + 0.5*l_mse_h2 + l_mse_lr l_feat_ht = feat_loss_calc(pred_ht, t_target) l_feat_h1 = feat_loss_calc(pred_h1, t_target1) l_feat_h2 = feat_loss_calc(pred_h2, t_target2) l_feat_lr = feat_loss_calc(pred_l, t_target_l) l_feat = l_feat_ht + 0.5*l_feat_h1 + 0.5*l_feat_h2 + l_feat_lr loss = l_mse + 0.1*l_feat t1 = time.time() ###show sample predictiont = utils.denorm(pred_ht[0][0].cpu().data,vgg=True) prediction1 = utils.denorm(pred_h1[0][0].cpu().data,vgg=True) prediction2 = utils.denorm(pred_h2[0][0].cpu().data,vgg=True) pred_l = utils.denorm(pred_l[0][0].cpu().data,vgg=True) save_img(prediction1, "1") save_img(predictiont, "2") save_img(prediction2, "3") save_img(pred_l, "lr") epoch_loss += loss.data loss.backward() optimizer.step() print("===> Epoch[{}]({}/{}): Loss: {:.4f} MSE_HR: {:.4f} Feat_HR: {:.4f} MSE_LR: {:.4f} || Timer: {:.4f} sec.".format(epoch, iteration, len(training_data_loader), loss.data, l_mse_ht, l_feat_ht, l_mse_lr, (t1 - t0))) print("===> Epoch {} Complete: Avg. Loss: {:.4f}".format(epoch, epoch_loss / len(training_data_loader)))
def test(self): model_list = glob( os.path.join(self.result_dir, self.dataset, 'model', '*.pt')) if not len(model_list) == 0: model_list.sort() iter = int(model_list[-1].split('_')[-1].split('.')[0]) self.load(os.path.join(self.result_dir, self.dataset, 'model'), iter) print(" [*] Load SUCCESS") else: print(" [*] Load FAILURE") return self.genA2B.eval(), self.genB2A.eval() for n, (real_A, _) in enumerate(self.testA_loader): real_A = real_A.to(self.device) fake_A2B, _, fake_A2B_heatmap = self.genA2B(real_A) fake_A2B2A, _, fake_A2B2A_heatmap = self.genB2A(fake_A2B) fake_A2A, _, fake_A2A_heatmap = self.genB2A(real_A) A2B = np.concatenate(( RGB2BGR(tensor2numpy(denorm(real_A[0]))), cam(tensor2numpy(fake_A2A_heatmap[0]), self.img_size), RGB2BGR(tensor2numpy(denorm(fake_A2A[0]))), cam(tensor2numpy(fake_A2B_heatmap[0]), self.img_size), RGB2BGR(tensor2numpy(denorm(fake_A2B[0]))), cam(tensor2numpy(fake_A2B2A_heatmap[0]), self.img_size), RGB2BGR(tensor2numpy(denorm(fake_A2B2A[0]))), ), 0) cv2.imwrite( os.path.join(self.result_dir, self.dataset, 'test', 'A2B_%d.png' % (n + 1)), A2B * 255.0) for n, (real_B, _) in enumerate(self.testB_loader): real_B = real_B.to(self.device) fake_B2A, _, fake_B2A_heatmap = self.genB2A(real_B) fake_B2A2B, _, fake_B2A2B_heatmap = self.genA2B(fake_B2A) fake_B2B, _, fake_B2B_heatmap = self.genA2B(real_B) B2A = np.concatenate(( RGB2BGR(tensor2numpy(denorm(real_B[0]))), cam(tensor2numpy(fake_B2B_heatmap[0]), self.img_size), RGB2BGR(tensor2numpy(denorm(fake_B2B[0]))), cam(tensor2numpy(fake_B2A_heatmap[0]), self.img_size), RGB2BGR(tensor2numpy(denorm(fake_B2A[0]))), cam(tensor2numpy(fake_B2A2B_heatmap[0]), self.img_size), RGB2BGR(tensor2numpy(denorm(fake_B2A2B[0]))), ), 0) cv2.imwrite( os.path.join(self.result_dir, self.dataset, 'test', 'B2A_%d.png' % (n + 1)), B2A * 255.0)
def inference(): # Inference Path # make_dirs(config.inference_path) # Prepare Data Loader # test_loader_selfie, test_loader_anime = get_selfie2anime_loader('test', config.batch_size) # Prepare Generator # G_A2B = Generator(image_size=config.crop_size, num_blocks=config.num_blocks).to(device) G_A2B.load_state_dict(torch.load(os.path.join(config.weights_path, 'U-GAT-IT_G_A2B_Epoch_{}.pkl'.format(config.num_epochs)))) # Inference # print("U-GAT-IT | Generating Selfie2Anime images started...") with torch.no_grad(): for i, (selfie, anime) in enumerate(zip(test_loader_selfie, test_loader_anime)): # Prepare Data # real_A = selfie.to(device) # Generate Fake Images # fake_B = G_A2B(real_A)[0] # Save Images (Selfie -> Anime) # result = torch.cat((real_A, fake_B), dim=0) save_image(denorm(result.data), os.path.join(config.inference_path, 'U-GAT-IT_Selfie2Anime_Results_%03d.png' % (i + 1)) ) # Make a GIF file # make_gifs_test("U-GAT-IT", "Selfie2Anime", config.inference_path)
def interpolate(model, device, latent_dim, hair_classes, eye_classes, samples=10): z1 = torch.randn(1, latent_dim).to(device) h1 = torch.zeros(1, hair_classes).to(device) e1 = torch.zeros(1, eye_classes).to(device) h1[0][np.random.randint(hair_classes)] = 1 e1[0][np.random.randint(eye_classes)] = 1 c1 = torch.cat((h1, e1), 1) z2 = torch.randn(1, latent_dim).to(device) h2 = torch.zeros(1, hair_classes).to(device) e2 = torch.zeros(1, eye_classes).to(device) h2[0][np.random.randint(hair_classes)] = 1 e2[0][np.random.randint(eye_classes)] = 1 c2 = torch.cat((h2, e2), 1) z_diff = z2 - z1 c_diff = c2 - c1 z_step = z_diff / (samples + 1) c_step = c_diff / (samples + 1) img_list = [] for i in range(0, samples + 2): z = z1 + z_step * i c = c1 + c_step * i img_list.append(model(z, c)) output = torch.cat(img_list, 0) save_image(utils.denorm(output), '{}/interpolation_{}.png'.format(args.sample_dir, args.index), nrow=samples + 2)
def generate_faces(): # Test Path # make_dirs(config.inference_path) # Prepare Generator # G = Generator().to(device) G.load_state_dict( torch.load( os.path.join( config.weights_path, 'Face_Generator_Epoch_{}.pkl'.format(config.num_epochs)))) G.eval() # Start Generating Faces # count = 1 while (True): # Prepare Fixed Noise and Generator # noise = torch.randn(config.batch_size, config.noise_dim).to(device) generated = G(noise) for i in range(config.batch_size): save_image( denorm(generated[i].data), os.path.join(config.inference_path, "Generated_CelebA_Faces_{}.png".format(count)), ) count += 1 if count > config.limit: print("Generating fake CelebA faces is finished.") break
def main(): latent_dim = 200 hair_classes, eye_classes, face_classes, glasses_classes = 6, 4, 3, 2 num_classes = hair_classes + eye_classes + face_classes + glasses_classes G_path = './checkpoint/ACGAN_generator_FID166.ckpt' prev_state = torch.load(G_path, map_location=device) G = Generator(latent_dim=latent_dim, class_dim=num_classes).to(device) G.load_state_dict(prev_state['model']) G = G.eval() # root_dir = '../sample_test/sample_fid_testing_labels.txt' root_dir = args.label_dir transform = Transform.Compose([Transform.ToTensor(), Transform.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))]) dataset = Anime(root_dir=root_dir, transform=transform, batch_size=64) for i in tqdm(range(dataset.length())): # print("Step: ", i) hair_tags, eye_tags, face_tags, glasses_tags = dataset.get_item(i) hair_tags, eye_tags, face_tags, glasses_tags = hair_tags.to(device), \ eye_tags.to(device), \ face_tags.to(device), \ glasses_tags.to(device) z = torch.randn(1, latent_dim).to(device) fake_tag = torch.cat((hair_tags, eye_tags, face_tags, glasses_tags)).unsqueeze(0).float().to(device) fake_img = G(z, fake_tag).to(device) save_image(utils.denorm(fake_img), '{}/{}.png'.format(args.output_dir, i))
def main(): latent_dim = 200 device = args.device hair_classes, eye_classes, face_classes, glasses_classes = 6, 4, 3, 2 num_classes = hair_classes + eye_classes + face_classes + glasses_classes # G_path = '{}/ACGAN-batch_size-[64]-steps-[{}]/G_{}.ckpt'.format(args.model_dir, args.steps, args.n_ckpt) G_path = '{}/G_{}.ckpt'.format(args.model_dir, args.n_ckpt) prev_state = torch.load(G_path) G = Generator(latent_dim=latent_dim, class_dim=num_classes).to(device) G.load_state_dict(prev_state['model']) G = G.eval() if not os.path.exists('{}/fid_gen_steps{}_with_G_{}'.format( args.out_dir, args.steps, args.n_ckpt)): os.makedirs('{}/fid_gen_steps{}_with_G_{}'.format( args.out_dir, args.steps, args.n_ckpt)) z = torch.randn(1, latent_dim).to(device) # fix noise img_i = 0 with open('./sample_test/sample_fid_testing_labels.txt') as f: lines = f.readlines() for line in lines[2:]: label = list(map(int, line.split(' '))) fake_tag = torch.tensor(label).unsqueeze(0).float().to(device) fake_img = G(z, fake_tag).to(device) save_image( utils.denorm(fake_img), '{}/fid_gen_steps{}_with_G_{}/{}.png'.format( args.out_dir, args.steps, args.n_ckpt, img_i)) print('{}.png saved'.format(img_i)) img_i += 1
def __init__(self, args, model, criterion): self.args = args self.model = model self.criterion = criterion self.use_cuda = torch.cuda.is_available() if self.use_cuda: self.model = self.model.to("cuda") self.criterion = self.criterion.to("cuda") # optimizer if args.train_what == 'last': # fine-tune???? for name, param in model.module.resnet.named_parameters(): param.requires_grad = False else: pass # train all layers # print('\n===========Check Grad============') # for name, param in model.named_parameters(): # print(name, param.requires_grad) # print('=================================\n') params = list( filter(lambda p: p.requires_grad, self.model.parameters())) self.optimizer = torch.optim.Adam(params, lr=args.lr, weight_decay=args.wd) print('| num. module params: {} (num. trained: {})'.format( sum(p.numel() for p in params), sum(p.numel() for p in params if p.requires_grad), )) self.de_normalize = denorm()
def eval(): model.eval() for batch in testing_data_loader: input, name = batch[0], batch[2] input[0] = utils.norm(input[0], vgg=True) with torch.no_grad(): input = Variable(input) if cuda: input = input.cuda(gpus_list[0]) t0 = time.time() if opt.chop_forward: with torch.no_grad(): prediction = chop_forward(input, model, opt.upscale_factor) else: if opt.self_ensemble: with torch.no_grad(): prediction = x8_forward(input, model) else: with torch.no_grad(): prediction = model(input) t1 = time.time() print("===> Processing: %s || Timer: %.4f sec." % (name[0], (t1 - t0))) prediction = utils.denorm(prediction.data[0], vgg=True) save_img(prediction.cpu(), name[0])
def train(self): """overrode default base function for custom function - logs total duration of training in seconds """ tik = time.time() for it in trange(self.config.num_iter): g_loss, d_loss, d_loss1, d_loss2, d_loss3 = self.train_step() if it % self.config.save_iter == 0: self.model.save(self.sess) if it % self.config.sample_iter == 0: images = self.sess.run([self.model.sample_image]) for i, image in enumerate(images[0]): image = denorm(np.squeeze(image)) sample_path = os.path.join( self.config.sample_dir, '{}-{}-sample.jpg'.format(i, it)) skimage.io.imsave(sample_path, image) if it % 100 == 0: summaries_dict = {} summaries_dict['g_loss'] = g_loss summaries_dict['d_loss'] = d_loss summaries_dict['d_real_loss'] = d_loss1 summaries_dict['d_wrong_loss'] = d_loss2 summaries_dict['d_fake_loss'] = d_loss3 self.logger.summarize(it, summaries_dict=summaries_dict) tok = time.time() logging.info('Duration: {} seconds'.format(tok - tik))
def generate_images(out_dir, G, n_images, config, device): im_batch_size = config['n_classes'] z_bound = config['trunc_z'] if z_bound > 0.0: print(f'Truncating z to (-{z_bound}, {z_bound})') for i_batch in range(0, n_images, im_batch_size): with torch.no_grad(): if z_bound > 0.0: z = trunc_trick(im_batch_size, G.dim_z, bound=z_bound).to(device) else: z = torch.randn(im_batch_size, G.dim_z, device=device) y = torch.arange(im_batch_size).to(device) images = G(z, G.shared(y)).float().cpu() if i_batch + im_batch_size > n_images: n_last_images = n_images - i_batch print(f'Taking only {n_last_images} images from the last batch...') images = images[:n_last_images] for i_image, image in enumerate(images): fname = os.path.join(out_dir, f'image_{i_batch+i_image:05d}.png') image = utils.denorm(image) torchvision.utils.save_image(image, fname)
def logging(self, step): self.loss = {} self.images = {} self.loss['D/Total'] = self.d_loss self.loss['G/Total'] = self.g_loss self.loss['G/adv_loss'] = self.g_adv_loss self.loss['G/percep_loss'] = self.g_percep_loss self.loss['G/idt_loss'] = self.g_idt_loss self.images['Train_realExpIdt/realExp_realExpIdt'] = torch.cat([ denorm(self.real_exp.cpu()), denorm(self.real_exp_idt.detach().cpu()) ], 3) self.images['Train_compare/realRaw_fakeExp_realExp'] = torch.cat([ denorm(self.real_raw.cpu()), denorm(self.fake_exp.detach().cpu()), denorm(self.real_exp.cpu()) ], 3) self.images['Train_fakeExp/fakeExp'] = denorm( self.fake_exp.detach().cpu()) self.images['Train_fakeExpStore/fakeExpStore'] = denorm( self.fake_exp_store.detach().cpu()) if (step + 1) % self.args.log_step == 0: if self.args.use_tensorboard: for tag, value in self.loss.items(): self.logger.scalar_summary(tag, value, step + 1) for tag, image in self.images.items(): self.logger.images_summary(tag, image, step + 1)
def main(): f = open("./sample_test/sample_human_testing_labels.txt", 'r') num_of_samples = f.readline() print(f.readline()) tags_list = [] #results=[] for line in f: #print(line) #print(line.split()) line = line.split() #print(line[0]) tags = list(map(int, line[:])) # print(tags) #print(len(tags)) #print(np.asarray(tags)) tags_list.append(torch.from_numpy(np.asarray(tags))) # print(tags_list) # input() f.close() #print(len(tags_list)) #correct #input() if not os.path.exists(args.sample_dir): os.mkdir(args.sample_dir) latent_dim = 100 hair_classes = 6 eye_classes = 4 face_classes = 3 glass_classes = 2 batch_size = 1 device = 'cuda' G_path = '{}/G_50000.ckpt'.format(args.model_dir) #load trained generator here G = ACGAN_split.Generator(latent_dim=latent_dim, class_dim=hair_classes + eye_classes + face_classes + glass_classes) prev_state = torch.load(G_path) G.load_state_dict(prev_state['model']) G = G.eval() G = G.to(device) num = 0 for tag in tags_list: #randomly generate a noise z = torch.randn(latent_dim).unsqueeze(0).to(device) #print(z.shape) tag = tag.unsqueeze(0).to(device) tag = tag.type(torch.float32) #print(tag.shape) #print(tag) #input() img = G(z, tag) save_image(utils.denorm(img), '{}/{}.png'.format(args.sample_dir, num)) num += 1 print(num)
def inference(): # Inference Path # make_dirs(config.inference_path) # Prepare Data Loader # val_loader = get_edges2handbags_loader('val', config.val_batch_size) # Prepare Generator # G = Generator(z_dim=config.z_dim).to(device) G.load_state_dict( torch.load( os.path.join( config.weights_path, 'BicycleGAN_Generator_Epoch_{}.pkl'.format( config.num_epochs)))) G.eval() # Fixed Noise # fixed_noise = torch.randn(config.test_size, config.num_images, config.z_dim).to(device) # Test # print("BiCycleGAN | Generating Edges2Handbags Images started...") for iters, (sketch, ground_truth) in enumerate(val_loader): # Prepare Data # N = sketch.size(0) sketch = sketch.to(device) results = torch.FloatTensor(N * (1 + config.num_images), 3, config.crop_size, config.crop_size) # Generate Fake Images # for i in range(N): results[i * (1 + config.num_images)] = sketch[i].data for j in range(config.num_images): image = sketch[i].unsqueeze(dim=0) noise_to_generator = fixed_noise[i, j, :].unsqueeze(dim=0) out = G(image, noise_to_generator) results[i * (1 + config.num_images) + j + 1] = out # Save Images # save_image( denorm(results.data), os.path.join( config.inference_path, 'BicycleGAN_Edges2Handbags_Results_%03d.png' % (iters + 1)), nrow=(1 + config.num_images), ) make_gifs_test("BicycleGAN", config.inference_path)
def generate_celeba_synthesize(sfs_net_model, dl, train_epoch_num=0, use_cuda=False, out_folder=None, wandb=None): # debugging flag to dump image fix_bix_dump = 0 recon_loss = nn.L1Loss() if use_cuda: recon_loss = recon_loss.cuda() tloss = 0 # Total loss rloss = 0 # Reconstruction loss for bix, data in enumerate(dl): face = data if use_cuda: face = face.cuda() # predicted_face == reconstruction predicted_normal, predicted_albedo, predicted_sh, predicted_shading, predicted_face = sfs_net_model( face) # save predictions in log folder file_name = out_folder + str(train_epoch_num) + '_' + str(bix) # log images predicted_normal = denorm(predicted_normal) save_image(predicted_normal, path=file_name + '_normal.png') save_image(predicted_albedo, path=file_name + '_albedo.png') save_image(predicted_shading, path=file_name + '_shading.png') save_image(predicted_face, path=file_name + '_recon.png') save_image(face, path=file_name + '_face.png') np.savetxt(file_name + '_light.txt', predicted_sh.cpu().detach().numpy(), delimiter='\t') # Loss computation # Reconstruction loss total_loss = recon_loss(predicted_face, face) # Logging for display and debugging purposes tloss += total_loss.item() len_dl = len(dl) f = open(out_folder + 'readme.txt', 'w') f.write('Average Reconstruction Loss: ' + str(tloss / len_dl)) f.close() # return average loss over dataset return tloss / len_dl
def save_test_images(opt, preds, batch_dict, path, index): preds = preds.cpu().detach() if opt.dataset == 'hurricane': gt = batch_dict['orignal_data_to_predict'].cpu().detach() else: gt = batch_dict['data_to_predict'].cpu().detach() b, t, c, h, w = gt.shape if opt.input_norm: preds = utils.denorm(preds) gt = utils.denorm(gt) os.makedirs(os.path.join(path, 'pred'), exist_ok=True) os.makedirs(os.path.join(path, 'gt'), exist_ok=True) for i in range(b): for j in range(t): save_image(preds[i, j, ...], os.path.join(path, 'pred', f"pred_{index + i:03d}_{j:03d}.png")) save_image(gt[i, j, ...], os.path.join(path, 'gt', f"gt_{index + i:03d}_{j:03d}.png"))
def inference(): # Inference Path # make_dirs(config.inference_path) # Prepare Data Loader # val_loader = get_edges2shoes_loader(purpose='val', batch_size=config.val_batch_size) # Prepare Generator # G_A2B = Generator().to(device) G_B2A = Generator().to(device) G_A2B.load_state_dict( torch.load( os.path.join( config.weights_path, 'DiscoGAN_Generator_A2B_Epoch_{}.pkl'.format( config.num_epochs)))) G_B2A.load_state_dict( torch.load( os.path.join( config.weights_path, 'DiscoGAN_Generator_B2A_Epoch_{}.pkl'.format( config.num_epochs)))) # Test # print("DiscoGAN | Generating Edges2Shoes images started...") for i, (real_A, real_B) in enumerate(val_loader): # Prepare Data # real_A = real_A.to(device) real_B = real_B.to(device) # Generate Fake Images # fake_B = G_A2B(real_A) fake_A = G_B2A(real_B) # Generated Reconstructed Images # fake_ABA = G_B2A(fake_B) fake_BAB = G_A2B(fake_A) # Save Images # result = torch.cat( (real_A, fake_A, fake_BAB, real_B, fake_B, fake_ABA), dim=0) save_image(denorm(result.data), os.path.join( config.inference_path, 'DiscoGAN_Edges2Shoes_Results_%03d.png' % (i + 1)), nrow=8, normalize=True) # Make a GIF file # make_gifs_test("DiscoGAN", config.inference_path)
def test_change(self): model_list = os.listdir( os.path.join(self.result_dir, self.dataset, 'model')) if not len(model_list) == 0: model_list.sort() iter = int(model_list[-1].split('/')[-1]) self.load(os.path.join(self.result_dir, self.dataset, 'model'), iter) print("[*] Load SUCCESS") else: print("[*] Load FAILURE") return self.genA2B.eval(), self.genB2A.eval() for n, (real_A, fname) in enumerate(self.testA_loader()): real_A = np.array([real_A[0].reshape(3, 256, 256)]).astype("float32") real_A = to_variable(real_A) fake_A2B, _, _ = self.genA2B(real_A) A2B = RGB2BGR(tensor2numpy(denorm(fake_A2B[0]))) cv2.imwrite( os.path.join( self.result_dir, self.dataset, 'test', 'testA2B', '%s_fake.%s' % (fname.split('.')[0], fname.split('.')[-1])), A2B * 255.0) for n, (real_B, fname) in enumerate(self.testB_loader()): real_B = np.array([real_B[0].reshape(3, 256, 256)]).astype("float32") real_B = to_variable(real_B) fake_B2A, _, _ = self.genB2A(real_B) B2A = RGB2BGR(tensor2numpy(denorm(fake_B2A[0]))) cv2.imwrite( os.path.join( self.result_dir, self.dataset, 'test', 'testB2A', '%s_fake.%s' % (fname.split('.')[0], fname.split('.')[-1])), B2A * 255.0)
def save_images(engine): if needs_save: if engine.state.epoch % config.save.save_epoch_interval == 0: image = norm(engine.state.batch['image']) with torch.no_grad(): z, _, _ = E(image) x_r = D(z) x_p = D(fixed_z) image = denorm(image).detach().cpu() x_r = denorm(x_r).detach().cpu() x_p = denorm(x_p).detach().cpu() image = image[:config.save.n_save_images, ...] x_r = x_r[:config.save.n_save_images, ...] x_p = x_p[:config.save.n_save_images, ...] save_path = os.path.join( output_dir, 'result_{}.png'.format(engine.state.epoch)) save_image(torch.cat([image, x_r, x_p]).data, save_path)
def generate_images(model, label, output_dir, device): #torch.manual_seed(42) all_tag = torch.FloatTensor(label).to(device) #z = torch.randn(1, 100).to(device) # tag = all_tag[:64,:] # output_img = model(z, tag) # save_image(utils.denorm(output_img), './test.png', nrow=16) for i in tqdm(range(len(label))): z = torch.randn(1, 100).to(device) tag = all_tag[i].unsqueeze(0) output = model(z, tag) save_image(utils.denorm(output), os.path.join(output_dir, '{}.png'.format(i)))
def test_single(self, img_fn): # networks self.G = Generator(num_channels=self.num_channels, base_filter=64, num_residuals=16) if self.gpu_mode: self.G.cuda() # load model self.load_model() # load data img = Image.open(img_fn).convert('RGB') if self.num_channels == 1: img = img.convert('YCbCr') img_y, img_cb, img_cr = img.split() input = ToTensor()(img_y) y_ = Variable(utils.norm(input.unsqueeze(1), vgg=True)) else: input = ToTensor()(img).view(1, -1, img.height, img.width) y_ = Variable(utils.norm(input, vgg=True)) if self.gpu_mode: y_ = y_.cuda() # prediction self.G.eval() recon_img = self.G(y_) recon_img = utils.denorm(recon_img.cpu().data[0].clamp(0, 1), vgg=True) recon_img = ToPILImage()(recon_img) if self.num_channels == 1: # merge color channels with super-resolved Y-channel recon_y = recon_img recon_cb = img_cb.resize(recon_y.size, Image.BICUBIC) recon_cr = img_cr.resize(recon_y.size, Image.BICUBIC) recon_img = Image.merge( 'YCbCr', [recon_y, recon_cb, recon_cr]).convert('RGB') # save img result_dir = os.path.join(self.save_dir, 'test_result') if not os.path.exists(result_dir): os.makedirs(result_dir) save_fn = result_dir + '/SR_result.png' recon_img.save(save_fn) print('Single test result image is saved.') return recon_img
def test(self): nrow = 8 n_samples = nrow * nrow with torch.no_grad(): for i, (x_real, noise, label) in enumerate(tqdm(self.test_loader)): if i == 10: break x_real = x_real.to(self.device) noise = noise.to(self.device) label = label.to(self.device) ''' test flop ''' if i == 0: from thop import profile flops, params = profile(self.G, inputs=(noise[0].unsqueeze(0), label[0].unsqueeze(0))) print( '=======================================================================' ) print('FLOPS: {:.2f} B, Params.: {:.1f} M'.format( flops / 10**9, params / 10**6)) print( '=======================================================================' ) # recon x_fake = self.G(noise, label) comparison = torch.cat([x_real[:nrow], x_fake[:nrow].float()]) sample_path = os.path.join(self.result_dir, '{}-rec.png'.format(i + 1)) save_image(utils.denorm(comparison.cpu()), sample_path) # sample noise2 = torch.FloatTensor(utils.truncated_normal(n_samples*self.z_dim)) \ .view(n_samples, self.z_dim).to(self.device) label = label[:nrow].repeat(nrow) x_sample = self.G(noise2, label) sample_path = os.path.join(self.result_dir, '{}-sample.png'.format(i + 1)) save_image(utils.denorm(x_sample), sample_path, nrow=nrow)
def project_on_generator( G: Generator, pix2pix: networks.UnetGenerator, target_image: np.ndarray, E: Encoder, dcgan_img_size: int = 64, pix2pix_img_size: int = 128) -> Tuple[np.ndarray, torch.Tensor]: """Projects the input image onto the manifold span by the GAN. It operates as follows: 1. reshape and normalize the image 2. run the encoder to obtain a latent vector 3. run the DCGAN generator to obtain a low resolution image 4. run the Pix2Pix model to obtain a high resulution image Arguments: G {Generator} -- DCGAN generator pix2pix {networks.UnetGenerator} -- Low resolution to high resolution Pix2Pix model target_image {np.ndarray} -- The image to project E {Encoder} -- The DCGAN encoder Keyword Arguments: dcgan_img_size {int} -- Low resolution image size (default: {64}) pix2pix_img_size {int} -- High resolution image size (default: {128}) Returns: Tuple[np.ndarray, torch.Tensor] -- The projected high resolution image and the latent vector that was used to generate it. """ # reshape and normalize image target_image = torch.Tensor(target_image).cuda().reshape( 1, 3, pix2pix_img_size, pix2pix_img_size) target_image = F.interpolate(target_image, scale_factor=dcgan_img_size / pix2pix_img_size, mode='bilinear') target_image = target_image.clamp(min=0) target_image = target_image / target_image.max() target_image = (target_image - 0.5) / 0.5 # Run dcgan z = E(target_image) dcgan_image = G(z) # run pix2pix pix_input = F.interpolate(dcgan_image, scale_factor=pix2pix_img_size / dcgan_img_size, mode='bilinear') pix_outputs = pix2pix(pix_input) out_image = utils.denorm(pix_outputs.detach()).clamp( 0, 1).cpu().numpy().reshape(3, -1, 1) return out_image, z
def generateUsingHairEye(model, device, hairClasses, eyeClasses, lDim, hColor, eColor): vecSize=64 htag = torch.zeros(vecSize, hairClasses).to(device) etag = torch.zeros(vecSize, eyeClasses).to(device) hairLabelIndex = hair_dict[hColor] eyeLabelIndex = eye_dict[eColor] for i in range(vecSize): htag[i][hairLabelIndex] = 1 etag[i][eyeLabelIndex] = 1 fulltag = torch.cat((htag, etag), 1) z = torch.randn(vecSize, lDim).to(device) output = model(z, fulltag) save_image(utils.denorm(output), '../generated/{} hair {} eyes.png'.format(hairClassesList[hairLabelIndex], eyeClassesList[eyeLabelIndex]))
def make_save_sequence(opt, batch_dict, res): """ 4 cases: (interp, extrap) | (regular, irregular) """ b, t, c, h, w = batch_dict['observed_data'].size() # Filter out / Select by mask if opt.irregular: observed_mask = batch_dict["observed_mask"] mask_predicted_data = batch_dict["mask_predicted_data"] selected_timesteps = int(observed_mask[0].sum()) if opt.dataset in ['hurricane']: batch_dict['observed_data'] = batch_dict['observed_data'][observed_mask.squeeze(-1).byte(), ...].view(b, selected_timesteps, c, h, w) batch_dict['data_to_predict'] = batch_dict['data_to_predict'][mask_predicted_data.squeeze(-1).byte(), ...].view(b, selected_timesteps, c, h, w) else: batch_dict['observed_data'] = batch_dict['observed_data'] * observed_mask.unsqueeze(-1).unsqueeze(-1) batch_dict['data_to_predict'] = batch_dict['data_to_predict'] * mask_predicted_data.unsqueeze(-1).unsqueeze(-1) # Make sequence to save pred = res['pred_y'].cpu().detach() if opt.extrap: inputs = batch_dict['observed_data'].cpu().detach() gt_to_predict = batch_dict['data_to_predict'].cpu().detach() gt = torch.cat([inputs, gt_to_predict], dim=1) else: gt = batch_dict['data_to_predict'].cpu().detach() time_steps = None if opt.input_norm: gt = utils.denorm(gt) pred = utils.denorm(pred) return gt, pred, time_steps