def test(args): # device device = torch.device("cuda:%d" % args.gpu if torch.cuda.is_available() else "cpu") torch.backends.cudnn.benchmark = True # data testset = SonyTestDataset(args.input_dir, args.gt_dir) test_loader = DataLoader(testset, batch_size=args.batch_size, shuffle=False, num_workers=args.num_workers) # model model = Unet() model.load_state_dict(torch.load(args.model)) model.to(device) model.eval() # testing for i, databatch in tqdm(enumerate(test_loader), total=len(test_loader)): input_full, scale_full, gt_full, test_id, ratio = databatch scale_full, gt_full = torch.squeeze(scale_full), torch.squeeze(gt_full) # processing inputs = input_full.to(device) outputs = model(inputs) outputs = outputs.cpu().detach() outputs = torch.squeeze(outputs) outputs = outputs.permute(1, 2, 0) # scaling can clipping outputs, scale_full, gt_full = outputs.numpy(), scale_full.numpy( ), gt_full.numpy() scale_full = scale_full * np.mean(gt_full) / np.mean( scale_full ) # scale the low-light image to the same mean of the ground truth outputs = np.minimum(np.maximum(outputs, 0), 1) # saving if not os.path.isdir(os.path.join(args.result_dir, 'eval')): os.makedirs(os.path.join(args.result_dir, 'eval')) scipy.misc.toimage( scale_full * 255, high=255, low=0, cmin=0, cmax=255).save( os.path.join( args.result_dir, 'eval', '%05d_00_train_%d_scale.jpg' % (test_id[0], ratio[0]))) scipy.misc.toimage( outputs * 255, high=255, low=0, cmin=0, cmax=255).save( os.path.join( args.result_dir, 'eval', '%05d_00_train_%d_out.jpg' % (test_id[0], ratio[0]))) scipy.misc.toimage( gt_full * 255, high=255, low=0, cmin=0, cmax=255).save( os.path.join( args.result_dir, 'eval', '%05d_00_train_%d_gt.jpg' % (test_id[0], ratio[0])))
def test(args): # device device = torch.device("cuda:%d" % args.gpu if torch.cuda.is_available() else "cpu") torch.backends.cudnn.benchmark = True # images path fns = glob.glob(path.join(args.imgdir, '*.DNG')) n = len(fns) # model model = Unet() model.load_state_dict(torch.load(args.model)) model.to(device) model.eval() # ratio ratio = 200 for idx in range(n): fn = fns[idx] print(fn) raw = rawpy.imread(fn) input = np.expand_dims(pack_raw(raw), axis=0) * ratio scale_full = np.expand_dims(np.float32(input / 65535.0), axis=0) input = crop_center(input, 1024, 1024) input = torch.from_numpy(input) input = torch.squeeze(input) input = input.permute(2, 0, 1) input = torch.unsqueeze(input, dim=0) input = input.to(device) outputs = model(input) outputs = outputs.cpu().detach() outputs = torch.squeeze(outputs) outputs = outputs.permute(1, 2, 0) outputs = outputs.numpy() outputs = np.minimum(np.maximum(outputs, 0), 1) scale_full = torch.from_numpy(scale_full) scale_full = torch.squeeze(scale_full) scipy.misc.toimage(outputs * 255, high=255, low=0, cmin=0, cmax=255).save( path.join(args.imgdir, path.basename(fn) + '_out.jpg'))
def main(args): logging.basicConfig(level=logging.DEBUG if args.verbose else logging.INFO) device = torch.device( 'cuda' if args.use_cuda and torch.cuda.is_available() else 'cpu') # Create output folder if (args.output_folder is not None): if not os.path.exists(args.output_folder): os.makedirs(args.output_folder) logging.debug('Creating folder `{0}`'.format(args.output_folder)) output_folder = os.path.join(args.output_folder, time.strftime('%Y-%m-%d_%H%M%S')) os.makedirs(output_folder) logging.debug('Creating folder `{0}`'.format(output_folder)) args.datafolder = os.path.abspath(args.datafolder) args.model_path = os.path.abspath( os.path.join(output_folder, 'model.th')) # Save the configuration in a config.json file with open(os.path.join(output_folder, 'config.json'), 'w') as f: json.dump(vars(args), f, indent=2) logging.info('Saving configuration file in `{0}`'.format( os.path.abspath(os.path.join(output_folder, 'config.json')))) # Get datasets and load into meta learning format meta_train_dataset, meta_val_dataset, _ = get_datasets( args.dataset, args.datafolder, args.num_ways, args.num_shots, args.num_shots_test, augment=augment, fold=args.fold, download=download_data) meta_train_dataloader = BatchMetaDataLoader(meta_train_dataset, batch_size=args.batch_size, shuffle=True, num_workers=args.num_workers, pin_memory=True) meta_val_dataloader = BatchMetaDataLoader(meta_val_dataset, batch_size=args.batch_size, shuffle=True, num_workers=args.num_workers, pin_memory=True) # Define model model = Unet(device=device, feature_scale=args.feature_scale) model = model.to(device) print(f'Using device: {device}') # Define optimizer meta_optimizer = torch.optim.Adam(model.parameters(), lr=args.meta_lr) #, weight_decay=1e-5) #meta_optimizer = torch.optim.RMSprop(model.parameters(), lr=learning_rate, momentum = 0.99) # Define meta learner metalearner = ModelAgnosticMetaLearning( model, meta_optimizer, first_order=args.first_order, num_adaptation_steps=args.num_adaption_steps, step_size=args.step_size, learn_step_size=False, loss_function=loss_function, device=device) best_value = None # Training loop epoch_desc = 'Epoch {{0: <{0}d}}'.format(1 + int(math.log10(args.num_epochs))) train_losses = [] val_losses = [] train_ious = [] train_accuracies = [] val_accuracies = [] val_ious = [] start_time = time.time() for epoch in range(args.num_epochs): print('start epoch ', epoch + 1) print('start train---------------------------------------------------') train_loss, train_accuracy, train_iou = metalearner.train( meta_train_dataloader, max_batches=args.num_batches, verbose=args.verbose, desc='Training', leave=False) print(f'\n train accuracy: {train_accuracy}, train loss: {train_loss}') print('end train---------------------------------------------------') train_losses.append(train_loss) train_accuracies.append(train_accuracy) train_ious.append(train_iou) # Evaluate in given intervals if epoch % args.val_step_size == 0: print( 'start evaluate-------------------------------------------------' ) results = metalearner.evaluate(meta_val_dataloader, max_batches=args.num_batches, verbose=args.verbose, desc=epoch_desc.format(epoch + 1), is_test=False) val_acc = results['accuracy'] val_loss = results['mean_outer_loss'] val_losses.append(val_loss) val_accuracies.append(val_acc) val_ious.append(results['iou']) print( f'\n validation accuracy: {val_acc}, validation loss: {val_loss}' ) print( 'end evaluate-------------------------------------------------' ) # Save best model if 'accuracies_after' in results: if (best_value is None) or (best_value < results['accuracies_after']): best_value = results['accuracies_after'] save_model = True elif (best_value is None) or (best_value > results['mean_outer_loss']): best_value = results['mean_outer_loss'] save_model = True else: save_model = False if save_model and (args.output_folder is not None): with open(args.model_path, 'wb') as f: torch.save(model.state_dict(), f) print('end epoch ', epoch + 1) elapsed_time = time.time() - start_time print('Finished after ', time.strftime('%H:%M:%S', time.gmtime(elapsed_time))) r = {} r['train_losses'] = train_losses r['train_accuracies'] = train_accuracies r['train_ious'] = train_ious r['val_losses'] = val_losses r['val_accuracies'] = val_accuracies r['val_ious'] = val_ious r['time'] = time.strftime('%H:%M:%S', time.gmtime(elapsed_time)) with open(os.path.join(output_folder, 'train_results.json'), 'w') as g: json.dump(r, g) logging.info('Saving results dict in `{0}`'.format( os.path.abspath(os.path.join(output_folder, 'train_results.json')))) # Plot results plot_errors(args.num_epochs, train_losses, val_losses, val_step_size=args.val_step_size, output_folder=output_folder, save=True, bce_dice_focal=bce_dice_focal) plot_accuracy(args.num_epochs, train_accuracies, val_accuracies, val_step_size=args.val_step_size, output_folder=output_folder, save=True) plot_iou(args.num_epochs, train_ious, val_ious, val_step_size=args.val_step_size, output_folder=output_folder, save=True) if hasattr(meta_train_dataset, 'close'): meta_train_dataset.close() meta_val_dataset.close()