def load_torch_model(args, use_cuda):
model = CNNGeometric(use_cuda=use_cuda, geometric_model='affine',
feature_extraction_cnn=args.feature_extraction_cnn)
# Load trained weights
print('Loading trained model weights...')
checkpoint = torch.load(args.pretrained, map_location=lambda storage, loc: storage)
checkpoint['state_dict'] = OrderedDict(
[(k.replace('vgg', 'model'), v) for k, v in checkpoint['state_dict'].items()])
model.load_state_dict(checkpoint['state_dict'])
return 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)
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)
def load_model(aff_params_path = '', aff_feat_ext = 'wormbrain_1', aff_feat_reg = 'simpler',
tps_params_path = '', tps_feat_ext = 'wormbrain_1', tps_feat_reg = 'simpler'):
"""
Loads a model. Assumes that each model (Affine and Thin-Plate Spline)
have been trained separately. Must specify the architecture used for feature_extraction
By default, it is resnet101
"""
use_cuda = torch.cuda.is_available()
#Only create a model for which weights have been provided
do_aff = not aff_params_path==''
do_tps = not tps_params_path==''
if not do_aff and not do_tps:
print("No weights found. Models not created, exiting.")
return
print("Creating CNN model.")
if do_aff:
model_aff = CNNGeometric(output_dim=6,use_cuda=use_cuda,
feature_extraction_cnn= aff_feat_ext,
feature_regression=aff_feat_reg)
if do_tps:
model_tps = CNNGeometric(output_dim=18,use_cuda=use_cuda,
feature_extraction_cnn= tps_feat_ext,
feature_regression = tps_feat_reg)
print("Loading trained model weights.")
if do_aff: #Loading affine model
if aff_feat_ext == 'resnet101': aff_feat_ext = 'resnet'
checkpoint = torch.load(aff_params_path, map_location=lambda storage, loc: storage)
checkpoint['state_dict'] = OrderedDict([(k.replace(aff_feat_ext, 'model'), v) for k, v in checkpoint['state_dict'].items()])
model_aff.load_state_dict(checkpoint['state_dict'])
print('Weights for Affine model loaded.')
if do_tps: #Loading thin plate spline model
if tps_feat_ext == 'resnet101': aff_feat_ext = 'resnet'
checkpoint = torch.load(tps_params_path, map_location=lambda storage, loc: storage)
checkpoint['state_dict'] = OrderedDict([(k.replace(aff_feat_ext, 'model'), v) for k, v in checkpoint['state_dict'].items()])
model_tps.load_state_dict(checkpoint['state_dict'])
print('Weights for Thin-Plate Spline model loaded.')
print('Returning model(s).')
if do_aff and not do_tps:
return model_aff
if do_tps and not do_aff:
return model_tps
if do_aff and do_tps:
return model_aff, model_tps
print('Creating CNN model...')
# check type of given model and create model
# if two_stage:
# model = TwoStageCNNGeometric(use_cuda=use_cuda,
# **arg_groups['model'])
# if not two_stage and do_aff:
# model = CNNGeometric(use_cuda=use_cuda,
# output_dim=6,
# **arg_groups['model'])
# if not two_stage and do_tps:
# model_tps = CNNGeometric(use_cuda=use_cuda,
# output_dim=18,
# **arg_groups['model'])
model = CNNGeometric(use_cuda=use_cuda, **arg_groups['model'])
# load pretrained weights
# if two_stage and args.model!='':
# checkpoint = torch.load(args.model, map_location=lambda storage, loc: storage)
# checkpoint['state_dict'] = OrderedDict([(k.replace('vgg', 'model'), v) for k, v in checkpoint['state_dict'].items()])
# for name, param in model.FeatureExtraction.state_dict().items():
# model.FeatureExtraction.state_dict()[name].copy_(checkpoint['state_dict']['FeatureExtraction.' + name])
# for name, param in model.FeatureRegression.state_dict().items():
# model.FeatureRegression.state_dict()[name].copy_(checkpoint['state_dict']['FeatureRegression.' + name])
# for name, param in model.FeatureRegression2.state_dict().items():
# model.FeatureRegression2.state_dict()[name].copy_(checkpoint['state_dict']['FeatureRegression2.' + name])
# for name, param in model.GridDeformation2.state_dict().items():
# model.GridDeformation2.state_dict()[name].copy_(checkpoint['state_dict']['GridDeformation2.' + name])
if args.dataset_csv_path == '' and args.geometric_model=='affine':
args.dataset_csv_path = 'training_data/pascal-synth-aff'
elif args.dataset_csv_path == '' and args.geometric_model=='tps':
args.dataset_csv_path = 'training_data/pascal-synth-tps'
arg_groups['dataset']['dataset_image_path']=args.dataset_image_path
# CNN model and loss
print('Creating CNN model...')
if args.geometric_model=='affine':
cnn_output_dim = 6
elif args.geometric_model=='tps':
cnn_output_dim = 18
model = CNNGeometric(use_cuda=use_cuda,
output_dim=cnn_output_dim,
**arg_groups['model'])
# Load pretrained model
if args.model!='':
checkpoint = torch.load(args.model, map_location=lambda storage, loc: storage)
model.load_state_dict(checkpoint['state_dict'])
#checkpoint['state_dict'] = OrderedDict([(k.replace('vgg', 'model'), v) for k, v in checkpoint['state_dict'].items()])
#for name, param in model.FeatureExtraction.state_dict().items():
# model.FeatureExtraction.state_dict()[name].copy_(checkpoint['state_dict']['FeatureExtraction.' + name])
#for name, param in model.FeatureRegression.state_dict().items():
# model.FeatureRegression.state_dict()[name].copy_(checkpoint['state_dict']['FeatureRegression.' + name])
if args.use_mse_loss:
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!')
args = parser.parse_args()
use_cuda = torch.cuda.is_available()
do_aff = not args.model_aff == ''
do_tps = not args.model_tps == ''
# Download dataset if needed
download_PF_willow('datasets/')
# Create model
print('Creating CNN model...')
if do_aff:
model_aff = CNNGeometric(
use_cuda=use_cuda,
geometric_model='affine',
feature_extraction_cnn=args.feature_extraction_cnn)
if do_tps:
model_tps = CNNGeometric(
use_cuda=use_cuda,
geometric_model='tps',
feature_extraction_cnn=args.feature_extraction_cnn)
# Load trained weights
print('Loading trained model weights...')
if do_aff:
checkpoint = torch.load(args.model_aff,
map_location=lambda storage, loc: storage)
checkpoint['state_dict'] = OrderedDict([
(k.replace('vgg', 'model'), v)
for k, v in checkpoint['state_dict'].items()
args.dataset_csv_path = 'training_data/pascal-synth-aff'
elif args.dataset_csv_path == '' and args.geometric_model == 'tps':
args.dataset_csv_path = 'training_data/pascal-synth-tps'
arg_groups['dataset']['dataset_image_path'] = args.dataset_image_path
# CNN model and loss
print('Creating CNN model...')
if args.geometric_model == 'affine':
cnn_output_dim = 6
elif args.geometric_model == 'tps':
cnn_output_dim = 256
# cnn_output_dim = 256
model = CNNGeometric(use_cuda=use_cuda,
output_dim=cnn_output_dim,
**arg_groups['model'])
# Load pretrained model
if args.model != '':
checkpoint = torch.load(args.model,
map_location=lambda storage, loc: storage)
checkpoint['state_dict'] = OrderedDict([
(k.replace('vgg', 'model'), v)
for k, v in checkpoint['state_dict'].items()
])
for name, param in model.FeatureExtraction.state_dict().items():
model.FeatureExtraction.state_dict()[name].copy_(
checkpoint['state_dict']['FeatureExtraction.' + name])
for name, param in model.FeatureRegression.state_dict().items():
model.FeatureRegression.state_dict()[name].copy_(
def main():
args, arg_groups = ArgumentParser(mode='train').parse()
print(args)
use_cuda = torch.cuda.is_available()
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)
# Download dataset if needed and set paths
if args.training_dataset == 'pascal':
if args.dataset_image_path == '' and not os.path.exists(
'datasets/pascal-voc11/TrainVal'):
download_pascal('datasets/pascal-voc11/')
if args.dataset_image_path == '':
args.dataset_image_path = 'datasets/pascal-voc11/'
args.dataset_csv_path = 'training_data/pascal-random'
# CNN model and loss
print('Creating CNN model...')
if args.geometric_model == 'affine':
cnn_output_dim = 6
elif args.geometric_model == 'hom' and args.four_point_hom:
cnn_output_dim = 8
elif args.geometric_model == 'hom' and not args.four_point_hom:
cnn_output_dim = 9
elif args.geometric_model == 'tps':
cnn_output_dim = 18
model = CNNGeometric(use_cuda=use_cuda,
output_dim=cnn_output_dim,
**arg_groups['model'])
if args.geometric_model == 'hom' and not args.four_point_hom:
init_theta = torch.tensor([1, 0, 0, 0, 1, 0, 0, 0, 1], device=device)
model.FeatureRegression.linear.bias.data += init_theta
if args.geometric_model == 'hom' and args.four_point_hom:
init_theta = torch.tensor([-1, -1, 1, 1, -1, 1, -1, 1], device=device)
model.FeatureRegression.linear.bias.data += init_theta
if args.use_mse_loss:
print('Using MSE loss...')
loss = nn.MSELoss()
else:
print('Using grid loss...')
loss = TransformedGridLoss(use_cuda=use_cuda,
geometric_model=args.geometric_model)
# Initialize Dataset objects
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
pair_generation_tnf = SynthPairTnf(geometric_model=args.geometric_model,
use_cuda=use_cuda)
# 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 and eventual scheduler
optimizer = optim.Adam(model.FeatureRegression.parameters(), lr=args.lr)
if args.lr_scheduler:
scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(
optimizer, T_max=args.lr_max_iter, eta_min=1e-6)
# scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(optimizer, 'min')
else:
scheduler = False
# Train
# Set up names for checkpoints
if args.use_mse_loss:
ckpt = args.trained_model_fn + '_' + args.geometric_model + '_mse_loss' + args.feature_extraction_cnn
checkpoint_path = os.path.join(args.trained_model_dir,
args.trained_model_fn,
ckpt + '.pth.tar')
else:
ckpt = args.trained_model_fn + '_' + args.geometric_model + '_grid_loss' + args.feature_extraction_cnn
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
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([16, 2, 3], device=device)
}
logs_writer.add_graph(model, dummy_input)
# Start of training
print('Starting training...')
best_val_loss = float("inf")
for epoch in range(1, args.num_epochs + 1):
# 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,
tb_writer=logs_writer)
val_loss = validate_model(model, loss, dataloader_val,
pair_generation_tnf, epoch, logs_writer)
# 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 __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)
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
if args.dataset_csv_path == '' and args.geometric_model == 'affine':
args.dataset_csv_path = 'training_data/pascal-synth-aff'
elif args.dataset_csv_path == '' and args.geometric_model == 'tps':
args.dataset_csv_path = 'training_data/pascal-synth-tps'
arg_groups['dataset']['dataset_image_path'] = args.dataset_image_path
# CNN model and loss
print('Creating CNN model...')
if args.geometric_model == 'affine':
cnn_output_dim = 6
elif args.geometric_model == 'tps':
cnn_output_dim = 18
model = CNNGeometric(use_cuda=use_cuda,
output_dim=cnn_output_dim,
**arg_groups['model'])
# Load pretrained model
if args.model != '':
checkpoint = torch.load(args.model,
map_location=lambda storage, loc: storage)
checkpoint['state_dict'] = OrderedDict([
(k.replace('vgg', 'model'), v)
for k, v in checkpoint['state_dict'].items()
])
for name, param in model.FeatureExtraction.state_dict().items():
model.FeatureExtraction.state_dict()[name].copy_(
checkpoint['state_dict']['FeatureExtraction.' + name])
for name, param in model.FeatureRegression.state_dict().items():
training_image_path=args.pf_path,
transform=NormalizeImageDict(
['source_image', 'target_image']))
batchTensorToVars = BatchTensorToVars()
# Instantiate point transformer
pt = PointTnf()
# Instatiate image transformers
affTnf = GeometricTnf(geometric_model='affine')
with tf.Graph().as_default():
# Create model
print('Creating CNN model...')
model_aff = CNNGeometric(
feature_extraction_cnn=args.feature_extraction_cnn)
init_op = tf.global_variables_initializer()
saver = tf.train.Saver()
with tf.Session() as sess:
sess.run(init_op)
# Load trained weights
print('Loading trained model weights...')
saver.restore(sess, args.model_aff)
# reload_checkpoint(sess, saver, args.model_aff)
while True:
# get random batch of size 1
batch = choice(dataset)
batch = batchTensorToVars(batch)
use_cuda = torch.cuda.is_available()
# Argument parsing
parser = argparse.ArgumentParser(description='CNNGeometric PyTorch implementation')
# Paths
parser.add_argument('--model', type=str, default='trained_models/best_checkpoint_resnet18_adam_pose_mse_loss.pth.tar', help='Trained affine model filename')
#parser.add_argument('--model-tps', type=str, default='trained_models/best_pascal_checkpoint_adam_tps_grid_loss.pth.tar', help='Trained TPS model filename')
parser.add_argument('--path', type=str, default='/home/develop/Work/Datasets/', help='Path to PF dataset')
parser.add_argument('--pairs', type=str, default='/home/develop/Work/Datasets/gardens_pairs_path_samples_sift_RANSAC_12kps.csv', help='Path to PF dataset')
args = parser.parse_args()
dataset_path=args.path
dataset_pairs_file = args.pairs
# Create model
print('Creating CNN model...')
model = CNNGeometric(use_cuda=use_cuda,geometric_model='pose',arch = 'resnet18')
print('Load CNN Weights ...')
checkpoint = torch.load(args.model, map_location=lambda storage, loc: storage)
model.load_state_dict(checkpoint['state_dict'])
# Dataset and dataloader
dataset = PlacesDataset(csv_file=dataset_pairs_file,
training_image_path=dataset_path,
transform=NormalizeImageDict(['source_image','target_image']))
dataloader = DataLoader(dataset, batch_size=1,
shuffle=True, num_workers=4)
batchTensorToVars = BatchTensorToVars(use_cuda=use_cuda)
for source_im_path,target_im_path,batch in dataloader:
# get random batch of size 1
if args.dataset_csv_path == '' and args.geometric_model=='affine':
args.dataset_csv_path = 'training_data/pascal-synth-aff'
elif args.dataset_csv_path == '' and args.geometric_model=='tps':
args.dataset_csv_path = 'training_data/pascal-synth-tps'
arg_groups['dataset']['dataset_image_path']=args.dataset_image_path
# CNN model and loss
print('Creating CNN model...')
if args.geometric_model=='affine':
cnn_output_dim = 6
elif args.geometric_model=='tps':
cnn_output_dim = 18
model = CNNGeometric(use_cuda=use_cuda,
output_dim=cnn_output_dim,
**arg_groups['model'])
# Load pretrained model
if args.model!='':
checkpoint = torch.load(args.model, map_location=lambda storage, loc: storage)
checkpoint['state_dict'] = OrderedDict([(k.replace('vgg', 'model'), v) for k, v in checkpoint['state_dict'].items()])
for name, param in model.FeatureExtraction.state_dict().items():
model.FeatureExtraction.state_dict()[name].copy_(checkpoint['state_dict']['FeatureExtraction.' + name])
for name, param in model.FeatureRegression.state_dict().items():
model.FeatureRegression.state_dict()[name].copy_(checkpoint['state_dict']['FeatureRegression.' + name])
if args.use_mse_loss:
print('Using MSE loss...')
tf.set_random_seed(args.seed)
# Download dataset if needed and set paths
if args.training_dataset == 'pascal':
if args.training_image_path == '':
download_pascal('datasets/pascal-voc11/')
args.training_image_path = 'datasets/pascal-voc11/'
if args.training_tnf_csv == '' and args.geometric_model == 'affine':
args.training_tnf_csv = 'training_data/pascal-synth-aff'
elif args.training_tnf_csv == '' and args.geometric_model == 'tps':
args.training_tnf_csv = 'training_data/pascal-synth-tps'
# CNN model and loss
print('Creating CNN model...')
model = CNNGeometric(feature_extraction_cnn=args.feature_extraction_cnn,
training=False)
if args.use_mse_loss:
print('Using MSE loss...')
loss = tf.losses
else:
print('Using grid loss...')
loss = TransformedGridLoss(geometric_model=args.geometric_model)
# Dataset and dataloader
dataset = SynthDataset(geometric_model=args.geometric_model,
csv_file=os.path.join(args.training_tnf_csv,
'train.csv'),
training_image_path=args.training_image_path,
transform=NormalizeImageDict(['image']),
random_sample=args.random_sample)
def main():
args = parse_flags()
use_cuda = torch.cuda.is_available()
# Seed
torch.manual_seed(args.seed)
if use_cuda:
torch.cuda.manual_seed(args.seed)
# Download dataset if needed and set paths
if args.training_dataset == 'pascal':
if args.training_image_path == '':
download_pascal('datasets/pascal-voc11/')
args.training_image_path = 'datasets/pascal-voc11/'
if args.training_tnf_csv == '' and args.geometric_model == 'affine':
args.training_tnf_csv = 'training_data/pascal-synth-aff'
elif args.training_tnf_csv == '' and args.geometric_model == 'tps':
args.training_tnf_csv = 'training_data/pascal-synth-tps'
# CNN model and loss
if not args.pretrained:
if args.light_model:
print('Creating light CNN model...')
model = LightCNN(use_cuda=use_cuda,
geometric_model=args.geometric_model)
else:
print('Creating CNN model...')
model = CNNGeometric(
use_cuda=use_cuda,
geometric_model=args.geometric_model,
feature_extraction_cnn=args.feature_extraction_cnn)
else:
model = load_torch_model(args, use_cuda)
if args.loss == 'mse':
print('Using MSE loss...')
loss = MSELoss()
elif args.loss == 'sum':
print('Using the sum of MSE and grid loss...')
loss = GridLossWithMSE(use_cuda=use_cuda,
geometric_model=args.geometric_model)
else:
print('Using grid loss...')
loss = TransformedGridLoss(use_cuda=use_cuda,
geometric_model=args.geometric_model)
# Initialize csv paths
train_csv_path_list = glob(
os.path.join(args.training_tnf_csv, '*train.csv'))
if len(train_csv_path_list) > 1:
print(
"!!!!WARNING!!!! multiple train csv files found, using first in glob order"
)
elif not len(train_csv_path_list):
raise FileNotFoundError(
"No training csv where found in the specified path!!!")
train_csv_path = train_csv_path_list[0]
val_csv_path_list = glob(os.path.join(args.training_tnf_csv, '*val.csv'))
if len(val_csv_path_list) > 1:
print(
"!!!!WARNING!!!! multiple train csv files found, using first in glob order"
)
elif not len(val_csv_path_list):
raise FileNotFoundError(
"No training csv where found in the specified path!!!")
val_csv_path = val_csv_path_list[0]
# Initialize Dataset objects
if args.coupled_dataset:
# Dataset for train and val if dataset is already coupled
dataset = CoupledDataset(geometric_model=args.geometric_model,
csv_file=train_csv_path,
training_image_path=args.training_image_path,
transform=NormalizeImageDict(
['image_a', 'image_b']))
dataset_val = CoupledDataset(
geometric_model=args.geometric_model,
csv_file=val_csv_path,
training_image_path=args.training_image_path,
transform=NormalizeImageDict(['image_a', 'image_b']))
# Set Tnf pair generation func
pair_generation_tnf = CoupledPairTnf(use_cuda=use_cuda)
else:
# Standard Dataset for train and val
dataset = SynthDataset(geometric_model=args.geometric_model,
csv_file=train_csv_path,
training_image_path=args.training_image_path,
transform=NormalizeImageDict(['image']),
random_sample=args.random_sample)
dataset_val = SynthDataset(
geometric_model=args.geometric_model,
csv_file=val_csv_path,
training_image_path=args.training_image_path,
transform=NormalizeImageDict(['image']),
random_sample=args.random_sample)
# Set Tnf pair generation func
pair_generation_tnf = SynthPairTnf(
geometric_model=args.geometric_model, use_cuda=use_cuda)
# 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 and eventual scheduler
optimizer = Adam(model.FeatureRegression.parameters(), lr=args.lr)
if args.lr_scheduler:
if args.scheduler_type == 'cosine':
print('Using cosine learning rate scheduler')
scheduler = CosineAnnealingLR(optimizer,
T_max=args.lr_max_iter,
eta_min=args.lr_min)
elif args.scheduler_type == 'decay':
print('Using decay learning rate scheduler')
scheduler = ReduceLROnPlateau(optimizer, 'min')
else:
print(
'Using truncated cosine with decay learning rate scheduler...')
scheduler = TruncateCosineScheduler(optimizer, len(dataloader),
args.num_epochs - 1)
else:
scheduler = False
# Train
# Set up names for checkpoints
if args.loss == 'mse':
ckpt = args.trained_models_fn + '_' + args.geometric_model + '_mse_loss' + args.feature_extraction_cnn
checkpoint_path = os.path.join(args.trained_models_dir,
args.trained_models_fn,
ckpt + '.pth.tar')
elif args.loss == 'sum':
ckpt = args.trained_models_fn + '_' + args.geometric_model + '_sum_loss' + args.feature_extraction_cnn
checkpoint_path = os.path.join(args.trained_models_dir,
args.trained_models_fn,
ckpt + '.pth.tar')
else:
ckpt = args.trained_models_fn + '_' + args.geometric_model + '_grid_loss' + args.feature_extraction_cnn
checkpoint_path = os.path.join(args.trained_models_dir,
args.trained_models_fn,
ckpt + '.pth.tar')
if not os.path.exists(args.trained_models_dir):
os.mkdir(args.trained_models_dir)
# Set up TensorBoard writer
if not args.log_dir:
tb_dir = os.path.join(args.trained_models_dir,
args.trained_models_fn + '_tb_logs')
else:
tb_dir = os.path.join(args.log_dir,
args.trained_models_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
dummy_input = {
'source_image': torch.rand([args.batch_size, 3, 240, 240]),
'target_image': torch.rand([args.batch_size, 3, 240, 240]),
'theta_GT': torch.rand([16, 2, 3])
}
logs_writer.add_graph(model, dummy_input)
# START OF TRAINING #
print('Starting training...')
best_val_loss = float("inf")
for epoch in range(1, args.num_epochs + 1):
# 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,
tb_writer=logs_writer)
val_loss = validate_model(model,
loss,
dataloader_val,
pair_generation_tnf,
epoch,
logs_writer,
coupled=args.coupled_dataset)
# 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!')
# Download dataset if needed and set paths
if args.training_dataset == 'pascal':
if args.training_image_path == '':
download_pascal('datasets/pascal-voc11/')
args.training_image_path = 'datasets/pascal-voc11/'
if args.training_tnf_csv == '' and args.geometric_model == 'affine':
args.training_tnf_csv = 'training_data/pascal-synth-aff'
elif args.training_tnf_csv == '' and args.geometric_model == 'tps':
args.training_tnf_csv = 'training_data/pascal-synth-tps'
# CNN model and loss
print('Creating CNN model...')
model = CNNGeometric(use_cuda=use_cuda,
geometric_model=args.geometric_model,
feature_extraction_cnn=args.feature_extraction_cnn)
if args.use_mse_loss:
print('Using MSE loss...')
loss = nn.MSELoss()
else:
print('Using grid loss...')
loss = TransformedGridLoss(use_cuda=use_cuda,
geometric_model=args.geometric_model)
# Dataset and dataloader
dataset = SynthDataset(geometric_model=args.geometric_model,
csv_file=os.path.join(args.training_tnf_csv,
'train.csv'),
training_image_path=args.training_image_path,
torch.cuda.manual_seed(args.seed)
# Download dataset if needed and set paths
if args.training_dataset == 'pascal':
if args.training_image_path == '':
download_pascal('datasets/pascal-voc11/')
args.training_image_path = 'datasets/pascal-voc11/'
if args.training_tnf_csv == '' and args.geometric_model=='affine':
args.training_tnf_csv = 'training_data/pascal-synth-aff'
elif args.training_tnf_csv == '' and args.geometric_model=='tps':
args.training_tnf_csv = 'training_data/pascal-synth-tps'
# CNN model and loss
print('Creating CNN model...')
model = CNNGeometric(use_cuda=use_cuda,geometric_model=args.geometric_model)
if args.use_mse_loss:
print('Using MSE loss...')
loss = nn.MSELoss()
else:
print('Using grid loss...')
loss = TransformedGridLoss(use_cuda=use_cuda,geometric_model=args.geometric_model)
# Dataset and dataloader
dataset = SynthDataset(geometric_model=args.geometric_model,
csv_file=os.path.join(args.training_tnf_csv,'train.csv'),
training_image_path=args.training_image_path,
transform=NormalizeImageDict(['image']))
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)
elif args.eval_dataset=='pascal-parts' and not exists(args.eval_dataset_path):
download_pascal_parts(args.eval_dataset_path)
# Create model
print('Creating CNN model...')
# check type of given model and create model
if two_stage:
model = TwoStageCNNGeometric(use_cuda=use_cuda,
**arg_groups['model'])
if not two_stage and do_aff:
model = CNNGeometric(use_cuda=use_cuda,
output_dim=6,
**arg_groups['model'])
if not two_stage and do_tps:
model_tps = CNNGeometric(use_cuda=use_cuda,
output_dim=18,
**arg_groups['model'])
# load pretrained weights
if two_stage and args.model!='':
checkpoint = torch.load(args.model, map_location=lambda storage, loc: storage)
checkpoint['state_dict'] = OrderedDict([(k.replace('vgg', 'model'), v) for k, v in checkpoint['state_dict'].items()])
for name, param in model.FeatureExtraction.state_dict().items():
model.FeatureExtraction.state_dict()[name].copy_(checkpoint['state_dict']['FeatureExtraction.' + name])
for name, param in model.FeatureRegression.state_dict().items():
args = parser.parse_args()
use_cuda = torch.cuda.is_available()
do_aff = not args.model_aff == ''
do_tps = not args.model_tps == ''
# Download dataset if needed
download_PF_willow('datasets/')
# Create model
print('Creating CNN model...')
if do_aff:
model_aff = CNNGeometric(
use_cuda=use_cuda,
output_dim=6,
feature_extraction_cnn=args.feature_extraction_cnn)
if do_tps:
model_tps = CNNGeometric(
use_cuda=use_cuda,
output_dim=18,
feature_extraction_cnn=args.feature_extraction_cnn)
# Load trained weights
print('Loading trained model weights...')
if do_aff:
checkpoint = torch.load(args.model_aff,
map_location=lambda storage, loc: storage)
checkpoint['state_dict'] = OrderedDict([
(k.replace('vgg', 'model'), v)
for k, v in checkpoint['state_dict'].items()
default='Pairs CSV file',
help='Path to PF dataset')
args = parser.parse_args()
use_cuda = torch.cuda.is_available()
do_aff = not args.model_aff == ''
do_tps = False #not args.model_tps==''
show_diffmap = True
dataset_path = args.path
dataset_pairs_file = args.pairs
# Create model
print('Creating CNN model...')
if do_aff:
model_aff = CNNGeometric(use_cuda=use_cuda, geometric_model='affine')
if do_tps:
model_tps = CNNGeometric(use_cuda=use_cuda, geometric_model='tps')
# Load trained weights
print('Loading trained model weights...')
if do_aff:
checkpoint = torch.load(args.model_aff,
map_location=lambda storage, loc: storage)
model_aff.load_state_dict(checkpoint['state_dict'])
if do_tps:
checkpoint = torch.load(args.model_tps,
map_location=lambda storage, loc: storage)
model_tps.load_state_dict(checkpoint['state_dict'])
# Dataset and dataloader
torch.cuda.manual_seed(args.seed)
# Download dataset if needed and set paths
if args.training_dataset == 'pascal':
if args.training_image_path == '':
download_pascal('datasets/pascal-voc11/')
args.training_image_path = 'datasets/pascal-voc11/'
if args.training_tnf_csv == '' and args.geometric_model=='affine':
args.training_tnf_csv = 'training_data/pascal-synth-aff'
elif args.training_tnf_csv == '' and args.geometric_model=='tps':
args.training_tnf_csv = 'training_data/pascal-synth-tps'
# CNN model and loss
print('Creating CNN model...')
model = CNNGeometric(use_cuda=use_cuda,geometric_model=args.geometric_model,arch=args.arch,featext_weights=args.fw)
if args.use_mse_loss:
print('Using MSE loss...')
loss = nn.MSELoss()
else:
print('Using grid loss...')
loss = TransformedGridLoss(use_cuda=use_cuda,geometric_model=args.geometric_model)
if args.geometric_model == 'pose':
dataset = PoseDataset(
geometric_model=args.geometric_model,
csv_file=os.path.join(args.training_tnf_csv, 'train.csv'),
training_image_path=args.training_image_path,
output_size=(240, 240),
