def inlier_count(batch,
batch_start_idx,
theta_tps,
theta_aff_tps,
corr_aff,
corr_aff_tps,
model_tps,
stats,
args,
use_cuda=True):
inliersComposed = TwoStageWeakInlierCount(
use_cuda=torch.cuda.is_available(),
dilation_filter=0,
normalize_inlier_count=True)
inliers_comp = inliersComposed(matches=corr_aff,
theta_aff=theta_aff,
theta_aff_tps=theta_aff_tps,
model_tps=model_tps)
current_batch_size = batch['source_im_size'].size(0)
indices = range(batch_start_idx, batch_start_idx + current_batch_size)
stats['aff_tps']['inlier_count'][indices] = inliers_comp.unsqueeze(
1).cpu().data.numpy()
return stats
if args.train_fr and name.find('bn') != -1:
param.requires_grad = args.train_bn
# define loss
print('Using weak loss...')
if args.dilation_filter == 0:
dilation_filter = 0
else:
dilation_filter = generate_binary_structure(2, args.dilation_filter)
if args.inlier_loss:
inliersAffine = WeakInlierCount(geometric_model='affine',
seg_mask=args.mask,
**arg_groups['weak_loss'])
inliersComposed = TwoStageWeakInlierCount(use_cuda=use_cuda,
seg_mask=args.mask,
**arg_groups['weak_loss'])
if args.consistency_loss:
aff_ConsisLoss = ConsistencyLoss(use_cuda=use_cuda, transform='affine')
tps_ConsisLoss = ConsistencyLoss(use_cuda=use_cuda, transform='tps')
if args.transitivity_loss:
aff_TransLoss = TransitivityLoss(use_cuda=use_cuda, transform='affine')
tps_TransLoss = TransitivityLoss(use_cuda=use_cuda, transform='tps')
if args.consistency_loss or args.transitivity_loss:
coord = []
for i in range(config.NUM_OF_COORD):
for j in range(config.NUM_OF_COORD):
xx = []
if args.train_fr and name.find('bn') != -1:
param.requires_grad = args.train_bn
for name, param in model.GridDeformation2.named_parameters():
print(name.ljust(30), param.requires_grad)
# define loss
print('Using weak loss...')
if args.dilation_filter == 0:
dilation_filter = 0
else:
dilation_filter = generate_binary_structure(2, args.dilation_filter)
inliersAffine = WeakInlierCount(geometric_model='affine',
**arg_groups['weak_loss'])
inliersTps = WeakInlierCount(geometric_model='tps', **arg_groups['weak_loss'])
inliersComposed = TwoStageWeakInlierCount(use_cuda=use_cuda,
**arg_groups['weak_loss'])
def inlier_score_function(theta_aff,
theta_aff_tps,
corr_aff,
corr_aff_tps,
model_tps,
minimize_outliers=False):
inliers_comp = inliersComposed(matches=corr_aff,
theta_aff=theta_aff,
theta_aff_tps=theta_aff_tps,
model_tps=model_tps)
inliers_aff = inliersAffine(matches=corr_aff, theta=theta_aff)