class SuperMatcher:
resize = [640, 480]
superglue = "outdoor"
max_keypoints = -1
nms_radius = 4
keypoint_threshold = 0.1
sinkhorn_iterations = 20
match_threshold = 0.4
device = "cuda" if torch.cuda.is_available() else "cpu"
show_keypoints = False
config = {
"superpoint": {
"nms_radius": nms_radius,
"keypoint_threshold": keypoint_threshold,
"max_keypoints": max_keypoints,
},
"superglue": {
"weights": superglue,
"sinkhorn_iterations": sinkhorn_iterations,
"match_threshold": match_threshold,
},
}
def __init__(self):
self.matching = Matching(self.config).eval().to(self.device)
self.keys = ["keypoints", "scores", "descriptors"]
### Anchor Frame and associated data
self.anchor_frame_tensor = None
self.anchor_data = None
self.anchor_frame = None
self.anchor_image_id = None
def set_anchor(self, frame):
# Frame will be divided by 255
self.anchor_frame_tensor = frame2tensor(frame, self.device)
self.anchor_data = self.matching.superpoint(
{"image": self.anchor_frame_tensor})
self.anchor_data = {k + "0": self.anchor_data[k] for k in self.keys}
self.anchor_data["image0"] = self.anchor_frame_tensor
self.anchor_frame = frame
self.anchor_image_id = 0
def process(self, frame):
if self.anchor_frame_tensor is None:
print("Please set anchor frame first...")
return None
frame_tensor = frame2tensor(frame, self.device)
pred = self.matching({**self.anchor_data, "image1": frame_tensor})
kpts0 = self.anchor_data["keypoints0"][0].cpu().numpy()
kpts1 = pred["keypoints1"][0].cpu().numpy()
matches = pred["matches0"][0].cpu().numpy()
confidence = pred["matching_scores0"][0].cpu().numpy()
valid = matches > -1
mkpts0 = kpts0[valid]
mkpts1 = kpts1[matches[valid]]
color = cm.jet(confidence[valid])
text = [
"SuperGlue",
"Keypoints: {}:{}".format(len(kpts0), len(kpts1)),
"Matches: {}".format(len(mkpts0)),
]
k_thresh = self.matching.superpoint.config["keypoint_threshold"]
m_thresh = self.matching.superglue.config["match_threshold"]
small_text = [
"Keypoint Threshold: {:.4f}".format(k_thresh),
"Match Threshold: {:.2f}".format(m_thresh),
]
out = make_matching_plot_fast(
self.anchor_frame,
frame,
kpts0,
kpts1,
mkpts0,
mkpts1,
color,
text,
path=None,
show_keypoints=self.show_keypoints,
small_text=small_text,
)
return out, mkpts0, mkpts1
'superglue': {
'weights': opt.superglue,
'sinkhorn_iterations': opt.sinkhorn_iterations,
'match_threshold': opt.match_threshold,
}
}
matching = Matching(config).eval().to(device)
keys = ['keypoints', 'scores', 'descriptors']
vs = VideoStreamer(opt.input, opt.resize, opt.skip, opt.image_glob,
opt.max_length)
frame, ret = vs.next_frame()
assert ret, 'Error when reading the first frame (try different --input?)'
frame_tensor = frame2tensor(frame, device)
last_data = matching.superpoint({'image': frame_tensor})
last_data = {k + '0': last_data[k] for k in keys}
last_data['image0'] = frame_tensor
last_frame = frame
last_image_id = 0
if opt.output_dir is not None:
print('==> Will write outputs to {}'.format(opt.output_dir))
Path(opt.output_dir).mkdir(exist_ok=True)
# Create a window to display the demo.
if not opt.no_display:
cv2.namedWindow('SuperGlue matches', cv2.WINDOW_NORMAL)
cv2.resizeWindow('SuperGlue matches', (640 * 2, 480))
else:
print('Skipping visualization, will not show a GUI.')
K[1, 2] = h / 2
def backproject(depth, pts_2d, T):
d = depth[tuple(np.round(pts_2d).astype(int).T)[::-1]]
valid = d > 0
pts_3d = (to_homogeneous(pts_2d) @ np.linalg.inv(K).T) * d[:, None]
pts_3d = (pts_3d @ T[:3, :3].T) + T[:3, 3][None]
pts_3d[~valid] = None
return pts_3d
all_pts_3d = []
all_pts_color = []
viz_poses = []
is_ref_frame = []
pred = matching.superpoint({'image': frame_tensor})
pred = {k + '1': v for k, v in pred.items()}
def update_reference(i):
last_data = {k + '0': pred[k + '1'] for k in keys}
last_data['image0'] = frame_tensor
last_frame = frame
last_id = i
num = len(last_data['keypoints0'][0])
colors = cm.hsv(np.random.rand(num))
T = poses[associate_ts(frames[i][0], poses)]
pts_3d = backproject(depth, last_data['keypoints0'][0].cpu().numpy(),
T)
write_ply(Path(opt.output_dir, f'model_{i}.ply'), pts_3d, colors)
return last_data, last_frame, last_id, colors, pts_3d, T
