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