# Instantiate point transformer
pt = PointTnf(use_cuda=use_cuda)
# Instatiate image transformers
tpsTnf = GeometricTnf(geometric_model='tps', use_cuda=use_cuda)
affTnf = GeometricTnf(geometric_model='affine', use_cuda=use_cuda)
for i, batch in enumerate(dataloader):
# get random batch of size 1
batch = batchTensorToVars(batch)
source_im_size = batch['source_im_size']
target_im_size = batch['target_im_size']
if do_aff:
model_aff.eval()
if do_tps:
model_tps.eval()
# Evaluate models
if do_aff:
theta_aff = model_aff(batch)
warped_image_aff = affTnf(batch['source_image'],
theta_aff.view(-1, 2, 3))
if do_tps:
theta_tps = model_tps(batch)
warped_image_tps = tpsTnf(batch['source_image'], theta_tps)
if do_aff and do_tps:
theta_aff_tps = model_tps({
if batch_idx % log_interval == 0:
print(mode.capitalize() +
' Epoch: {} [{}/{} ({:.0f}%)]\t\tLoss: {:.6f}'.format(
epoch, batch_idx, len(dataloader), 100. * batch_idx /
len(dataloader), loss_np))
epoch_loss /= len(dataloader)
print(mode.capitalize() + ' set: Average loss: {:.4f}'.format(epoch_loss))
return epoch_loss
train_loss = np.zeros(args.num_epochs)
test_loss = np.zeros(args.num_epochs)
print('Starting training...')
model.eval()
for epoch in range(1, args.num_epochs + 1):
model.train()
train_loss[epoch - 1] = process_epoch('train',
epoch,
model,
loss,
optimizer,
dataloader,
pair_generation_tnf,
log_interval=100)
model.eval()
test_loss[epoch - 1] = process_epoch('test',
epoch,
model,
class CNNGeometricMatcher:
def __init__(self,
use_extracted_features=False,
geometric_affine_model=None,
geometric_tps_model=None,
arch='resnet18',
featext_weights=None,
min_mutual_keypoints=6,
min_reprojection_error=200):
self.min_mutual_keypoints = min_mutual_keypoints
self.min_reprojection_error = min_reprojection_error
self.__do_affine = geometric_affine_model is not None
self.__do_tps = not use_extracted_features and geometric_tps_model is not None
self.__affTnf = GeometricTnf(geometric_model='affine',
use_cuda=_use_cuda)
if self.__do_affine:
checkpoint = torch.load(geometric_affine_model,
map_location=lambda storage, loc: storage)
print('Loading CNN Affine Geometric Model')
if use_extracted_features:
self.model_affine = CNNGeometricRegression(
use_cuda=use_cuda,
geometric_model='affine',
arch=arch,
featext_weights=featext_weights)
model_dict = self.model_affine.state_dict()
pretrained_dict = {
k: v
for k, v in checkpoint['state_dict'].items()
if k in model_dict
}
model_dict.update(pretrained_dict)
self.model_affine.load_state_dict(model_dict)
else:
self.model_affine = CNNGeometric(
use_cuda=_use_cuda,
geometric_model='affine',
arch=arch,
featext_weights=featext_weights)
self.model_affine.load_state_dict(checkpoint['state_dict'])
self.model_affine.eval()
if self.__do_tps:
self.model_tps = CNNGeometric(use_cuda=use_cuda,
geometric_model='tps',
arch=arch,
featext_weights=featext_weights)
checkpoint = torch.load(geometric_tps_model,
map_location=lambda storage, loc: storage)
print('Loading CNN TPS Geometric Model')
#self.model_tps.load_state_dict(checkpoint['state_dict'])
model_dict = self.model_tps.state_dict()
pretrained_dict = {
k: v
for k, v in checkpoint['state_dict'].items() if k in model_dict
}
model_dict.update(pretrained_dict)
self.model_tps.load_state_dict(model_dict)
self.model_tps.eval()
self.pt = PointTnf(use_cuda=_use_cuda)
def run(self, batch):
if self.__do_affine:
theta_aff, correlationAB, correlationBA = self.model_affine(batch)
if self.__do_tps:
if self.__do_affine:
warped_image_aff = self.__affTnf(batch['source_image'],
theta_aff.view(-1, 2, 3))
theta_affine_tps, correlationAB, correlationBA = self.model_tps(
{
'source_image': warped_image_aff,
'target_image': batch['target_image']
})
else:
theta_tps, correlationAB, correlationBA = self.model_tps(batch)
keypoints_A, keypoints_B = find_mutual_matached_keypoints(
correlationAB, correlationBA)
num_mutual_keypoints = keypoints_A.shape[0]
if num_mutual_keypoints < self.min_mutual_keypoints:
matched = False
reprojection_error = -1
else:
source_im_size = batch['source_im_size']
target_im_size = batch['target_im_size']
source_im_shape_np = source_im_size.data.numpy()
target_im_shape_np = target_im_size.data.numpy()
tensor_shape = correlationAB.data.numpy()[0].shape
im_keypointsA = tensorPointstoPixels(
keypoints_A,
tensor_size=tensor_shape,
im_size=(source_im_shape_np[0][1], source_im_shape_np[0][0]))
im_keypointsB = tensorPointstoPixels(
keypoints_B,
tensor_size=tensor_shape,
im_size=(target_im_shape_np[0][1], target_im_shape_np[0][0]))
torch_keypointsA_var = Variable(
torch.Tensor(
im_keypointsA.reshape(1, 2, -1).astype(np.float32)))
torch_keypointsB_var = Variable(
torch.Tensor(
im_keypointsB.reshape(1, 2, -1).astype(np.float32)))
target_points_norm = PointsToUnitCoords(torch_keypointsB_var,
target_im_size)
if self.__do_affine and self.__do_tps:
warped_points_aff_tps_norm = self.pt.tpsPointTnf(
theta_affine_tps, target_points_norm)
warped_points_norm = self.pt.affPointTnf(
theta_aff, warped_points_aff_tps_norm)
elif self.__do_affine:
warped_points_norm = self.pt.affPointTnf(
theta_aff, target_points_norm)
elif self.__do_tps:
warped_points_norm = self.pt.tpsPointTnf(
theta_tps, target_points_norm)
warped_points_aff = PointsToPixelCoords(warped_points_norm,
source_im_size)
reprojection_error = compute_reprojection_error(
torch_keypointsA_var, warped_points_aff)
matched = reprojection_error <= self.min_reprojection_error
return reprojection_error, matched, num_mutual_keypoints
batch = batchTensorToVars(batch)
print("batch", batch, batch['source_image'])
source_im_size = batch['source_im_size']
target_im_size = batch['target_im_size']
source_points = batch['source_points']
target_points = batch['target_points']
# warp points with estimated transformations
target_points_norm = PointsToUnitCoords(target_points, target_im_size)
if do_aff:
pass
#model_aff.eval()
if do_tps:
model_tps.eval()
# Evaluate models
if do_aff:
theta_aff = model_aff(batch)
print("theta_aff", theta_aff)
warped_image_aff = affTnf(batch['source_image'],
theta_aff.view(-1, 2, 3))
if do_tps:
theta_tps = model_tps(batch)
warped_image_tps = tpsTnf(batch['source_image'], theta_tps)
if do_aff and do_tps:
theta_aff_tps = model_tps({
'source_image': warped_image_aff,
for i, batch in enumerate(dataloader):
# get random batch of size 1
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']
# warp points with estimated transformations
target_points_norm = PointsToUnitCoords(target_points,target_im_size)
if do_aff:
model_aff.eval()
# Evaluate models
if do_aff:
theta_aff=model_aff(batch)
warped_image_aff = affTnf(batch['source_image'],theta_aff.view(-1,2,3))
print(theta_aff)
print(theta_aff.view(-1,2,3))
# Un-normalize images and convert to numpy
source_image = normalize_image(batch['source_image'],forward=False)
source_image = source_image.data.squeeze(0).transpose(0,1).transpose(1,2).cpu().numpy()
target_image = normalize_image(batch['target_image'],forward=False)
target_image = target_image.data.squeeze(0).transpose(0,1).transpose(1,2).cpu().numpy()
