def test(): opt.device = 'cuda:0' opt.data_root = 'demo/input/' # The location of your testing data opt.mask_root = 'demo/mask/' # The location of your testing data mask testset = MyDataLoader(opt) print('Test with %d' % (len(testset))) model = MyModel() model.initialize(opt) model.load_networks('places_irregular') # For irregular mask inpainting # model.load_networks('celebahq_center') # For centering mask inpainting, i.e., 120*120 hole in 256*256 input val_ssim, val_psnr, val_mae, val_losses_G = [], [], [], [] with torch.no_grad(): for i, data in enumerate(testset): fname = data['fname'][0] model.set_input(data) I_g, I_o, val_loss_G = model.optimize_parameters(val=True) val_s, val_p, val_m = metrics(I_g, I_o) val_ssim.append(val_s) val_psnr.append(val_p) val_mae.append(val_m) val_losses_G.append(val_loss_G.detach().item()) cv2.imwrite('demo/output/' + fname[:-4] + '.png', postprocess(I_o).numpy()[0]) print('Val (%d/%d) G:%5.4f, S:%4.4f, P:%4.2f, M:%4.4f' % ( i + 1, len(testset), np.mean(val_losses_G), np.mean(val_ssim), np.mean(val_psnr), np.mean(val_mae)), end='\r') print('Val G:%5.4f, S:%4.4f, P:%4.2f, M:%4.4f' % (np.mean(val_losses_G), np.mean(val_ssim), np.mean(val_psnr), np.mean(val_mae)))
def train(): opt.device = 'cuda:0' opt.data_root = 'demo/input/' # The location of your training data opt.mask_root = 'demo/mask/' # The location of your training data mask train_set = MyDataLoader(opt) opt.data_root = 'demo/input/' # The location of your validation data opt.mask_root = 'demo/mask/' # The location of your validation data mask val_set = MyDataLoader(opt) model = MyModel() model.initialize(opt) print('Train/Val with %d/%d' % (len(train_set), len(val_set))) for epoch in range(1, 1000): print('Epoch: %d' % epoch) epoch_iter = 0 losses_G, ssim, psnr, mae = [], [], [], [] for i, data in enumerate(train_set): epoch_iter += opt.batchSize model.set_input(data) I_g, I_o, loss_G = model.optimize_parameters() s, p, m = metrics(I_g, I_o) ssim.append(s) psnr.append(p) mae.append(m) losses_G.append(loss_G.detach().item()) print('Tra (%d/%d) G:%5.4f, S:%4.4f, P:%4.2f, M:%4.4f' % (epoch_iter, len(train_set), np.mean(losses_G), np.mean(ssim), np.mean(psnr), np.mean(mae)), end='\r') if epoch_iter == len(train_set): val_ssim, val_psnr, val_mae, val_losses_G = [], [], [], [] with torch.no_grad(): for i, data in enumerate(val_set): fname = data['fname'][0] model.set_input(data) I_g, I_o, val_loss_G = model.optimize_parameters(val=True) val_s, val_p, val_m = metrics(I_g, I_o) val_ssim.append(val_s) val_psnr.append(val_p) val_mae.append(val_m) val_losses_G.append(val_loss_G.item()) if i+1 <= 200: cv2.imwrite('./demo/output/' + fname[:-4] + '.png', postprocess(I_o).numpy()[0]) print('Val (%d/%d) G:%5.4f, S:%4.4f, P:%4.2f, M:%4.4f' % (epoch_iter, len(train_set), np.mean(val_losses_G), np.mean(val_ssim), np.mean(val_psnr), np.mean(val_mae))) losses_G, ssim, psnr, mae = [], [], [], [] model.save_networks('Model_weights')
def train(): opt.device = 'cuda:0' if not os.path.exists(opt.checkpoints_dir): os.mkdir(opt.checkpoints_dir) from dataset_tfrecord import define_dataset tfrecord_path = "/content/generator_layers_v2.1_categories.record" batch_size = opt.batchSize trainset, trainset_length = define_dataset(tfrecord_path, batch_size, train=True) valset, valset_length = define_dataset(tfrecord_path, batch_size, train=False) model = MyModel() model.initialize(opt) dpp = Preprocess() # data pre-process (dpp) print('Train/Val with %d/%d' % (trainset_length, valset_length)) for epoch in range(1, 1000): print('Epoch: %d' % epoch) train_iterator = iter(trainset) num_iterations = int(trainset_length / batch_size) epoch_iter = 0 losses_G, ssim, psnr, mae = [], [], [], [] for i in range(num_iterations): epoch_iter += opt.batchSize data, model_inputs = next(train_iterator) inpaint_region = data["inpaint_region"] person_cloth = data["person_cloth"] # warped_cloth_input = model_inputs["warped_cloth"] # Not using masked cloth. (person_cloth*inpaint_region) data = dpp(person_cloth, inpaint_region) try: model.set_input(data) I_g, I_o, loss_G = model.optimize_parameters() s, p, m = metrics(I_g, I_o) ssim.append(s) psnr.append(p) mae.append(m) losses_G.append(loss_G.detach().item()) if i % 100 == 0: print('Tra (%d/%d) G:%5.4f, S:%4.4f, P:%4.2f, M:%4.4f' % (epoch_iter, trainset_length, np.mean(losses_G), np.mean(ssim), np.mean(psnr), np.mean(mae))) #, end='\r') if epoch_iter == trainset_length: # val_ssim, val_psnr, val_mae, val_losses_G = [], [], [], [] # with torch.no_grad(): # for i, data in enumerate(val_set): # fname = data['fname'][0] # model.set_input(data) # I_g, I_o, val_loss_G = model.optimize_parameters(val=True) # val_s, val_p, val_m = metrics(I_g, I_o) # val_ssim.append(val_s) # val_psnr.append(val_p) # val_mae.append(val_m) # val_losses_G.append(val_loss_G.item()) # if i+1 <= 200: # cv2.imwrite('./demo/output/' + fname[:-4] + '.png', postprocess(I_o).numpy()[0]) # print('Val (%d/%d) G:%5.4f, S:%4.4f, P:%4.2f, M:%4.4f' % # (epoch_iter, len(train_set), np.mean(val_losses_G), np.mean(val_ssim), np.mean(val_psnr), np.mean(val_mae))) losses_G, ssim, psnr, mae = [], [], [], [] except: print("Error Occured") pass model.save_networks(epoch)
def test(): opt.device = 'cuda:0' result_dir = 'results' if not os.path.exists(result_dir): os.mkdir(result_dir) from dataset_tfrecord import define_dataset tfrecord_path = "/content/generator_layers_v2.1_categories.record" batch_size = 1 testset, testset_length = define_dataset(tfrecord_path, batch_size, train=False, test=True) dpp = Preprocess() # data pre-process (dpp) print('Test with %d' % (testset_length)) model = MyModel() model.initialize(opt) model.load_networks(str(38)) # For irregular mask inpainting val_ssim, val_psnr, val_mae, val_losses_G = [], [], [], [] ids = [] test_iterator = iter(testset) num_iterations = int(testset_length / batch_size) def tensor2array(xx): xx = xx.detach().cpu()[0] xx = xx.permute(1, 2, 0).numpy() xx = (xx + 1) * 0.5 xx = np.clip(xx, 0, 1) if xx.shape[2] == 1: return np.concatenate([xx, xx, xx], -1)[:, 32:32 + 192, :] if xx.shape[2] == 3: return xx[:, 32:32 + 192, :] with torch.no_grad(): for i in range(num_iterations): # if i == 1: break try: data, model_inputs = next(test_iterator) except: break inpaint_region = data["inpaint_region"] person_cloth = data["person_cloth"] cloth_no = int(data['clothno'].numpy()[0]) person_no = int(data['personno'].numpy()[0]) data = dpp(person_cloth, inpaint_region) model.set_input(data) I_g, I_o, I_i, val_loss_G, I_raw, L_o, mask = model.optimize_parameters( val=True) plt.figure(figsize=(12, 10)) plt.subplot(2, 3, 1) plt.imshow(tensor2array(I_i)) plt.title("Input", fontsize=20) plt.subplot(2, 3, 2) plt.imshow(tensor2array(mask)) plt.title("Inpaint Region (M)", fontsize=20) plt.subplot(2, 3, 3) plt.imshow(tensor2array(I_g)) plt.title("Ground Truth (GT)", fontsize=20) plt.subplot(2, 3, 4) plt.imshow(tensor2array(L_o)) plt.title("LBP Output", fontsize=20) plt.subplot(2, 3, 5) plt.imshow(tensor2array(I_raw)) plt.title("Generator Output (GO)", fontsize=20) plt.subplot(2, 3, 6) plt.imshow(tensor2array(I_o)) plt.title("GO*M+GT*(1-M)", fontsize=20) plt.savefig(f"{result_dir}/{i}_result.jpg") plt.imsave(f"{result_dir}/{i}_input.jpg", tensor2array(I_i)) plt.imsave(f"{result_dir}/{i}_mask.png", tensor2array(mask)) plt.imsave(f"{result_dir}/{i}_output.jpg", tensor2array(I_o)) val_s, val_p, val_m = metrics(I_g, I_o) val_ssim.append(val_s) val_psnr.append(val_p) val_mae.append(val_m) val_losses_G.append(val_loss_G.detach().item()) ids.append(str(cloth_no) + '_' + str(person_no)) losses = { 'ssim': val_ssim, 'val_mae': val_mae, 'psnr': val_psnr, 'loss_G': val_losses_G, 'ids': ids } import pandas as pd csv = pd.DataFrame(losses) csv.to_csv(f"{result_dir}/losses.csv") cmd = f"gsutil -m cp -r {result_dir} gs://vinit_helper/cloth_inpainting_gan/cloth_inpainting_local_binary_pattern/{opt.checkpoints_dir.split('/')[1]}" os.system(cmd)
def train(): opt.device = 'cuda:0' if not os.path.exists(opt.checkpoints_dir): os.mkdir(opt.checkpoints_dir) log_dir = 'tensorboard' writer = SummaryWriter(f"{opt.checkpoints_dir}/{log_dir}") from dataset_tfrecord import define_dataset tfrecord_path = "/content/generator_layers_v2.2_categories_lbp.record" batch_size = opt.batchSize trainset, trainset_length = define_dataset(tfrecord_path, batch_size, train=True) valset, valset_length = define_dataset(tfrecord_path, batch_size, train=False) model = MyModel() model.initialize(opt) dpp = Preprocess() # data pre-process (dpp) output_channel_idx = 0 input_channel_idx = 0 weights_b4 = copy.deepcopy( model.netG.dn11[0].weight[output_channel_idx, input_channel_idx].detach().cpu()) from torchsummary import summary # summary(model.netG, (4, 256, 256)) print('Train/Val with %d/%d' % (trainset_length, valset_length)) for epoch in range(1, 2): print('Epoch: %d' % epoch) train_iterator = iter(trainset) num_iterations = int(trainset_length / batch_size) epoch_iter = 0 losses_G, ssim, psnr, mae = [], [], [], [] for i in range(num_iterations): if i == 1: break epoch_iter += opt.batchSize data, model_inputs = next(train_iterator) inpaint_region = data["inpaint_region"] person_cloth = data["person_cloth"] data = dpp(person_cloth, inpaint_region) try: model.set_input(data) plt.imsave( 'input.png', model.I_i.permute(0, 2, 3, 1)[0].cpu().numpy() * 0.5 + 0.5) plt.imsave( 'inpaint_gt.png', model.I_g.permute(0, 2, 3, 1)[0].cpu().numpy() * 0.5 + 0.5) plt.imsave( 'inpaint_region.png', model.mask.permute(0, 2, 3, 1)[0, :, :, 0].cpu().numpy()) I_g, I_o, loss_G, loss_G_L2, loss_G_GAN, loss_style, loss_perceptual, loss_multi, loss_D_I_o, loss_D_I_g = model.optimize_parameters( ) s, p, m = metrics(I_g, I_o) ssim.append(s) psnr.append(p) mae.append(m) losses_G.append(loss_G.detach().item()) if i % 100 == 0: print( '\n\nTra (%d/%d) G:%5.4f, S:%4.4f, P:%4.2f, M:%4.4f' % (epoch_iter, trainset_length, np.mean(losses_G), np.mean(ssim), np.mean(psnr), np.mean(mae))) print('D1_loss_gen: %4.4f | D1_loss_real: %4.4f | D1_G_loss: %4.4f | recon_loss: %4.4f |\ vgg_loss: %4.4f | features_loss: %4.4f | style_loss: %4.4f' %(loss_D_I_o.item(), loss_D_I_g.item(), \ loss_G_GAN.item(), loss_G_L2.item(), loss_perceptual.item(), loss_multi.item(), loss_style.item())) except: print("Error") pass print('\n\n') print('Before the training: ') print(weights_b4) print('\n') print('After training of first batch with batch size 1: ') weights_after = model.netG.dn11[0].weight[ output_channel_idx, input_channel_idx].detach().cpu() print(weights_after) print('\n\n') print('Gradients: ') gradients = model.netG.dn11[0].weight.grad[ output_channel_idx, input_channel_idx].detach().cpu() print(gradients) print('\n\n')