def main(): """ Initialize model """ model = init_model(args, arg_groups, use_cuda) """ Initialize dataloader """ train_data, train_loader = init_train_data(args) eval_data, eval_loader = init_eval_data(args) """ Initialize optimizer """ model_opt = init_model_optim(args, model) batch_tnf = BatchTensorToVars(use_cuda=use_cuda) """ Evaluate initial condition """ ''' eval_categories = np.array(range(20)) + 1 eval_flag = np.zeros(len(eval_data)) for i in range(len(eval_data)): eval_flag[i] = sum(eval_categories == eval_data.category[i]) eval_idx = np.flatnonzero(eval_flag) model.eval() eval_stats = compute_metric(args.eval_metric, model, eval_data, eval_loader, batch_tnf, args) best_eval_pck = np.mean(eval_stats['aff_tps'][args.eval_metric][eval_idx]) ''' best_epoch = 1 """ Start training """ for epoch in range(1, args.num_epochs + 1): model.eval() process_epoch(epoch, model, model_opt, train_loader, batch_tnf) ''' model.eval() eval_stats = compute_metric(args.eval_metric, model, eval_data, eval_loader, batch_tnf, args) eval_pck = np.mean(eval_stats['aff_tps'][args.eval_metric][eval_idx]) is_best = eval_pck > best_eval_pck if eval_pck > best_eval_pck: best_eval_pck = eval_pck best_epoch = epoch print('eval: {:.3f}'.format(eval_pck), 'best eval: {:.3f}'.format(best_eval_pck), 'best epoch: {}'.format(best_epoch)) """ Early stopping """ if eval_pck < (best_eval_pck - 0.05): break ''' save_model(args, model, is_best)
def __call__(self, fname, fname2): image = io.imread(fname) image = np.expand_dims(image.transpose((2, 0, 1)), 0) image = torch.Tensor(image.astype(np.float32)) image_var = Variable(image, requires_grad=False) image_A = self.affTnf(image_var).data.squeeze(0) image_A_demo = self.affTnf_demo(image_var).data.squeeze(0) image_A_origin = self.affTnf_origin(image_var).data.squeeze(0) image2 = io.imread(fname2) image2 = np.expand_dims(image2.transpose((2, 0, 1)), 0) image2 = torch.Tensor(image2.astype(np.float32)) image_var2 = Variable(image2, requires_grad=False) image_B = self.affTnf(image_var2).data.squeeze(0) sample = {'source_image': image_A, 'target_image': image_B, 'demo': image_A_demo, 'origin_image': image_A_origin} sample = self.transform(sample) batchTensorToVars = BatchTensorToVars(use_cuda=self.use_cuda) batch = batchTensorToVars(sample) batch['source_image'] = torch.unsqueeze(batch['source_image'],0) batch['target_image'] = torch.unsqueeze(batch['target_image'],0) batch['origin_image'] = torch.unsqueeze(batch['origin_image'],0) batch['demo'] = torch.unsqueeze(batch['demo'],0) if self.do_aff: self.model_aff.eval() # Evaluate models if self.do_aff: theta_aff = self.model_aff(batch) warped_image_aff_demo = self.affTnf_demo(batch['demo'], theta_aff.view(-1, 2, 3)) if self.do_aff: warped_image_aff_demo = normalize_image(warped_image_aff_demo, forward=False) warped_image_aff_demo = warped_image_aff_demo.data.squeeze(0).transpose(0, 1).transpose(1, 2).cpu().numpy() print("Done") imsave('result.jpg', warped_image_aff_demo) return warped_image_aff_demo
map_location=lambda storage, loc: storage) checkpoint_tps['state_dict'] = OrderedDict([ (k.replace('vgg', 'model'), v) for k, v in checkpoint_tps['state_dict'].items() ]) for name, param in model.FeatureRegression2.state_dict().items(): model.FeatureRegression2.state_dict()[name].copy_( checkpoint_tps['state_dict']['FeatureRegression.' + name]) # Dataset and dataloader dataset = PFPascalDataset( csv_file=os.path.join(args.pf_path, 'test_pairs_pf_pascal.csv'), dataset_path=args.pf_path, transform=NormalizeImageDict(['source_image', 'target_image'])) dataloader = DataLoader(dataset, batch_size=1, shuffle=True, num_workers=4) batchTensorToVars = BatchTensorToVars(use_cuda=use_cuda) # Instatiate image transformers affTnf = GeometricTnf(geometric_model='affine', use_cuda=use_cuda) def affTpsTnf(source_image, theta_aff, theta_aff_tps, use_cuda=use_cuda): tpstnf = GeometricTnf(geometric_model='tps', use_cuda=use_cuda) sampling_grid = tpstnf(image_batch=source_image, theta_batch=theta_aff_tps, return_sampling_grid=True)[1] X = sampling_grid[:, :, :, 0].unsqueeze(3) Y = sampling_grid[:, :, :, 1].unsqueeze(3) Xp = X * theta_aff[:, 0].unsqueeze(1).unsqueeze( 2) + Y * theta_aff[:, 1].unsqueeze(1).unsqueeze( 2) + theta_aff[:, 2].unsqueeze(1).unsqueeze(2)
dataset = Dataset(csv_file=os.path.join(args.csv_path, csv_file), dataset_path=args.eval_dataset_path, transform=NormalizeImageDict(['source_image','target_image']), output_size=cnn_image_size) if use_cuda: batch_size=8 else: batch_size=1 dataloader = DataLoader(dataset, batch_size=batch_size, shuffle=False, num_workers=4, collate_fn=collate_fn) batch_tnf = BatchTensorToVars(use_cuda=use_cuda) if args.eval_dataset=='pf' or args.eval_dataset=='pf_pascal' or args.eval_dataset == 'pf_willow' or args.eval_dataset == 'tss-pck': metric = 'pck' elif args.eval_dataset=='caltech': metric = 'area' elif args.eval_dataset=='pascal-parts': metric = 'pascal_parts' elif args.eval_dataset=='tss': metric = 'flow' model.eval() stats=compute_metric(metric,model,dataset,dataloader,batch_tnf,batch_size,two_stage,do_aff,do_tps,args)
def main(): args, arg_groups = ArgumentParser(mode='train').parse() print(args) use_cuda = torch.cuda.is_available() use_me = args.use_me device = torch.device('cuda') if use_cuda else torch.device('cpu') # Seed # torch.manual_seed(args.seed) # if use_cuda: # torch.cuda.manual_seed(args.seed) # CNN model and loss print('Creating CNN model...') if args.geometric_model == 'affine_simple': cnn_output_dim = 3 elif args.geometric_model == 'affine_simple_4': cnn_output_dim = 4 else: raise NotImplementedError('Specified geometric model is unsupported') model = CNNGeometric(use_cuda=use_cuda, output_dim=cnn_output_dim, **arg_groups['model']) if args.geometric_model == 'affine_simple': init_theta = torch.tensor([0.0, 1.0, 0.0], device=device) elif args.geometric_model == 'affine_simple_4': init_theta = torch.tensor([0.0, 1.0, 0.0, 0.0], device=device) try: model.FeatureRegression.linear.bias.data += init_theta except: model.FeatureRegression.resnet.fc.bias.data += init_theta args.load_images = False if args.loss == 'mse': print('Using MSE loss...') loss = nn.MSELoss() elif args.loss == 'weighted_mse': print('Using weighted MSE loss...') loss = WeightedMSELoss(use_cuda=use_cuda) elif args.loss == 'reconstruction': print('Using reconstruction loss...') loss = ReconstructionLoss( int(np.rint(args.input_width * (1 - args.crop_factor) / 16) * 16), int(np.rint(args.input_height * (1 - args.crop_factor) / 16) * 16), args.input_height, use_cuda=use_cuda) args.load_images = True elif args.loss == 'combined': print('Using combined loss...') loss = CombinedLoss(args, use_cuda=use_cuda) if args.use_reconstruction_loss: args.load_images = True elif args.loss == 'grid': print('Using grid loss...') loss = SequentialGridLoss(use_cuda=use_cuda) else: raise NotImplementedError('Specifyed loss %s is not supported' % args.loss) # Initialize Dataset objects if use_me: dataset = MEDataset(geometric_model=args.geometric_model, dataset_csv_path=args.dataset_csv_path, dataset_csv_file='train.csv', dataset_image_path=args.dataset_image_path, input_height=args.input_height, input_width=args.input_width, crop=args.crop_factor, use_conf=args.use_conf, use_random_patch=args.use_random_patch, normalize_inputs=args.normalize_inputs, random_sample=args.random_sample, load_images=args.load_images) dataset_val = MEDataset(geometric_model=args.geometric_model, dataset_csv_path=args.dataset_csv_path, dataset_csv_file='val.csv', dataset_image_path=args.dataset_image_path, input_height=args.input_height, input_width=args.input_width, crop=args.crop_factor, use_conf=args.use_conf, use_random_patch=args.use_random_patch, normalize_inputs=args.normalize_inputs, random_sample=args.random_sample, load_images=args.load_images) else: dataset = SynthDataset(geometric_model=args.geometric_model, dataset_csv_path=args.dataset_csv_path, dataset_csv_file='train.csv', dataset_image_path=args.dataset_image_path, transform=NormalizeImageDict(['image']), random_sample=args.random_sample) dataset_val = SynthDataset(geometric_model=args.geometric_model, dataset_csv_path=args.dataset_csv_path, dataset_csv_file='val.csv', dataset_image_path=args.dataset_image_path, transform=NormalizeImageDict(['image']), random_sample=args.random_sample) # Set Tnf pair generation func if use_me: pair_generation_tnf = BatchTensorToVars(use_cuda=use_cuda) elif args.geometric_model == 'affine_simple' or args.geometric_model == 'affine_simple_4': pair_generation_tnf = SynthPairTnf(geometric_model='affine', use_cuda=use_cuda) else: raise NotImplementedError('Specified geometric model is unsupported') # Initialize DataLoaders dataloader = DataLoader(dataset, batch_size=args.batch_size, shuffle=True, num_workers=4) dataloader_val = DataLoader(dataset_val, batch_size=args.batch_size, shuffle=True, num_workers=4) # Optimizer optimizer = optim.Adam(model.FeatureRegression.parameters(), lr=args.lr) # Train # Set up names for checkpoints ckpt = args.trained_model_fn + '_' + args.geometric_model + '_' + args.loss + '_loss_' checkpoint_path = os.path.join(args.trained_model_dir, args.trained_model_fn, ckpt + '.pth.tar') if not os.path.exists(args.trained_model_dir): os.mkdir(args.trained_model_dir) # Set up TensorBoard writer if not args.log_dir: tb_dir = os.path.join(args.trained_model_dir, args.trained_model_fn + '_tb_logs') else: tb_dir = os.path.join(args.log_dir, args.trained_model_fn + '_tb_logs') logs_writer = SummaryWriter(tb_dir) # add graph, to do so we have to generate a dummy input to pass along with the graph if use_me: dummy_input = { 'mv_L2R': torch.rand([args.batch_size, 2, 216, 384], device=device), 'mv_R2L': torch.rand([args.batch_size, 2, 216, 384], device=device), 'grid_L2R': torch.rand([args.batch_size, 2, 216, 384], device=device), 'grid_R2L': torch.rand([args.batch_size, 2, 216, 384], device=device), 'grid': torch.rand([args.batch_size, 2, 216, 384], device=device), 'conf_L': torch.rand([args.batch_size, 1, 216, 384], device=device), 'conf_R': torch.rand([args.batch_size, 1, 216, 384], device=device), 'theta_GT': torch.rand([args.batch_size, 4], device=device), } if args.load_images: dummy_input['img_R_orig'] = torch.rand( [args.batch_size, 1, 216, 384], device=device) dummy_input['img_R'] = torch.rand([args.batch_size, 1, 216, 384], device=device) else: dummy_input = { 'source_image': torch.rand([args.batch_size, 3, 240, 240], device=device), 'target_image': torch.rand([args.batch_size, 3, 240, 240], device=device), 'theta_GT': torch.rand([args.batch_size, 2, 3], device=device) } logs_writer.add_graph(model, dummy_input) # Start of training print('Starting training...') best_val_loss = float("inf") max_batch_iters = len(dataloader) print('Iterations for one epoch:', max_batch_iters) epoch_to_change_lr = int(args.lr_max_iter / max_batch_iters * 2 + 0.5) # Loading checkpoint model, optimizer, start_epoch, best_val_loss, last_epoch = load_checkpoint( checkpoint_path, model, optimizer, device) # Scheduler if args.lr_scheduler == 'cosine': is_cosine_scheduler = True scheduler = torch.optim.lr_scheduler.CosineAnnealingLR( optimizer, T_max=args.lr_max_iter, eta_min=1e-7, last_epoch=last_epoch) elif args.lr_scheduler == 'cosine_restarts': is_cosine_scheduler = True scheduler = torch.optim.lr_scheduler.CosineAnnealingWarmRestarts( optimizer, T_0=args.lr_max_iter, T_mult=2, last_epoch=last_epoch) elif args.lr_scheduler == 'exp': is_cosine_scheduler = False if last_epoch > 0: last_epoch /= max_batch_iters scheduler = torch.optim.lr_scheduler.ExponentialLR( optimizer, gamma=args.lr_decay, last_epoch=last_epoch) # elif args.lr_scheduler == 'step': # step_scheduler = torch.optim.lr_scheduler.StepLR(optimizer, 10, gamma=0.1) # scheduler = False else: is_cosine_scheduler = False scheduler = False for epoch in range(1, start_epoch): if args.lr_scheduler == 'cosine' and (epoch % epoch_to_change_lr == 0): scheduler.state_dict()['base_lrs'][0] *= args.lr_decay torch.autograd.set_detect_anomaly(True) for epoch in range(start_epoch, args.num_epochs + 1): print('Current epoch: ', epoch) # we don't need the average epoch loss so we assign it to _ _ = train(epoch, model, loss, optimizer, dataloader, pair_generation_tnf, log_interval=args.log_interval, scheduler=scheduler, is_cosine_scheduler=is_cosine_scheduler, tb_writer=logs_writer) # Step non-cosine scheduler if scheduler and not is_cosine_scheduler: scheduler.step() val_loss = validate_model(model, loss, dataloader_val, pair_generation_tnf, epoch, logs_writer) # Change lr_max in cosine annealing if args.lr_scheduler == 'cosine' and (epoch % epoch_to_change_lr == 0): scheduler.state_dict()['base_lrs'][0] *= args.lr_decay if (epoch % epoch_to_change_lr == epoch_to_change_lr // 2) or epoch == 1: compute_metric('absdiff', model, args.geometric_model, None, None, dataset_val, dataloader_val, pair_generation_tnf, args.batch_size, args) # remember best loss is_best = val_loss < best_val_loss best_val_loss = min(val_loss, best_val_loss) save_checkpoint( { 'epoch': epoch + 1, 'args': args, 'state_dict': model.state_dict(), 'best_val_loss': best_val_loss, 'optimizer': optimizer.state_dict(), }, is_best, checkpoint_path) logs_writer.close() print('Done!')
def main(): # Argument parsing args, arg_groups = ArgumentParser(mode='eval').parse() print(args) # check provided models and deduce if single/two-stage model should be used two_stage = args.model_2 != '' if args.eval_dataset_path == '' and args.eval_dataset == 'pf': args.eval_dataset_path = 'datasets/proposal-flow-willow/' if args.eval_dataset_path == '' and args.eval_dataset == 'pf-pascal': args.eval_dataset_path = 'datasets/proposal-flow-pascal/' if args.eval_dataset_path == '' and args.eval_dataset == 'caltech': args.eval_dataset_path = 'datasets/caltech-101/' if args.eval_dataset_path == '' and args.eval_dataset == 'tss': args.eval_dataset_path = 'datasets/tss/' use_cuda = torch.cuda.is_available() # Download dataset if needed if args.eval_dataset == 'pf' and not exists(args.eval_dataset_path): download_PF_willow(args.eval_dataset_path) elif args.eval_dataset == 'pf-pascal' and not exists( args.eval_dataset_path): download_PF_pascal(args.eval_dataset_path) elif args.eval_dataset == 'caltech' and not exists(args.eval_dataset_path): download_caltech(args.eval_dataset_path) elif args.eval_dataset == 'tss' and not exists(args.eval_dataset_path): download_TSS(args.eval_dataset_path) print('Creating CNN model...') def create_model(model_filename): checkpoint = torch.load(model_filename, map_location=lambda storage, loc: storage) checkpoint['state_dict'] = OrderedDict([ (k.replace('vgg', 'model'), v) for k, v in checkpoint['state_dict'].items() ]) output_size = checkpoint['state_dict'][ 'FeatureRegression.linear.bias'].size()[0] if output_size == 6: geometric_model = 'affine' elif output_size == 8 or output_size == 9: geometric_model = 'hom' else: geometric_model = 'tps' model = CNNGeometric(use_cuda=use_cuda, output_dim=output_size, **arg_groups['model']) for name, param in model.FeatureExtraction.state_dict().items(): if not name.endswith('num_batches_tracked'): model.FeatureExtraction.state_dict()[name].copy_( checkpoint['state_dict']['FeatureExtraction.' + name]) for name, param in model.FeatureRegression.state_dict().items(): if not name.endswith('num_batches_tracked'): model.FeatureRegression.state_dict()[name].copy_( checkpoint['state_dict']['FeatureRegression.' + name]) return (model, geometric_model) # Load model for stage 1 model_1, geometric_model_1 = create_model(args.model_1) if two_stage: # Load model for stage 2 model_2, geometric_model_2 = create_model(args.model_2) else: model_2, geometric_model_2 = None, None #import pdb; pdb.set_trace() print('Creating dataset and dataloader...') # Dataset and dataloader if args.eval_dataset == 'pf': Dataset = PFDataset collate_fn = default_collate csv_file = 'test_pairs_pf.csv' if args.eval_dataset == 'pf-pascal': Dataset = PFPascalDataset collate_fn = default_collate csv_file = 'all_pairs_pf_pascal.csv' elif args.eval_dataset == 'caltech': Dataset = CaltechDataset collate_fn = default_collate csv_file = 'test_pairs_caltech_with_category.csv' elif args.eval_dataset == 'tss': Dataset = TSSDataset collate_fn = default_collate csv_file = 'test_pairs_tss.csv' cnn_image_size = (args.image_size, args.image_size) dataset = Dataset(csv_file=os.path.join(args.eval_dataset_path, csv_file), dataset_path=args.eval_dataset_path, transform=NormalizeImageDict( ['source_image', 'target_image']), output_size=cnn_image_size) if use_cuda: batch_size = args.batch_size else: batch_size = 1 dataloader = DataLoader(dataset, batch_size=batch_size, shuffle=False, num_workers=0, collate_fn=collate_fn) batch_tnf = BatchTensorToVars(use_cuda=use_cuda) if args.eval_dataset == 'pf' or args.eval_dataset == 'pf-pascal': metric = 'pck' elif args.eval_dataset == 'caltech': metric = 'area' elif args.eval_dataset == 'tss': metric = 'flow' model_1.eval() if two_stage: model_2.eval() print('Starting evaluation...') stats = compute_metric(metric, model_1, geometric_model_1, model_2, geometric_model_2, dataset, dataloader, batch_tnf, batch_size, args)
def main(passed_arguments=None): # Argument parsing args,arg_groups = ArgumentParser(mode='eval').parse(passed_arguments) print(args) # check provided models and deduce if single/two-stage model should be used two_stage = args.model_2 != '' use_cuda = torch.cuda.is_available() use_me = args.use_me print('Creating CNN model...') def create_model(model_filename): checkpoint = torch.load(model_filename, map_location=lambda storage, loc: storage) checkpoint['state_dict'] = OrderedDict([(k.replace('vgg', 'model'), v) for k, v in checkpoint['state_dict'].items()]) try: output_size = checkpoint['state_dict']['FeatureRegression.linear.bias'].size()[0] except: output_size = checkpoint['state_dict']['FeatureRegression.resnet.fc.bias'].size()[0] if output_size == 4: geometric_model = 'affine_simple_4' elif output_size == 3: geometric_model = 'affine_simple' else: raise NotImplementedError('Geometric model deducted from output layer is unsupported') model = CNNGeometric(use_cuda=use_cuda, output_dim=output_size, **arg_groups['model']) if use_me is False: for name, param in model.FeatureExtraction.state_dict().items(): if not name.endswith('num_batches_tracked'): model.FeatureExtraction.state_dict()[name].copy_(checkpoint['state_dict']['FeatureExtraction.' + name]) for name, param in model.FeatureRegression.state_dict().items(): if not name.endswith('num_batches_tracked'): model.FeatureRegression.state_dict()[name].copy_(checkpoint['state_dict']['FeatureRegression.' + name]) return (model,geometric_model) # Load model for stage 1 model_1, geometric_model_1 = create_model(args.model_1) if two_stage: # Load model for stage 2 model_2, geometric_model_2 = create_model(args.model_2) else: model_2,geometric_model_2 = None, None #import pdb; pdb.set_trace() print('Creating dataset and dataloader...') # Dataset and dataloader if args.eval_dataset == '3d' and use_me is False: cnn_image_size=(args.image_size,args.image_size) dataset = Dataset3D(csv_file = os.path.join(args.eval_dataset_path, 'all_pairs.csv'), dataset_path = args.eval_dataset_path, transform = NormalizeImageDict(['source_image','target_image']), output_size = cnn_image_size) collate_fn = default_collate elif args.eval_dataset == '3d' and use_me is True: cnn_image_size=(args.input_height, args.input_width) dataset = MEDataset(dataset_csv_path=args.eval_dataset_path, dataset_csv_file='all_pairs_3d.csv', dataset_image_path=args.eval_dataset_path, input_height=args.input_height, input_width=args.input_width, crop=args.crop_factor, use_conf=args.use_conf, use_random_patch=args.use_random_patch, normalize_inputs=args.normalize_inputs, geometric_model='EVAL', random_sample=False) collate_fn = default_collate else: raise NotImplementedError('Dataset is unsupported') if use_cuda: batch_size = args.batch_size else: batch_size = 1 dataloader = DataLoader(dataset, batch_size = batch_size, shuffle = False, num_workers=0, collate_fn = collate_fn) batch_tnf = BatchTensorToVars(use_cuda = use_cuda) if args.eval_dataset == '3d': metric = 'absdiff' else: raise NotImplementedError('Dataset is unsupported') model_1.eval() if two_stage: model_2.eval() print(os.path.basename(args.model_1)) print('Starting evaluation...', flush=True) stats=compute_metric(metric, model_1, geometric_model_1, model_2, geometric_model_2, dataset, dataloader, batch_tnf, batch_size, args) if args.eval_dataset_path.find('merged') >= 0: stats_fn = 'stats_merged.pkl' else: stats_fn = 'stats.pkl' stats_fn = os.path.join(os.path.dirname(args.model_1), stats_fn) save_dict(stats_fn, stats) return stats
def main(): print("eval pf dataset") os.environ["CUDA_VISIBLE_DEVICES"] = "0" # ntg_checkpoint_path = "/home/zlk/project/registration_cnn_ntg/trained_weight/output/voc2012_coco2014_NTG_resnet101.pth.tar" # ntg_checkpoint_path = "/home/zlk/project/registration_cnn_ntg/trained_weight/voc2011/checkpoint_voc2011_NTG_resnet101.pth.tar" # ntg_checkpoint_path = "/home/zlk/project/registration_cnn_ntg/trained_weight/voc2011/checkpoint_voc2011_20r_NTG_resnet101.pth.tar" # ntg_checkpoint_path = '/home/zlk/project/registration_cnn_ntg/trained_weight/three_channel/checkpoint_NTG_resnet101.pth.tar' small_aff_ntg_checkpoint_path = '/home/zlk/project/registration_cnn_ntg/trained_weight/three_channel/coco2014_small_aff_checkpoint_NTG_resnet101.pth.tar' ntg_checkpoint_path = '/home/zlk/project/registration_cnn_ntg/trained_weight/voc2011/best_checkpoint_voc2011_three_channel_paper_NTG_resnet101.pth.tar' # ntg_checkpoint_path = '/home/zlk/project/registration_cnn_ntg/trained_weight/voc2011_paper_affine/best_checkpoint_voc2011_NTG_resnet101.pth.tar' #ntg_checkpoint_path = "/home/zlk/project/registration_cnn_ntg/trained_weight/voc2011/checkpoint_voc2011_30r_NTG_resnet101.pth.tar" # image_path = '../datasets/row_data/VOC/3 # label_path = '../datasets/row_data/label_file/aff_param2.csv' #image_path = '../datasets/row_data/COCO/' #label_path = '../datasets/row_data/label_file/aff_param_coco.csv' pf_data_path = 'datasets/row_data/pf_data' batch_size = 128 # 加载模型 use_cuda = torch.cuda.is_available() ntg_model = createModel(ntg_checkpoint_path, use_cuda=use_cuda) small_aff_ntg_model = createModel(small_aff_ntg_checkpoint_path, use_cuda=use_cuda) dataset = PFDataset( csv_file=os.path.join(pf_data_path, 'test_pairs_pf.csv'), training_image_path=pf_data_path, transform=NormalizeImageDict(['source_image', 'target_image'])) dataloader = DataLoader(dataset, batch_size=batch_size, shuffle=False, num_workers=4) batchTensorToVars = BatchTensorToVars(use_cuda=use_cuda) pt = PointTnf(use_cuda=use_cuda) print('Computing PCK...') total_correct_points_aff = 0 ntg_total_correct_points_aff = 0 cnn_ntg_total_correct_points_aff = 0 total_correct_points_tps = 0 total_correct_points_aff_tps = 0 total_points = 0 ntg_total_points = 0 cnn_ntg_total_points = 0 for i, batch in enumerate(dataloader): batch = batchTensorToVars(batch) source_im_size = batch['source_im_size'] target_im_size = batch['target_im_size'] source_points = batch['source_points'] target_points = batch['target_points'] source_image_batch = batch['source_image'] target_image_batch = batch['target_image'] # warp points with estimated transformations target_points_norm = PointsToUnitCoords(target_points, target_im_size) theta_estimate_batch = ntg_model(batch) #warped_image_batch = affine_transform_pytorch(source_image_batch, theta_estimate_batch) #batch['source_image'] = warped_image_batch #theta_estimate_batch = small_aff_ntg_model(batch) # 将pytorch的变换参数转为opencv的变换参数 #theta_opencv = theta2param(theta_estimate_batch.view(-1, 2, 3), 240, 240, use_cuda=use_cuda) # P5使用传统NTG方法进行优化cnn的结果 #cnn_ntg_param_batch = estimate_param_batch(source_image_batch, target_image_batch, theta_opencv,itermax = 600) #theta_pytorch = param2theta(cnn_ntg_param_batch.view(-1, 2, 3),240,240,use_cuda=use_cuda) # theta_opencv = theta2param(theta_estimate_batch.view(-1, 2, 3), 240, 240, use_cuda=use_cuda) # with torch.no_grad(): # ntg_param_batch = estimate_aff_param_iterator(source_image_batch[:, 0, :, :].unsqueeze(1), # target_image_batch[:, 0, :, :].unsqueeze(1), # None, use_cuda=use_cuda, itermax=600) # # cnn_ntg_param_batch = estimate_aff_param_iterator(source_image_batch[:, 0, :, :].unsqueeze(1), # target_image_batch[:, 0, :, :].unsqueeze(1), # theta_opencv, use_cuda=use_cuda, itermax=600) # # ntg_param_pytorch_batch = param2theta(ntg_param_batch,240, 240, use_cuda=use_cuda) # cnn_ntg_param_pytorch_batch = param2theta(cnn_ntg_param_batch,240, 240, use_cuda=use_cuda) warped_points_aff_norm = pt.affPointTnf(theta_estimate_batch, target_points_norm) warped_points_aff = PointsToPixelCoords(warped_points_aff_norm, source_im_size) # ntg_warped_points_aff_norm = pt.affPointTnf(ntg_param_pytorch_batch, target_points_norm) # ntg_warped_points_aff = PointsToPixelCoords(ntg_warped_points_aff_norm, source_im_size) # # cnn_ntg_warped_points_aff_norm = pt.affPointTnf(cnn_ntg_param_pytorch_batch, target_points_norm) # cnn_ntg_warped_points_aff = PointsToPixelCoords(cnn_ntg_warped_points_aff_norm, source_im_size) L_pck = batch['L_pck'].data correct_points_aff, num_points = correct_keypoints( source_points.data, warped_points_aff.data, L_pck) # ntg_correct_points_aff, ntg_num_points = correct_keypoints(source_points.data, # ntg_warped_points_aff.data, L_pck) # cnn_ntg_correct_points_aff, cnn_ntg_num_points = correct_keypoints(source_points.data, # cnn_ntg_warped_points_aff.data, L_pck) total_correct_points_aff += correct_points_aff total_points += num_points # ntg_total_correct_points_aff += ntg_correct_points_aff # ntg_total_points += ntg_num_points # # cnn_ntg_total_correct_points_aff += cnn_ntg_correct_points_aff # cnn_ntg_total_points += cnn_ntg_num_points print('Batch: [{}/{} ({:.0f}%)]'.format(i, len(dataloader), 100. * i / len(dataloader))) total_correct_points_aff = total_correct_points_aff.__float__() # ntg_total_correct_points_aff = ntg_total_correct_points_aff.__float__() # cnn_ntg_total_correct_points_aff = cnn_ntg_total_correct_points_aff.__float__() PCK_aff = total_correct_points_aff / total_points # ntg_PCK_aff=ntg_total_correct_points_aff/ntg_total_points # cnn_ntg_PCK_aff=cnn_ntg_total_correct_points_aff/cnn_ntg_total_points print('PCK affine:', PCK_aff) # print('ntg_PCK affine:',ntg_PCK_aff) # print('cnn_ntg_PCK affine:',cnn_ntg_PCK_aff) print('Done!')