def main(args):
use_cuda = torch.cuda.is_available()
# Download dataset if needed
download_PF_willow('datasets/')
# Create model
print('Creating CNN model...')
model_aff = load_torch_model(args, use_cuda)
csv_path_list = glob(os.path.join(args.pf_path, '*.csv'))
if len(csv_path_list) > 1:
print("!!!!WARNING!!!! multiple csv files found, using first in glob order")
elif not len(csv_path_list):
raise FileNotFoundError("No csvs where found in the specified path!!!")
if not os.path.exists(args.out_dir):
os.mkdir(args.out_dir)
csv_path = csv_path_list[0]
# Dataset and dataloader
dataset = CoupledDataset(csv_file=csv_path,
training_image_path=args.pf_images_dir,
transform=NormalizeImageDict(['image_a', 'image_b'])) # template=args.template_path)
batch_size = 1
dataloader = DataLoader(dataset, batch_size=batch_size,
shuffle=True, num_workers=4)
# Set Tnf pair generation func
pair_generation_tnf = CoupledPairTnf(use_cuda=use_cuda)
for i, batch in enumerate(dataloader):
tnf_batch = pair_generation_tnf(batch)
theta = model_aff(tnf_batch)
test_and_save_alignment(batch, i, theta, args.out_dir)
default='resnet101',
help='Feature extraction architecture: vgg/resnet101')
parser.add_argument('--pf-path',
type=str,
default='datasets/PF-dataset',
help='Path to PF dataset')
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