(k.replace('vgg', 'model'), v)
for k, v in checkpoint['state_dict'].items()
])
model_aff.load_state_dict(checkpoint['state_dict'])
if do_tps:
checkpoint = torch.load(args.model_tps,
map_location=lambda storage, loc: storage)
checkpoint['state_dict'] = OrderedDict([
(k.replace('vgg', 'model'), v)
for k, v in checkpoint['state_dict'].items()
])
model_tps.load_state_dict(checkpoint['state_dict'])
# Dataset and dataloader
dataset = PFDataset(csv_file=os.path.join(args.pf_path, 'test_pairs_pf.csv'),
training_image_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)
# 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)
type=str,
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()
# Download dataset if needed
download_PF_willow('datasets/')
# Dataset and dataloader
dataset = PFDataset(csv_file='datasets/PF-dataset/test_pairs_pf.csv',
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)
dloader_train = torch.utils.data.DataLoader(Subset(ds_train, train_index),
batch_size=args.batch_size,
shuffle=True,
num_workers=args.n_workers,
pin_memory=True)
dloader_valid = torch.utils.data.DataLoader(
Subset(ds_train, valid_index),
batch_size=args.batch_size_valid,
shuffle=False,
num_workers=args.n_workers,
pin_memory=True)
# eval 用データ
ds_eval = PFDataset(args,
args.dataset_eval_dir,
image_height=args.image_height,
image_width=args.image_width,
debug=args.debug)
dloader_eval = torch.utils.data.DataLoader(ds_eval,
batch_size=1,
shuffle=False,
num_workers=args.n_workers,
pin_memory=True)
#================================
# モデルの構造を定義する。
#================================
# GeometricMatchingCNN モデル
if (args.geometric_model == "affine"):
model_G = GeometricMatchingCNN(n_out_channels=6).to(device)
elif (args.geometric_model == "tps"):
# NAN 値の検出
if (args.detect_nan):
torch.autograd.set_detect_anomaly(True)
# tensorboard 出力
board_test = SummaryWriter(
log_dir=os.path.join(args.tensorboard_dir, args.exper_name + "_test"))
#================================
# データセットの読み込み
#================================
# 学習用データセットとテスト用データセットの設定
ds_test = PFDataset(args,
args.dataset_dir,
image_height=args.image_height,
image_width=args.image_width,
debug=args.debug)
dloader_test = torch.utils.data.DataLoader(ds_test,
batch_size=args.batch_size_test,
shuffle=False,
num_workers=args.n_workers,
pin_memory=True)
#================================
# モデルの構造を定義する。
#================================
#------------------------
# stage 1
#------------------------
# GeometricMatchingCNN モデル