def main(): args = tools.get_args(parser) config = tools.get_config(args) tools.init(config) tb_logger, logger = tools.get_logger(config) tools.check_dist_init(config, logger) checkpoint = tools.get_checkpoint(config) runner = tools.get_model(config, checkpoint) loaders = tools.get_data_loader(config) if dist.is_master(): logger.info(config) if args.mode == 'train': train(config, runner, loaders, checkpoint, tb_logger) elif args.mode == 'evaluate': evaluate(runner, loaders) elif args.mode == 'calc_flops': if dist.is_master(): flops = tools.get_model_flops(config, runner.get_model()) logger.info('flops: {}'.format(flops)) elif args.mode == 'calc_params': if dist.is_master(): params = tools.get_model_parameters(runner.get_model()) logger.info('params: {}'.format(params)) else: assert checkpoint is not None from models.dmcp.utils import sample_model sample_model(config, runner.get_model()) if dist.is_master(): logger.info('Done')
def main(argv): (opt, args) = parser.parse_args(argv) config = get_config(opt.config) print(opt) if opt.manualSeed is None: opt.manualSeed = random.randint(1, 10000) print("Random Seed: ", opt.manualSeed) random.seed(opt.manualSeed) torch.manual_seed(opt.manualSeed) if config['cuda']: torch.cuda.manual_seed_all(opt.manualSeed) torch.cuda.set_device(opt.gpu_ids) cudnn.benchmark = True transform = transforms.Compose([ transforms.Resize((512, 512)), transforms.ToTensor(), transforms.Normalize(mean=(0.5, 0.5, 0.5), std=(0.5, 0.5, 0.5)) ]) dataset = Aligned_Dataset(config['datapath'], direction='AtoB', transform=transform) dataloader = torch.utils.data.DataLoader(dataset, batch_size=1, shuffle=True, num_workers=int(2)) trainer = GAN_Trainer(config, dataloader) trainer.train() return
def main(argv): (opt, args) = parser.parse_args(argv) config = get_config(opt.config) print(opt) if opt.manualSeed is None: opt.manualSeed = random.randint(1, 10000) print("Random Seed: ", opt.manualSeed) random.seed(opt.manualSeed) torch.manual_seed(opt.manualSeed) if config['cuda']: torch.cuda.manual_seed_all(opt.manualSeed) torch.cuda.set_device(opt.gpu_ids) cudnn.benchmark = True transform = transforms.Compose([transforms.Resize((512, 512)), transforms.ToTensor(), transforms.Normalize(mean=(0.5, 0.5, 0.5), std=(0.5, 0.5, 0.5))]) dataset = Aligned_Dataset(config['datapath'], subfolder='test', direction='AtoB', transform=transform) dataloader = torch.utils.data.DataLoader(dataset, batch_size=1, shuffle=False, num_workers=int(2)) model_dir = '/media/scw4750/AIwalker/stackgan-like/checkpoints/generator_epoch_160.pkl' trainer = GAN_Trainer(config, dataloader) # load the model trainer.G.load_state_dict(torch.load(model_dir)) trainer.test() return
def main(argv): (opt, args) = parser.parse_args(argv) print(opt) config = get_config(opt.config) if opt.manualSeed is None: opt.manualSeed = random.randint(1, 10000) print('Random Seed: ', opt.manualSeed) random.seed(opt.manualSeed) torch.manual_seed(opt.manualSeed) if opt.cuda: torch.cuda.manual_seed_all(opt.manualSeed) torch.cuda.set_device(opt.gpu_ids) cudnn.benchmark = True # loading data set transform = transforms.Compose([transforms.Resize((config['fineSizeH'], config['fineSizeW'])), transforms.ToTensor()]) dataset = Aligned_Dataset(config['dataPath'], direction='AtoB', transform=transform) train_loader = torch.utils.data.DataLoader(dataset, batch_size=config['batchSize'], shuffle=True, num_workers=int(4)) # setup model trainer = trainer_gan(config, train_loader, resume_epoch=opt.resume_epoch) if opt.cuda: trainer.cuda() if opt.resume_epoch: trainer.resume() # training for epoch in range(opt.resume_epoch, config['nepoch']): trainer.train(epoch) trainer.update_learning_rate(epoch) if epoch % 10 == 0: trainer.save(epoch)
def main(argv): (opt, args) = parser.parse_args(argv) print(opt) config = get_config(opt.config) if opt.manualSeed is None: opt.manualSeed = random.randint(1, 10000) print('Random Seed: ', opt.manualSeed) random.seed(opt.manualSeed) torch.manual_seed(opt.manualSeed) if opt.cuda: torch.cuda.manual_seed_all(opt.manualSeed) torch.cuda.set_device(opt.gpu_ids) cudnn.benchmark = True # loading data set transform = transforms.Compose([transforms.Resize((config['fineSizeH'], config['fineSizeW'])), transforms.ToTensor()]) dataset = Aligned_Dataset(config['dataPath'], subfolder='test', direction='AtoB', transform=transform) test_loader = torch.utils.data.DataLoader(dataset, batch_size=1, shuffle=False, num_workers=int(4)) # setup model trainer = trainer_gan(config, test_loader, resume_epoch=opt.resume_epoch) # load a model trainer.netG.load_state_dict(torch.load(opt.modeldir)) if opt.cuda: trainer.cuda() # testing trainer.test()
def main(argv): (opt, args) = parser.parse_args(argv) config = get_config(opt.config) print(opt) if opt.manualSeed is None: opt.manualSeed = random.randint(1, 10000) print("Random Seed: ", opt.manualSeed) random.seed(opt.manualSeed) torch.manual_seed(opt.manualSeed) if config['cuda']: torch.cuda.manual_seed_all(opt.manualSeed) torch.cuda.set_device(opt.gpu_ids) cudnn.benchmark = True transform = transforms.Compose([ transforms.Resize((512, 512)), transforms.ToTensor(), transforms.Normalize(mean=(0.5, 0.5, 0.5), std=(0.5, 0.5, 0.5)) ]) dataset = Aligned_Dataset(config['datapath'], subfolder='test', direction='AtoB', transform=transform) dataloader = torch.utils.data.DataLoader(dataset, batch_size=1, shuffle=False, num_workers=int(2)) model_dir = '/media/scw4750/AIwalker/stackgan-like/checkpoints/generator_epoch_160.pkl' trainer = GAN_Trainer(config, dataloader) # load the model trainer.G.load_state_dict(torch.load(model_dir)) trainer.test() return
def main(): args = parser.parse_args() config = get_config(args.config) # CUDA configuration cuda = config['cuda'] device_ids = config['gpu_ids'] if cuda: # os.environ['CUDA_VISIBLE_DEVICES'] = ','.join(str(i) for i in device_ids) device_ids = list(range(len(device_ids))) config['gpu_ids'] = device_ids cudnn.benchmark = True # datetime object containing current date and time # now = datetime.now() dt_string = now.strftime("%d-%m-%Y_%H.%M.%S") # Configure checkpoint path checkpoint_path = os.path.join( 'ckpts', config['dataset_name'] + '_' + dt_string, config['mask_type'] + '_' + config['expname'] ) if not os.path.exists(checkpoint_path): os.makedirs(checkpoint_path, exist_ok=True) shutil.copy(args.config, os.path.join(checkpoint_path, os.path.basename(args.config))) writer = SummaryWriter(logdir=checkpoint_path) logger = get_logger(checkpoint_path) # get logger and configure it at the first call logger.info("Arguments: {}".format(args)) # Set random seed if args.seed is None: args.seed = random.randint(1, 10000) logger.info("Random seed: {}".format(args.seed)) random.seed(args.seed) torch.manual_seed(args.seed) if cuda: torch.cuda.manual_seed_all(args.seed) # Log the configuration logger.info("Configuration: {}".format(config)) if args.distributed: print("Distributed training...") # train_distributed(config, logger, writer, checkpoint_path) train_distributed_v2(config, logger, writer, checkpoint_path) else: train(config, logger, checkpoint_path)
def loadGenerator(args): config = get_config(args.g_config) # CUDA configuration cuda = config['cuda'] device_ids = config['gpu_ids'] if cuda: os.environ['CUDA_VISIBLE_DEVICES'] = ','.join( str(i) for i in device_ids) device_ids = list(range(len(device_ids))) config['gpu_ids'] = device_ids cudnn.benchmark = True # Set checkpoint path if not args.checkpoint_path: checkpoint_path = os.path.join( 'checkpoints', config['dataset_name'], config['mask_type'] + '_' + config['expname']) else: checkpoint_path = args.checkpoint_path # Define the trainer netG = Generator(config['netG'], cuda, device_ids).cuda() # Resume weight last_model_name = get_model_list(checkpoint_path, "gen", iteration=args.iter) model_iteration = int(last_model_name[-11:-3]) netG.load_state_dict(torch.load(last_model_name)) print("Configuration: {}".format(config)) print("Resume from {} at iteration {}".format(checkpoint_path, model_iteration)) if cuda: netG = nn.parallel.DataParallel(netG, device_ids=device_ids) return netG
def main(): config = get_config(args.config) if config['cuda']: device = torch.device("cuda:{}".format(config['gpu_ids'][0])) else: device = torch.device("cpu") trainer = Trainer(config) trainer.load_state_dict(load_weights(args.model_path, device), strict=False) trainer.eval() image = imageio.imread(args.image) image = torch.FloatTensor(image).permute(2, 0, 1).unsqueeze(0).cuda() mask = imageio.imread(args.mask) mask = (torch.FloatTensor(mask[:, :, 0]) / 255).unsqueeze(0).unsqueeze(0).cuda() x = (image / 127.5 - 1) * (1 - mask).cuda() with torch.no_grad(): _, result, _ = trainer.netG(x, mask) imageio.imwrite(args.output, upcast(result[0].permute(1, 2, 0).detach().cpu().numpy()))
def main(argv): (opt, args) = parser.parse_args(argv) print(opt) config = get_config(opt.config) if opt.manualSeed is None: opt.manualSeed = random.randint(1, 10000) print('Random Seed: ', opt.manualSeed) random.seed(opt.manualSeed) torch.manual_seed(opt.manualSeed) if opt.cuda: torch.cuda.manual_seed_all(opt.manualSeed) torch.cuda.set_device(opt.gpu_ids) cudnn.benchmark = True # loading data set transform = transforms.Compose([ transforms.Resize((config['fineSize'], config['fineSize'])), transforms.ToTensor(), transforms.Normalize(mean=(0.5, 0.5, 0.5), std=(0.5, 0.5, 0.5)) ]) dataset = Aligned_Dataset(config['dataPath'], subfolder='test', direction='AtoB', transform=transform) test_loader = torch.utils.data.DataLoader(dataset, batch_size=1, shuffle=True, num_workers=int(4)) # setup model trainer = trainer_gan(config, test_loader) # load a model trainer.netG.load_state_dict(torch.load(opt.modeldir)) if opt.cuda: trainer.cuda() # testing trainer.test()
def generateInpaintedImage(args, netG, imagePath): config = get_config(args.g_config) occlusions = [] # CUDA configuration cuda = config['cuda'] device_ids = config['gpu_ids'] if cuda: os.environ['CUDA_VISIBLE_DEVICES'] = ','.join( str(i) for i in device_ids) device_ids = list(range(len(device_ids))) config['gpu_ids'] = device_ids cudnn.benchmark = True print("Arguments: {}".format(args)) # Set random seed if args.seed is None: args.seed = random.randint(1, 10000) print("Random seed: {}".format(args.seed)) random.seed(args.seed) torch.manual_seed(args.seed) if cuda: torch.cuda.manual_seed_all(args.seed) try: # for unexpected error logging with torch.no_grad(): # enter no grad context if is_image_file(imagePath): if args.mask and is_image_file(args.mask): # Test a multiple masked image with a given mask x = default_loader(imagePath) x = transforms.Resize([512, 1024])(x) mask = default_loader(args.mask) mask = transforms.Resize(config['image_shape'][:-1])(mask) mask = transforms.CenterCrop( config['image_shape'][:-1])(mask) mask = transforms.ToTensor()(mask)[0].unsqueeze(dim=0) mask = mask.unsqueeze(dim=0) w, h = x.size first = x.crop((0, 0, w // 3, h)) second = x.crop((w // 3, 0, ((w // 3) * 2) + 2, h)) third = x.crop(((w // 3) * 2, 0, w, h)) for y in [first, second, third]: y = transforms.CenterCrop( config['image_shape'][:-1])(y) y = transforms.ToTensor()(y) y = normalize(y) y = y * (1. - mask) occlusions.append(y) elif args.mask: raise TypeError("{} is not an image file.".format( args.mask)) default_image = default_loader(imagePath) di_w, di_h = default_image.size for idx, occlusion in enumerate(occlusions): if cuda: occlusion = occlusion.cuda() mask = mask.cuda() # Inference x1, x2, offset_flow = netG(occlusion, mask) inpainted_result = x2 * mask + occlusion * (1. - mask) inp_hw = config['image_shape'][1] if idx == 0: offset = ((di_w // 3 - inp_hw) // 2, (di_h - inp_hw) // 2) elif idx == 1: offset = ((di_w - inp_hw) // 2, (di_h - inp_hw) // 2) elif idx == 2: offset = ((((di_w - inp_hw) // 2) + (di_w // 3)), (di_h - inp_hw) // 2) grid = vutils.make_grid(inpainted_result, normalize=True) # Add 0.5 after unnormalizing to [0, 255] to round to nearest integer ndarr = grid.mul_(255).add_(0.5).clamp_(0, 255).permute( 1, 2, 0).to('cpu', torch.uint8).numpy() im = Image.fromarray(ndarr) im = transforms.CenterCrop(config['mask_shape'])(im) im = transforms.Resize(config['image_shape'][:-1])(im) default_image.paste(im, offset) return default_image else: raise TypeError("{} is not an image file.".format) # exit no grad context except Exception as e: # for unexpected error logging print("Error: {}".format(e)) raise e
def main(): # if you want to see output images with colored segmentation? VIEW_COLORED_SEGMENTATION = False # Config file reading args = parser.parse_args() config = get_config(args.config) # ------ CUDA configuration cuda = config['cuda'] device_ids = config['gpu_ids'] if cuda: os.environ['CUDA_VISIBLE_DEVICES'] = ','.join( str(i) for i in device_ids) device_ids = list(range(len(device_ids))) config['gpu_ids'] = device_ids cudnn.benchmark = True if not os.path.exists(config['output_test_dir']): os.makedirs(config['output_test_dir']) # Set random seed if args.seed is None: args.seed = random.randint(1, 10000) random.seed(args.seed) torch.manual_seed(args.seed) if cuda: torch.cuda.manual_seed_all(args.seed) # first, read images and pick labels with same name # we will train all images from HQ dataset # ---------- train and test dataset&loader try: # for unexpected error logging # Load the dataset print("Inference on dataset: {" + config['dataset_name'] + "}") test_dataset = Test_Dataset( data_path=config['test_data_path'], with_subfolder=config['data_with_subfolder'], image_shape=config['image_shape'], random_crop=config['random_crop'], return_name=True) test_loader = torch.utils.data.DataLoader( dataset=test_dataset, batch_size=config['batch_size'], shuffle=False, num_workers=config['num_workers']) # [Trainer] (in test, not use trainer class directly) netG = Parser(config, cuda, device_ids) # Get the resume iteration to restart training #================== <LOAD CHECKPOINT FILE starting with parser*.pt> ============================ last_checkpoint_file = get_model_list(config['resume'], "parser", config['resume_iter']) netG.load_state_dict(torch.load(last_checkpoint_file)) print("Resume from {}".format(config['resume'])) # CUDA AVAILABLE if cuda: netG = nn.parallel.DataParallel(netG, device_ids=device_ids) # connect loaders to iter() iterable_test_loader = iter(test_loader) # learing rate #lr = config['lr'] print('Inference Start.........') start_iteration = 0 # =============== TEST =================== for iteration in range(start_iteration, config['niter'] + 1): print('ITERATION {}..... [{}/{}]'.format( iteration, iteration * config['batch_size'], int(len(test_dataset.samples) / config['batch_size']))) try: test_img_names, test_orig_images = iterable_test_loader.next() except StopIteration: iterable_test_loader = iter(test_loader) test_img_names, test_orig_images = iterable_test_loader.next() if cuda: test_orig_images = test_orig_images.cuda() # <predict test set> test_predict = netG(test_orig_images) for test_idx in range(test_orig_images.shape[0]): pred_out = torch.argmax(test_predict[test_idx], dim=0) test_sam = pred_out.cpu().numpy() if VIEW_COLORED_SEGMENTATION: decoded = decode_segmap(test_sam) misc.imsave( config['output_test_dir'] + test_img_names[test_idx] + '.png', decoded) else: cv2.imwrite( config['output_test_dir'] + test_img_names[test_idx] + '.png', test_sam) except Exception as e: # for unexpected error logging (set print("{e}") raise e
def main(): args = parser.parse_args() config = get_config(args.config) # CUDA configuration cuda = config['cuda'] device_ids = config['gpu_ids'] if cuda: os.environ['CUDA_VISIBLE_DEVICES'] = ','.join( str(i) for i in device_ids) device_ids = list(range(len(device_ids))) config['gpu_ids'] = device_ids cudnn.benchmark = True print("Arguments: {}".format(args)) # Set random seed if args.seed is None: args.seed = random.randint(1, 10000) print("Random seed: {}".format(args.seed)) random.seed(args.seed) torch.manual_seed(args.seed) if cuda: torch.cuda.manual_seed_all(args.seed) print("Configuration: {}".format(config)) try: # for unexpected error logging with torch.no_grad(): # enter no grad context if is_image_file(args.image): if args.mask and is_image_file(args.mask): # Test a single masked image with a given mask x = default_loader(args.image) mask = default_loader(args.mask) x = transforms.Resize(config['image_shape'][:-1])(x) x = transforms.CenterCrop(config['image_shape'][:-1])(x) mask = transforms.Resize(config['image_shape'][:-1])(mask) mask = transforms.CenterCrop( config['image_shape'][:-1])(mask) x = transforms.ToTensor()(x) mask = transforms.ToTensor()(mask)[0].unsqueeze(dim=0) x = normalize(x) x = x * (1. - mask) x = x.unsqueeze(dim=0) mask = mask.unsqueeze(dim=0) elif args.mask: raise TypeError( "{} is not an image file.".format(args.mask)) else: # Test a single ground-truth image with a random mask ground_truth = default_loader(args.image) ground_truth = transforms.Resize( config['image_shape'][:-1])(ground_truth) ground_truth = transforms.CenterCrop( config['image_shape'][:-1])(ground_truth) ground_truth = transforms.ToTensor()(ground_truth) ground_truth = normalize(ground_truth) ground_truth = ground_truth.unsqueeze(dim=0) bboxes = random_bbox( config, batch_size=ground_truth.size(0)) x, mask = mask_image(ground_truth, bboxes, config) # Set checkpoint path if not args.checkpoint_path: checkpoint_path = os.path.join('checkpoints', config['dataset_name'], config['mask_type'] + '_' + config['expname']) else: checkpoint_path = args.checkpoint_path # Define the trainer netG = Generator(config['netG'], cuda, device_ids) # Resume weight last_model_name = get_model_list( checkpoint_path, "gen", iteration=args.iter) netG.load_state_dict(torch.load(last_model_name)) model_iteration = int(last_model_name[-11:-3]) print("Resume from {} at iteration {}".format( checkpoint_path, model_iteration)) if cuda: netG = nn.parallel.DataParallel( netG, device_ids=device_ids) x = x.cuda() mask = mask.cuda() # Inference x1, x2, offset_flow = netG(x, mask) inpainted_result = x2 * mask + x * (1. - mask) vutils.save_image(inpainted_result, args.output, padding=0, normalize=True) print("Saved the inpainted result to {}".format(args.output)) if args.flow: vutils.save_image(offset_flow, args.flow, padding=0, normalize=True) print("Saved offset flow to {}".format(args.flow)) else: raise TypeError("{} is not an image file.".format) # exit no grad context except Exception as e: # for unexpected error logging print("Error: {}".format(e)) raise e
def main(): args = parser.parse_args() config = get_config(args.config) # CUDA configuration cuda = config['cuda'] device_ids = config['gpu_ids'] if cuda: os.environ['CUDA_VISIBLE_DEVICES'] = ','.join( str(i) for i in device_ids) device_ids = list(range(len(device_ids))) config['gpu_ids'] = device_ids cudnn.benchmark = True # Set random seed if args.seed is None: args.seed = random.randint(1, 10000) print("Random seed: {}".format(args.seed)) random.seed(args.seed) torch.manual_seed(args.seed) if cuda: torch.cuda.manual_seed_all(args.seed) chunker = ImageChunker(config['image_shape'][0], config['image_shape'][1], args.overlap) try: # for unexpected error logging with torch.no_grad(): # enter no grad context if is_image_file(args.image): print("Loading image...") imgs, masks = [], [] img_ori = default_loader(args.image) img_w, img_h = img_ori.size # Load mask txt file fname = args.image.replace('.jpg', '.txt') bboxes, _ = load_bbox_txt(fname, img_w, img_h) mask_ori = create_mask(bboxes, img_w, img_h) chunked_images = chunker.dimension_preprocess( np.array(deepcopy(img_ori))) chunked_masks = chunker.dimension_preprocess( np.array(deepcopy(mask_ori))) for (x, msk) in zip(chunked_images, chunked_masks): x = transforms.ToTensor()(x) mask = transforms.ToTensor()(msk)[0].unsqueeze(dim=0) # x = normalize(x) x = x * (1. - mask) x = x.unsqueeze(dim=0) mask = mask.unsqueeze(dim=0) imgs.append(x) masks.append(mask) # Set checkpoint path if not args.checkpoint_path: checkpoint_path = os.path.join( 'checkpoints', config['dataset_name'], config['mask_type'] + '_' + config['expname']) else: checkpoint_path = args.checkpoint_path # Define the trainer netG = Generator(config['netG'], cuda, device_ids) # Resume weight last_model_name = get_model_list(checkpoint_path, "gen", iteration=args.iter) netG.load_state_dict(torch.load(last_model_name)) model_iteration = int(last_model_name[-11:-3]) print("Resume from {} at iteration {}".format( checkpoint_path, model_iteration)) pred_imgs = [] for (x, mask) in zip(imgs, masks): if torch.max(mask) == 1: if cuda: netG = nn.parallel.DataParallel( netG, device_ids=device_ids) x = x.cuda() mask = mask.cuda() # Inference x1, x2, offset_flow = netG(x, mask) inpainted_result = x2 * mask + x * (1. - mask) inpainted_result = inpainted_result.squeeze( dim=0).permute(1, 2, 0).cpu() pred_imgs.append(inpainted_result.numpy()) else: pred_imgs.append( x.squeeze(dim=0).permute(1, 2, 0).numpy()) pred_imgs = np.asarray(pred_imgs, dtype=np.float32) reconstructed_image = chunker.dimension_postprocess( pred_imgs, np.array(img_ori)) # plt.imshow(reconstructed_image); plt.show() reconstructed_image = torch.tensor( reconstructed_image).permute(2, 0, 1).unsqueeze(dim=0) vutils.save_image(reconstructed_image, args.output, padding=0, normalize=True) print("Saved the inpainted result to {}".format(args.output)) if args.flow: vutils.save_image(offset_flow, args.flow, padding=0, normalize=True) print("Saved offset flow to {}".format(args.flow)) else: raise TypeError("{} is not an image file.".format) # exit no grad context except Exception as e: # for unexpected error logging print("Error: {}".format(e)) raise e
def main(): # Config file reading args = parser.parse_args() config = get_config(args.config) # ------ CUDA configuration cuda = config['cuda'] device_ids = config['gpu_ids'] if cuda: os.environ['CUDA_VISIBLE_DEVICES'] = ','.join(str(i) for i in device_ids) device_ids = list(range(len(device_ids))) config['gpu_ids'] = device_ids cudnn.benchmark = True # ----- Directory to save checkpoint file checkpoint_path = os.path.join('checkpoints_', config['dataset_name'], 'parser' + '_' + config['expname']) if not os.path.exists(checkpoint_path): os.makedirs(checkpoint_path) shutil.copy(args.config, os.path.join(checkpoint_path, os.path.basename(args.config))) writer = SummaryWriter(log_dir=checkpoint_path) logger = get_logger(checkpoint_path) # get logger and configure it at the first call if not os.path.exists(config['output_test_dir']): os.makedirs(config['output_test_dir']) logger.info(f"Arguments: {args}") # Set random seed if args.seed is None: args.seed = random.randint(1, 10000) logger.info(f"Random seed: {args.seed}") random.seed(args.seed) torch.manual_seed(args.seed) if cuda: torch.cuda.manual_seed_all(args.seed) logger.info(f"configuration : {config}") # first, read images and pick labels with same name # we will train all images from HQ dataset # ---------- train and test dataset&loader try: # for unexpected error logging # Load the dataset logger.info(f"Training on dataset: {config['dataset_name']}") train_dataset = Parse_Dataset(data_path=config['all_data_path'], gt_path=config['gt_data_path'], with_subfolder=config['data_with_subfolder'], image_shape=config['image_shape'], random_crop=config['random_crop'], return_name=True) train_loader = torch.utils.data.DataLoader(dataset=train_dataset, batch_size=config['batch_size'], shuffle=True, num_workers=config['num_workers']) test_dataset = Test_Dataset(data_path=config['test_data_path'], with_subfolder=config['data_with_subfolder'], image_shape=config['image_shape'], random_crop=config['random_crop'], return_name=True) test_loader = torch.utils.data.DataLoader(dataset=test_dataset, batch_size=config['batch_size'], shuffle=False, num_workers=config['num_workers']) # [Trainer] trainer = Seg_Trainer(config) logger.info(f"\n{trainer.netParser}") # CUDA AVAILABLE if cuda: trainer = nn.parallel.DataParallel(trainer, device_ids=device_ids) trainer_module = trainer.module else: trainer_module = trainer # Get the resume iteration to restart training #start_iteration = trainer_module.resume(config['resume']) if config['resume'] else 1 start_iteration = trainer_module.resume(config['resume'], config['resume_iter']) if config['resume'] else 1 # connect loaders to iter() iterable_train_loader = iter(train_loader) iterable_test_loader = iter(test_loader) # learing rate lr = config['lr'] print('Training Start.........') for iteration in range(start_iteration, config['niter'] + 1): #=============== TRAIN =================== # ------ [ train batch loader ] --------- try: train_img_names, gt_images, gt_targets, orig_images = iterable_train_loader.next() except StopIteration: iterable_train_loader = iter(train_loader) train_img_names, gt_images, gt_targets, orig_images = iterable_train_loader.next() # ------ [ train batch ] --------- if cuda: orig_images = orig_images.cuda() gt_images = gt_images.cuda() # Forward loss, predict = trainer(orig_images, gt_images) if not loss.dim() == 0: loss = torch.mean(loss) # Backward (update optimizer) trainer_module.optimizerSGD.zero_grad() loss.backward() trainer_module.optimizerSGD.step() # [print loss] (in this, 1 print for 30 iteration) if iteration % 30 == 0: print("Epoch [%d/%d] Loss: %.10f lr:%.6f" % (iteration ,config['niter'], loss.data, lr)) #=============== TEST =================== if iteration % 50 == 0: try: test_img_names, test_orig_images = iterable_test_loader.next() except StopIteration: iterable_test_loader = iter(test_loader) test_img_names, test_orig_images = iterable_test_loader.next() if cuda: test_orig_images = test_orig_images.cuda() # <predict test set> test_predict = trainer.module.netParser(test_orig_images) for test_idx in range(test_orig_images.shape[0]): pred_out = torch.argmax(test_predict[test_idx], dim=0) test_sam = pred_out.cpu().numpy() cv2.imwrite(config['output_test_dir'] + test_img_names[test_idx] + '.png', test_sam) # <learning rate up> if iteration % 50000 == 0: lr = lr * config['lr_decay'] for param_group in trainer_module.optimizerSGD.param_groups: param_group['lr'] = lr # save the model if iteration % config['snapshot_save_iter'] == 0: trainer_module.save_model(checkpoint_path, iteration) except Exception as e: # for unexpected error logging logger.error(f"{e}") raise e
def main(argv): (opt, args) = parser.parse_args(argv) print(opt) if opt.manualSeed is None: opt.manualSeed = random.randint(1, 10000) print("random seed: ", opt.manualSeed) random.seed(opt.manualSeed) torch.manual_seed(opt.manualSeed) if opt.cuda: torch.cuda.manual_seed_all(opt.manualSeed) torch.cuda.set_device(opt.gpu_ids) cudnn.benchmark = True config = get_config(opt.config) # loading data sets datasetA = DATASET(os.path.join(config['dataPath'], 'trainA'), config['loadSize'], config['fineSize'], config['flip']) datasetB = DATASET(os.path.join(config['dataPath'], 'trainB'), config['loadSize'], config['fineSize'], config['flip']) loader_A = torch.utils.data.DataLoader(dataset=datasetA, batch_size=config['batchSize'], shuffle=True, num_workers=config['num_workers']) loaderA = iter(loader_A) loader_B = torch.utils.data.DataLoader(dataset=datasetB, batch_size=config['batchSize'], shuffle=True, num_workers=config['num_workers']) loaderB = iter(loader_B) # define the trainer trainer = ImGANTrainer(config) print(trainer.netG_ab, trainer.netD_ab) if opt.cuda: trainer.cuda() for iteration in range(1, config['niter'] + 1): try: imgA = loaderA.next() imgB = loaderB.next() except StopIteration: loaderA, loaderB = iter(loader_A), iter(loader_B) imgA = loaderA.next() imgB = loaderB.next() image_A = Variable(imgA.cuda()) image_B = Variable(imgB.cuda()) trainer.dis_upodate(image_A, image_B) trainer.gen_update(image_A, image_B) if iteration % 100 == 0: trainer.get_current_losses() losses = trainer.get_current_losses() message = '([%d/%d][%d/%d]) ' % (iteration, config['niter'], len(loader_A), len(loader_B)) for k, v in losses.items(): message += '%s: %.6f ' % (k, v) print(message) input_a, input_b, fake_ab, fake_ba, fake_bab, fake_aba = trainer.test( image_A, image_B) vutils.save_image(image_A.data, '%s/realA_niter_%03d_1.png' % (config['outf'], iteration), normalize=True) vutils.save_image(image_B.data, '%s/realB_niter_%03d_1.png' % (config['outf'], iteration), normalize=True) vutils.save_image(fake_aba.data, '%s/recA_niter_%03d_1.png' % (config['outf'], iteration), normalize=True) vutils.save_image(fake_bab.data, '%s/recB_niter_%03d_1.png' % (config['outf'], iteration), normalize=True) vutils.save_image(fake_ab.data, '%s/AB_niter_%03d_1.png' % (config['outf'], iteration), normalize=True) vutils.save_image(fake_ba.data, '%s/BA_niter_%03d_1.png' % (config['outf'], iteration), normalize=True)
parser.add_argument('--seed', type=int, default=0, help='manual seed') # parser.add_argument('--image', type=str) # parser.add_argument('--mask', type=str, default='') # parser.add_argument('--output', type=str, default='output.png') # parser.add_argument('--flow', type=str, default='') parser.add_argument('--checkpoint_path', type=str, default='checkpoints2/imagenet/hole_benchmark') parser.add_argument('--iter', type=int, default=0) # default means the latest iteration parser.add_argument( '--which_model', type=str, default='checkpoints2/imagenet/hole_benchmark/gen_00080000.pt') args = parser.parse_args() config = get_config(args.config) # CUDA configuration cuda = config['cuda'] if torch.cuda.device_count() > 0: cuda = True # memory problem else: cuda = False if not args.checkpoint_path: args.checkpoint_path = os.path.join( 'checkpoints', config['dataset_name'], config['mask_type'] + '_' + config['expname']) print("args.checkpoint_path = {}".format(args.checkpoint_path)) dataset_name = args.checkpoint_path.split("/")[1] device_ids = config['gpu_ids']
def main(): args = parser.parse_args() config = get_config(args.config) # CUDA configuration cuda = config['cuda'] device_ids = config['gpu_ids'] if cuda: os.environ['CUDA_VISIBLE_DEVICES'] = ','.join( str(i) for i in device_ids) device_ids = list(range(len(device_ids))) config['gpu_ids'] = device_ids cudnn.benchmark = True # Configure checkpoint path checkpoint_path = os.path.join( 'checkpoints', config['dataset_name'], config['mask_type'] + '_' + config['expname']) if not os.path.exists(checkpoint_path): os.makedirs(checkpoint_path) shutil.copy(args.config, os.path.join(checkpoint_path, os.path.basename(args.config))) writer = SummaryWriter(log_dir=checkpoint_path) logger = get_logger( checkpoint_path) # get logger and configure it at the first call logger.info(f"Arguments: {args}") # Set random seed if args.seed is None: args.seed = random.randint(1, 10000) logger.info(f"Random seed: {args.seed}") random.seed(args.seed) torch.manual_seed(args.seed) if cuda: torch.cuda.manual_seed_all(args.seed) # Log the configuration logger.info(f"Configuration: {config}") try: # for unexpected error logging # Load the dataset logger.info(f"Training on dataset: {config['dataset_name']}") train_dataset = Dataset(data_path=config['train_data_path'], with_subfolder=config['data_with_subfolder'], image_shape=config['image_shape'], random_crop=config['random_crop']) train_loader = torch.utils.data.DataLoader( dataset=train_dataset, batch_size=config['batch_size'], shuffle=True, num_workers=config['num_workers']) # Define the trainer trainer = Trainer(config) logger.info(f"\n{trainer.netG}") logger.info(f"\n{trainer.globalD}") if cuda: trainer = nn.parallel.DataParallel(trainer, device_ids=device_ids) trainer_module = trainer.module else: trainer_module = trainer # Get the resume iteration to restart training start_iteration = trainer_module.resume( config['checkpoint_dir'], args.iteration) if config['resume'] else 1 iterable_train_loader = iter(train_loader) time_count = time.time() for iteration in range(start_iteration, config['niter'] + 1): try: ground_truth = iterable_train_loader.next() except StopIteration: iterable_train_loader = iter(train_loader) ground_truth = iterable_train_loader.next() # Prepare the inputs x, mask = mask_image(ground_truth, config) if cuda: x = x.cuda() mask = mask.cuda() ground_truth = ground_truth.cuda() ###### Forward pass ###### losses, coarse_result, inpainted_result = trainer( x, mask, ground_truth) # Scalars from different devices are gathered into vectors for k in losses.keys(): if not losses[k].dim() == 0: losses[k] = torch.mean(losses[k]) ###### Backward pass ###### # Update D trainer_module.optimizer_d.zero_grad() losses['d'] = losses['d_loss_rel'] + losses['d_loss_loren'] losses['d'].backward(retain_graph=True) trainer_module.optimizer_d.step() # Update G trainer_module.optimizer_g.zero_grad() losses['g'] = losses['l1'] * config['l1_loss_alpha'] + losses[ 'g_loss_rel'] + losses['g_loss_loren'] losses['g'].backward() trainer_module.optimizer_g.step() # Log and visualization log_losses = ['l1', 'g', 'd'] if iteration % config['print_iter'] == 0: time_count = time.time() - time_count speed = config['print_iter'] / time_count speed_msg = 'speed: %.2f batches/s ' % speed time_count = time.time() message = 'Iter: [%d/%d] ' % (iteration, config['niter']) for k in log_losses: v = losses[k] writer.add_scalar(k, v, iteration) message += '%s: %.6f ' % (k, v) message += speed_msg logger.info(message) def denorm(x): out = (x + 1) / 2 # [-1,1] -> [0,1] return out.clamp_(0, 1) if iteration % (config['viz_iter']) == 0: ims = torch.cat( [x, coarse_result, inpainted_result, ground_truth], dim=3) writer.add_image('raw_masked_coarse_refine', denorm(ims), iteration) # if iteration % (config['save_image']) ==0: # viz_max_out = config['viz_max_out'] # if x.size(0) > viz_max_out: # viz_images = torch.stack([x[:viz_max_out], coarse_result[:viz_max_out], inpainted_result[:viz_max_out], # ], dim=1) # else: # viz_images = torch.stack([x, coarse_result, inpainted_result], dim=1) # viz_images = viz_images.view(-1, *list(x.size())[1:]) # vutils.save_image(viz_images, # '%s/niter_%03d.png' % (checkpoint_path, iteration), # nrow=3 * 4, # normalize=True) # Save the model if iteration % config['snapshot_save_iter'] == 0: trainer_module.save_model(checkpoint_path) except Exception as e: # for unexpected error logging logger.error(f"{e}") raise e
def main(): args = parser.parse_args() config = get_config(args.config) # CUDA configuration cuda = config['cuda'] device_ids = config['gpu_ids'] if cuda: os.environ['CUDA_VISIBLE_DEVICES'] = ','.join( str(i) for i in device_ids) device_ids = list(range(len(device_ids))) config['gpu_ids'] = device_ids cudnn.benchmark = True # Configure checkpoint path checkpoint_path = os.path.join( 'checkpoints', config['dataset_name'], config['mask_type'] + '_' + config['expname']) if not os.path.exists(checkpoint_path): os.makedirs(checkpoint_path) shutil.copy(args.config, os.path.join(checkpoint_path, os.path.basename(args.config))) writer = SummaryWriter(logdir=checkpoint_path) logger = get_logger( checkpoint_path) # get logger and configure it at the first call logger.info("Arguments: {}".format(args)) # Set random seed if args.seed is None: args.seed = random.randint(1, 10000) logger.info("Random seed: {}".format(args.seed)) random.seed(args.seed) torch.manual_seed(args.seed) if cuda: torch.cuda.manual_seed_all(args.seed) # Log the configuration logger.info("Configuration: {}".format(config)) try: # for unexpected error logging # Load the dataset logger.info("Training on dataset: {}".format(config['dataset_name'])) train_dataset = Dataset(data_path=config['train_data_path'], with_subfolder=config['data_with_subfolder'], image_shape=config['image_shape'], random_crop=config['random_crop']) # val_dataset = Dataset(data_path=config['val_data_path'], # with_subfolder=config['data_with_subfolder'], # image_size=config['image_size'], # random_crop=config['random_crop']) train_loader = torch.utils.data.DataLoader( dataset=train_dataset, batch_size=config['batch_size'], shuffle=True, num_workers=config['num_workers']) # val_loader = torch.utils.data.DataLoader(dataset=val_dataset, # batch_size=config['batch_size'], # shuffle=False, # num_workers=config['num_workers']) # Define the trainer trainer = Trainer(config) logger.info("\n{}".format(trainer.netG)) logger.info("\n{}".format(trainer.localD)) logger.info("\n{}".format(trainer.globalD)) if cuda: trainer = nn.parallel.DataParallel(trainer, device_ids=device_ids) trainer_module = trainer.module else: trainer_module = trainer # Get the resume iteration to restart training start_iteration = trainer_module.resume( "checkpoints/{}/hole_benchmark".format(config['dataset_name']), iteration=config['resume']) if config['resume'] else 1 iterable_train_loader = iter(train_loader) time_count = time.time() for iteration in range(start_iteration, config['niter'] + 1): try: ground_truth = next(iterable_train_loader) except StopIteration: iterable_train_loader = iter(train_loader) ground_truth = next(iterable_train_loader) # Prepare the inputs bboxes = random_bbox(config, batch_size=ground_truth.size(0)) x, mask = mask_image(ground_truth, bboxes, config) if cuda: x = x.cuda() mask = mask.cuda() ground_truth = ground_truth.cuda() ###### Forward pass ###### compute_g_loss = iteration % config['n_critic'] == 0 losses, inpainted_result, offset_flow = trainer( x, bboxes, mask, ground_truth, compute_g_loss) # Scalars from different devices are gathered into vectors for k in losses.keys(): if not losses[k].dim() == 0: losses[k] = torch.mean(losses[k]) ###### Backward pass ###### # Update D trainer_module.optimizer_d.zero_grad() losses['d'] = losses[ 'wgan_d'] + losses['wgan_gp'] * config['wgan_gp_lambda'] losses['d'].backward() trainer_module.optimizer_d.step() # Update G if compute_g_loss: trainer_module.optimizer_g.zero_grad() losses['g'] = losses['l1'] * config['l1_loss_alpha'] \ + losses['ae'] * config['ae_loss_alpha'] \ + losses['wgan_g'] * config['gan_loss_alpha'] losses['g'].backward() trainer_module.optimizer_g.step() # Log and visualization log_losses = ['l1', 'ae', 'wgan_g', 'wgan_d', 'wgan_gp', 'g', 'd'] if iteration % config['print_iter'] == 0: time_count = time.time() - time_count speed = config['print_iter'] / time_count speed_msg = 'speed: %.2f batches/s ' % speed time_count = time.time() message = 'Iter: [%d/%d] ' % (iteration, config['niter']) for k in log_losses: v = losses.get(k, 0.) writer.add_scalar(k, v, iteration) message += '%s: %.6f ' % (k, v) message += speed_msg logger.info(message) if iteration % (config['viz_iter']) == 0: viz_max_out = config['viz_max_out'] if x.size(0) > viz_max_out: viz_images = torch.stack([ x[:viz_max_out], inpainted_result[:viz_max_out], offset_flow[:viz_max_out] ], dim=1) else: viz_images = torch.stack( [x, inpainted_result, offset_flow], dim=1) viz_images = viz_images.view(-1, *list(x.size())[1:]) vutils.save_image(viz_images, '%s/niter_%03d.png' % (checkpoint_path, iteration), nrow=3 * 4, normalize=True) # Save the model if iteration % config['snapshot_save_iter'] == 0: trainer_module.save_model(checkpoint_path, iteration) except Exception as e: # for unexpected error logging logger.error("{}".format(e)) raise e
path2best_model = os.path.join(weights_folder, f"{run_name}.pth") log_metrics(metrics, epoch, "val", logger, config["log_mlflow"], path2best_model=path2best_model) if __name__ == '__main__': from datetime import datetime from utils.tools import get_config run_name = datetime.now().strftime("%d_%m_%Y.%H_%M") config = get_config() log_mlflow = config.get("mlflow", None) is not None config["log_mlflow"] = log_mlflow config["run_name"] = run_name setup_mlflow(config) weights_folder = os.path.join(os.getcwd(), "weights", config["description"]["model_name"], config["description"]["task_type"]) log_folder = os.path.join(os.getcwd(), "logs", config["description"]["model_name"], config["description"]["task_type"], run_name) os.makedirs(weights_folder, exist_ok=True) os.makedirs(log_folder, exist_ok=True)
def main(): args = parser.parse_args() config = get_config(args.config) # CUDA configuration cuda = config['cuda'] device_ids = config['gpu_ids'] if cuda: os.environ['CUDA_VISIBLE_DEVICES'] = ','.join( str(i) for i in device_ids) device_ids = list(range(len(device_ids))) config['gpu_ids'] = device_ids cudnn.benchmark = True print("Arguments: {}".format(args)) # Set random seed if args.seed is None: args.seed = random.randint(1, 10000) print("Random seed: {}".format(args.seed)) random.seed(args.seed) torch.manual_seed(args.seed) if cuda: torch.cuda.manual_seed_all(args.seed) print("Configuration: {}".format(config)) try: # for unexpected error logging with torch.no_grad(): # enter no grad context file = dataset_files(args.test_root, "*.jpg") mask_file = dataset_files(args.mask_dir, "*.png") for j in range(len(mask_file)): for i in range(len(file)): if is_image_file(file[i]): if mask_file and is_image_file(mask_file[j]): # Test a single masked image with a given mask x = default_loader(file[i]) mask = default_loader(mask_file[j]) # x = cv2.cvtColor(cv2.imread(file[i]), cv2.COLOR_BGR2RGB) # mask = cv2.cvtColor(cv2.imread(mask_file[j]), cv2.COLOR_BGR2RGB) # x = cv2.resize(x, (config['image_shape'][0], config['image_shape'][1])) # mask = cv2.resize(mask, (config['image_shape'][0], config['image_shape'][1])) x = transforms.Resize( config['image_shape'][:-1])(x) x = transforms.CenterCrop( config['image_shape'][:-1])(x) # mask = transforms.Resize(config['image_shape'][:-1])(mask) # mask = transforms.CenterCrop(config['image_shape'][:-1])(mask) x = transforms.ToTensor()(x) mask = transforms.ToTensor()(mask)[0].unsqueeze( dim=0) x = normalize(x) x = x * (1. - mask) x = x.unsqueeze(dim=0) # x_raw = x mask = mask.unsqueeze(dim=0) elif mask_file[j]: raise TypeError("{} is not an image file.".format( mask_file[j])) else: # Test a single ground-truth image with a random mask ground_truth = default_loader(file[i]) ground_truth = transforms.Resize( config['image_shape'][:-1])(ground_truth) ground_truth = transforms.CenterCrop( config['image_shape'][:-1])(ground_truth) ground_truth = transforms.ToTensor()(ground_truth) ground_truth = normalize(ground_truth) ground_truth = ground_truth.unsqueeze(dim=0) bboxes = test_bbox(config, batch_size=ground_truth.size(0), t=50, l=50) x, mask = mask_image(ground_truth, bboxes, config) # Set checkpoint path if not args.checkpoint_path: checkpoint_path = os.path.join( 'checkpoints', config['dataset_name'], config['mask_type'] + '_' + config['expname']) else: checkpoint_path = args.checkpoint_path # Define the trainer netG = Generator(config['netG'], cuda, device_ids) # Resume weight g_checkpoint = torch.load(f'{checkpoint_path}/gen.pt') netG.load_state_dict(g_checkpoint) # model_iteration = int(last_model_name[-11:-3]) print("Model Resumed".format(checkpoint_path)) if cuda: netG = nn.parallel.DataParallel( netG, device_ids=device_ids) x = x.cuda() mask = mask.cuda() # Inference x1, x2 = netG(x, mask) inpainted_result = x2 * mask + x * (1. - mask) inpainted_result_cpu = torch.Tensor.cpu( inpainted_result).detach().permute(0, 2, 3, 1) inpainted_result_cpu = np.asarray( inpainted_result_cpu[0]) inpainted_result_cpu = cv2.normalize( inpainted_result_cpu, inpainted_result_cpu, 0, 255, cv2.NORM_MINMAX) # cat_result = torch.cat([x, inpainted_result, ground_truth], dim=3).cuda() vutils.save_image(inpainted_result, args.output_dir + 'output_{}/'.format(j + 1) + 'output_{}.png'.format(i), padding=0, normalize=True) # cv2.imwrite(args.output_dir+ 'output_{}/'.format(j+1) + 'output_{}.png'.format(i), inpainted_result_cpu) # cv2.cvtColor(inpainted_result_cpu, cv2.COLOR_BGR2RGB)) print("{}th image saved".format(i)) else: raise TypeError("{} is not an image file.".format) # exit no grad context except Exception as e: # for unexpected error logging print("Error: {}".format(e)) raise e
from utils.tools import get_config, random_bbox, mask_image, is_image_file, default_loader, normalize, get_model_list import torch import math from multiprocessing.dummy import Pool irange = range SEED = 42 random.seed(SEED) np.random.seed(SEED) torch.manual_seed(SEED) torch.cuda.manual_seed_all(SEED) config = get_config('configs/config.yaml') # CUDA configuration cuda = config['cuda'] print(' ------------- cuda -------------') print(cuda) device_ids = config['gpu_ids'] def make_grid(tensor, nrow=8, padding=2, normalize=False, range=None, scale_each=False, pad_value=0): """Make a grid of images. Args: tensor (Tensor or list): 4D mini-batch Tensor of shape (B x C x H x W) or a list of images all of the same size. nrow (int, optional): Number of images displayed in each row of the grid.
def main(): args = parser.parse_args() config = get_config(args.config) # CUDA configuration cuda = config['cuda'] device_ids = config['gpu_ids'] if cuda: os.environ['CUDA_VISIBLE_DEVICES'] = ','.join(str(i) for i in device_ids) device_ids = list(range(len(device_ids))) config['gpu_ids'] = device_ids cudnn.benchmark = True # Set random seed if args.seed is None: args.seed = random.randint(1, 10000) # print("Random seed: {}".format(args.seed)) random.seed(args.seed) torch.manual_seed(args.seed) if cuda: torch.cuda.manual_seed_all(args.seed) t0 = time.time() dataset = datasets.LoadImages(args.image) chunker = ImageChunker(config['image_shape'][0], config['image_shape'][1], args.overlap) try: # for unexpected error logging with torch.no_grad(): # enter no grad context # Set checkpoint path if not args.checkpoint_path: checkpoint_path = os.path.join('checkpoints', config['dataset_name'], config['mask_type'] + '_' + config['expname']) else: checkpoint_path = args.checkpoint_path last_model_name = get_model_list(checkpoint_path, "gen", iteration=args.iter) prev_fname = '' vid_writer = None for fpath, img_ori, vid_cap in dataset : imgs, masks = [], [] if prev_fname == fpath : frame += 1 # increase frame number if still on the same file else : frame = 0 # start frame number _, img_h, img_w = img_ori.shape txtfile = pathlib.Path(fpath).with_suffix('.txt') # Load mask txt file txtfile = os.path.join(args.output, str(txtfile).split('/')[-1]) if os.path.exists(txtfile) : bboxes, bframes = load_bbox_txt(txtfile, img_w, img_h) assert len(bboxes) == len(bframes) idx = [ii for ii, val in enumerate(bframes) if val==frame] bndbxs = [bboxes[ii] for ii in idx] img_ori = np.moveaxis(img_ori, 0, -1) if len(bndbxs) > 0 : # if any logo detected mask_ori = create_mask(bndbxs, img_w, img_h) # fig, axes = plt.subplots(1,2); axes[0].imshow(img_ori[0]); axes[1].imshow(mask_ori); plt.show() chunked_images = chunker.dimension_preprocess(np.array(deepcopy(img_ori))) chunked_masks = chunker.dimension_preprocess(np.array(deepcopy(mask_ori))) for (x, msk) in zip(chunked_images, chunked_masks) : x = transforms.ToTensor()(x) mask = transforms.ToTensor()(msk)[0].unsqueeze(dim=0) # x = normalize(x) x = x * (1. - mask) x = x.unsqueeze(dim=0) mask = mask.unsqueeze(dim=0) imgs.append(x) masks.append(mask) # Define the trainer netG = Generator(config['netG'], cuda, device_ids) netG.load_state_dict(torch.load(last_model_name)) model_iteration = int(last_model_name[-11:-3]) # print("Resume from {} at iteration {}".format(checkpoint_path, model_iteration)) pred_imgs = [] for (x, mask) in zip(imgs, masks) : if torch.max(mask) == 1 : if cuda: netG = nn.parallel.DataParallel(netG, device_ids=device_ids) x = x.cuda() mask = mask.cuda() # Inference x1, x2, offset_flow = netG(x, mask) inpainted_result = x2 * mask + x * (1. - mask) inpainted_result = inpainted_result.squeeze(dim=0).permute(1,2,0).cpu() pred_imgs.append(inpainted_result.numpy()) else : pred_imgs.append(x.squeeze(dim=0).permute(1,2,0).numpy()) pred_imgs = np.asarray(pred_imgs, dtype=np.float32) reconstructed_image = chunker.dimension_postprocess(pred_imgs, np.array(img_ori)) reconstructed_image = np.uint8(reconstructed_image[:, :, ::-1]*255) # BGR to RGB, and rescaling else : # no logo detected reconstructed_image = img_ori[:, :, ::-1] # Save results (image with detections) outname = fpath.split('/')[-1] outname = outname.split('.')[0] + '-inp.' + outname.split('.')[-1] outpath = os.path.join(args.output, outname) if dataset.mode == 'images': cv2.imwrite(outpath, reconstructed_image) print("Saved the inpainted image to {}".format(outpath)) else : if fpath != prev_fname: # new video if isinstance(vid_writer, cv2.VideoWriter): vid_writer.release() # release previous video writer print("Saved the inpainted video to {}".format(outpath)) fps = vid_cap.get(cv2.CAP_PROP_FPS) w = int(vid_cap.get(cv2.CAP_PROP_FRAME_WIDTH)) h = int(vid_cap.get(cv2.CAP_PROP_FRAME_HEIGHT)) vid_writer = cv2.VideoWriter(outpath, cv2.VideoWriter_fourcc(*args.fourcc), fps, (w, h)) vid_writer.write(reconstructed_image) prev_fname = fpath # exit no grad context except Exception as err: # for unexpected error logging print("Error: {}".format(err)) pass print('Inpainting: (%.3fs)' % (time.time() - t0))