def __init__(
self,
config,
data_loader,
val_data_loader=None,
):
self.device = torch.device(
'cuda' if torch.cuda.is_available() else 'cpu')
self.config = config
self.max_epoch = config.max_epoch
self.save_freq = config.save_freq_epoch
self.train_max_iter = config.train_max_iter
self.val_max_iter = config.val_max_iter
self.val_epoch_freq = config.val_epoch_freq
self.best_val_metric = config.best_val_metric
self.best_val_epoch = -np.inf
self.best_val = -np.inf
self.initialize_model()
if config.use_gpu and not torch.cuda.is_available():
logging.warning(
'Warning: There\'s no CUDA support on this machine, '
'training is performed on CPU.')
raise ValueError('GPU not available, but cuda flag set')
self.start_epoch = 1
self.checkpoint_dir = config.out_dir
ensure_dir(self.checkpoint_dir)
json.dump(config,
open(os.path.join(self.checkpoint_dir, 'config.json'), 'w'),
indent=4,
sort_keys=False)
self.iter_size = config.iter_size
self.batch_size = config.batch_size
self.data_loader = data_loader
self.val_data_loader = val_data_loader
self.test_valid = True if self.val_data_loader is not None else False
self.log_step = int(np.sqrt(self.config.batch_size))
self.writer = SummaryWriter(logdir=config.out_dir)
self.initialize_optimiser_and_scheduler()
self.resume()
parser.add_argument('--with_cuda', action='store_false')
parser.add_argument('--num_rand_keypoints',
type=int,
default=5000,
help='Number of random keypoints for each scene')
args = parser.parse_args()
device = torch.device('cuda' if args.with_cuda else 'cpu')
if args.extract_features:
assert args.model is not None
assert args.source is not None
assert args.target is not None
ensure_dir(args.target)
checkpoint = torch.load(args.model)
config = checkpoint['config']
num_feats = 1
Model = load_model(config.model)
model = Model(num_feats,
config.model_n_out,
bn_momentum=0.05,
normalize_feature=config.normalize_feature,
conv1_kernel_size=config.conv1_kernel_size,
D=3)
model.load_state_dict(checkpoint['state_dict'])
model.eval()
model = model.to(device)
def __init__(
self,
config,
data_loader,
val_data_loader=None,
):
num_feats = 1 # occupancy only for 3D Match dataset. For ScanNet, use RGB 3 channels.
# Model initialization
Model = load_model(config.model)
model = Model(
num_feats,
config.model_n_out,
bn_momentum=config.bn_momentum,
normalize_feature=config.normalize_feature,
conv1_kernel_size=config.conv1_kernel_size,
D=3)
if config.weights:
checkpoint = torch.load(config.weights)
model.load_state_dict(checkpoint['state_dict'])
logging.info(model)
self.config = config
self.model = model
self.max_epoch = config.max_epoch
self.save_freq = config.save_freq_epoch
self.val_max_iter = config.val_max_iter
self.val_epoch_freq = config.val_epoch_freq
self.best_val_metric = config.best_val_metric
self.best_val_epoch = -np.inf
self.best_val = -np.inf
if config.use_gpu and not torch.cuda.is_available():
logging.warning('Warning: There\'s no CUDA support on this machine, '
'training is performed on CPU.')
raise ValueError('GPU not available, but cuda flag set')
self.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
self.optimizer = getattr(optim, config.optimizer)(
model.parameters(),
lr=config.lr,
momentum=config.momentum,
weight_decay=config.weight_decay)
self.scheduler = optim.lr_scheduler.ExponentialLR(self.optimizer, config.exp_gamma)
self.start_epoch = 1
self.checkpoint_dir = config.out_dir
ensure_dir(self.checkpoint_dir)
json.dump(
config,
open(os.path.join(self.checkpoint_dir, 'config.json'), 'w'),
indent=4,
sort_keys=False)
self.iter_size = config.iter_size
self.batch_size = data_loader.batch_size
self.data_loader = data_loader
self.val_data_loader = val_data_loader
self.test_valid = True if self.val_data_loader is not None else False
self.log_step = int(np.sqrt(self.config.batch_size))
self.model = self.model.to(self.device)
self.writer = SummaryWriter(logdir=config.out_dir)
if config.resume is not None:
if osp.isfile(config.resume):
logging.info("=> loading checkpoint '{}'".format(config.resume))
state = torch.load(config.resume)
self.start_epoch = state['epoch']
model.load_state_dict(state['state_dict'])
self.scheduler.load_state_dict(state['scheduler'])
self.optimizer.load_state_dict(state['optimizer'])
if 'best_val' in state.keys():
self.best_val = state['best_val']
self.best_val_epoch = state['best_val_epoch']
self.best_val_metric = state['best_val_metric']
else:
raise ValueError(f"=> no checkpoint found at '{config.resume}'")
def __init__(self, config, data_loader, val_data_loader=None):
# occupancy only for 3D Match dataset. For ScanNet, use RGB 3 channels.
num_feats = 3 if config.use_xyz_feature else 1
# Feature model initialization
if config.use_gpu and not torch.cuda.is_available():
logging.warning(
'Warning: There\'s no CUDA support on this machine, '
'training is performed on CPU.')
raise ValueError('GPU not available, but cuda flag set')
self.device = torch.device(
'cuda' if torch.cuda.is_available() else 'cpu')
self.config = config
# Training config
self.max_epoch = config.max_epoch
self.start_epoch = 1
self.checkpoint_dir = config.out_dir
self.data_loader = data_loader
self.train_data_loader_iter = self.data_loader.__iter__()
self.iter_size = config.iter_size
self.batch_size = data_loader.batch_size
# Validation config
self.val_max_iter = config.val_max_iter
self.val_epoch_freq = config.val_epoch_freq
self.best_val_metric = config.best_val_metric
self.best_val_epoch = -np.inf
self.best_val = -np.inf
self.val_data_loader = val_data_loader
self.test_valid = True if self.val_data_loader is not None else False
# Logging
self.log_step = int(np.sqrt(self.config.batch_size))
self.writer = SummaryWriter(config.out_dir)
# Model
FeatModel = load_model(config.feat_model)
InlierModel = load_model(config.inlier_model)
num_feats = 6 if self.config.inlier_feature_type == 'coords' else 1
self.feat_model = FeatModel(
num_feats,
config.feat_model_n_out,
bn_momentum=config.bn_momentum,
conv1_kernel_size=config.feat_conv1_kernel_size,
normalize_feature=config.normalize_feature).to(self.device)
logging.info(self.feat_model)
self.inlier_model = InlierModel(
num_feats,
1,
bn_momentum=config.bn_momentum,
conv1_kernel_size=config.inlier_conv1_kernel_size,
normalize_feature=False,
D=6).to(self.device)
logging.info(self.inlier_model)
# Loss and optimizer
self.clip_weight_thresh = self.config.clip_weight_thresh
if self.config.use_balanced_loss:
self.crit = BalancedLoss()
else:
self.crit = UnbalancedLoss()
self.optimizer = getattr(optim, config.optimizer)(
self.inlier_model.parameters(),
lr=config.lr,
momentum=config.momentum,
weight_decay=config.weight_decay)
self.scheduler = optim.lr_scheduler.ExponentialLR(
self.optimizer, config.exp_gamma)
# Output preparation
ensure_dir(self.checkpoint_dir)
json.dump(config,
open(os.path.join(self.checkpoint_dir, 'config.json'), 'w'),
indent=4,
sort_keys=False)
self._load_weights(config)
def visualize_image_correspondence(img0,
img1,
F0,
F1,
filename,
mode='gpu-all',
config=None,
visualize=True):
use_stability_test = True
use_cyclic_test = False
keypoint = 'sift'
if keypoint == 'sift':
sift = cv2.xfeatures2d.SIFT_create(
0,
9,
0.01, # Smaller more keypoints, default 0.04
100 # larger more keypoints, default 10
)
kp0 = sift.detect(img0, None)
kp1 = sift.detect(img1, None)
xy_kp0 = np.floor(np.array([k.pt for k in kp0]).T)
xy_kp1 = np.floor(np.array([k.pt for k in kp1]).T)
x0, y0 = xy_kp0[0], xy_kp0[1]
x1, y1 = xy_kp1[0], xy_kp1[1]
elif keypoint == 'all':
x0, y0 = None, None
x1, y1 = None, None
H0, W0 = img0.shape
H1, W1 = img1.shape
if mode == 'cpu-keypoints':
matches1 = util_2d.feature_match(F0[:, y0, x0].t().cpu().numpy(),
F1[:, y1, x1].t().cpu().numpy(),
ratio_test=True,
ratio=0.95)
# Convert the index to coordinate: BxCxHxW
x0 = x0[matches1[:, 0]]
y0 = y0[matches1[:, 0]]
xs1 = x1[matches1[:, 1]]
ys1 = y1[matches1[:, 1]]
# Test reciprocity
nn_inds0 = find_nn_gpu(F1[:, ys1, xs1],
F0[:, y0, x0],
nn_max_n=config.nn_max_n,
transposed=True)
# Convert the index to coordinate: BxCxHxW
xs0 = x0[nn_inds0.numpy()]
ys0 = y0[nn_inds0.numpy()]
dist_sq_nn = (x0 - xs0)**2 + (y0 - ys0)**2
mask = dist_sq_nn < (config.ucn_inlier_threshold_pixel**2)
elif mode == 'gpu-keypoints':
nn_inds1 = find_nn_gpu(F0[:, y0, x0],
F1[:, y1, x1],
nn_max_n=config.nn_max_n,
transposed=True).numpy()
# Convert the index to coordinate: BxCxHxW
xs1 = x1[nn_inds1]
ys1 = y1[nn_inds1]
if use_stability_test:
# Stability test: check stable under perturbation
noisex = 2 * (np.random.rand(len(xs1)) < 0.5) - 1
noisey = 2 * (np.random.rand(len(ys1)) < 0.5) - 1
xs1n = np.clip(xs1 + noisex, 0, W1 - 1)
ys1n = np.clip(ys1 + noisey, 0, H1 - 1)
else:
xs1n = xs1
ys1n = ys1
# Test reciprocity
nn_inds0 = find_nn_gpu(F1[:, ys1n, xs1n],
F0[:, y0, x0],
nn_max_n=config.nn_max_n,
transposed=True).numpy()
# Convert the index to coordinate: BxCxHxW
xs0 = x0[nn_inds0]
ys0 = y0[nn_inds0]
dist_sq_nn = (x0 - xs0)**2 + (y0 - ys0)**2
mask = dist_sq_nn < (config.ucn_inlier_threshold_pixel**2)
elif mode == 'gpu-all':
nn_inds1 = find_nn_faiss(
F0[:, y0, x0],
F1.view(F1.shape[0], -1),
)
# Convert the index to coordinate: BxCxHxW
xs1 = nn_inds1 % W1
ys1 = nn_inds1 // W1
if use_stability_test:
# Stability test: check stable under perturbation
noisex = 2 * (np.random.rand(len(xs1)) < 0.5) - 1
noisey = 2 * (np.random.rand(len(ys1)) < 0.5) - 1
xs1n = np.clip(xs1 + noisex, 0, W1 - 1)
ys1n = np.clip(ys1 + noisey, 0, H1 - 1)
else:
xs1n = xs1
ys1n = ys1
if use_cyclic_test:
# Test reciprocity
nn_inds0 = find_nn_faiss(
F1[:, ys1n, xs1n],
F0.view(F0.shape[0], -1),
)
# Convert the index to coordinate: BxCxHxW
xs0 = (nn_inds0 % W0)
ys0 = (nn_inds0 // W0)
# Test cyclic consistency
dist_sq_nn = (x0 - xs0)**2 + (y0 - ys0)**2
mask = dist_sq_nn < (config.ucn_inlier_threshold_pixel**2)
else:
xs0 = x0
ys0 = y0
mask = np.ones(len(x0)).astype(bool)
elif mode == 'gpu-all-all':
nn_inds1 = find_nn_faiss(
F0.view(F0.shape[0], -1),
F1.view(F1.shape[0], -1),
)
inds0 = np.arange(len(nn_inds1))
x0 = inds0 % W0
y0 = inds0 // W0
xs1 = nn_inds1 % W1
ys1 = nn_inds1 // W1
if use_stability_test:
# Stability test: check stable under perturbation
noisex = 2 * (np.random.rand(len(xs1)) < 0.5) - 1
noisey = 2 * (np.random.rand(len(ys1)) < 0.5) - 1
xs1n = np.clip(xs1 + noisex, 0, W1 - 1)
ys1n = np.clip(ys1 + noisey, 0, H1 - 1)
else:
xs1n = xs1
ys1n = ys1
# Test reciprocity
nn_inds0 = find_nn_faiss(
F1[:, ys1n, xs1n],
F0.view(F0.shape[0], -1),
)
# Convert the index to coordinate: BxCxHxW
xs0 = nn_inds0 % W0
ys0 = nn_inds0 // W0
# Filter out the points that fail the cycle consistency
dist_sq_nn = (x0 - xs0)**2 + (y0 - ys0)**2
mask = dist_sq_nn < (config.ucn_inlier_threshold_pixel**2)
if visualize:
color = x0[mask] + y0[mask] * W0
plt.clf()
fig, (ax0, ax1) = plt.subplots(nrows=1, ncols=2)
fig = plt.gcf()
fig.set_size_inches(9, 6)
ax0.imshow(img0 * 0.5, vmin=0, vmax=255, cmap='gray')
ax0.scatter(x=x0[mask], y=y0[mask], c=color, s=2, cmap="jet")
ax0.axis('off')
ax1.imshow(img1 * 0.5, vmin=0, vmax=255, cmap='gray')
ax1.scatter(x=xs1[mask], y=ys1[mask], c=color, s=2, cmap="jet")
ax1.axis('off')
fig.tight_layout()
ensure_dir('./ucn_outputs')
plt.savefig(f"./ucn_outputs/{filename:03d}.png", dpi=300)
else:
return x0[mask], y0[mask], xs1[mask], ys1[mask]
