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() # print("anchor",self.anchor_data.keys()) # print(self.anchor_data['keypoints0'][0].shape) # print(self.anchor_data['scores0'][0].shape) # print(self.anchor_data['descriptors0'][0].shape) # print("pred",pred.keys()) 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', 'Keypoints: {}:{}'.format(len(kpts0), len(kpts1)), 'Matches: {}'.format(len(mkpts0)) ] k_thresh = matching.superpoint.config['keypoint_threshold'] m_thresh = matching.superglue.config['match_threshold'] small_text = [ 'Keypoint Threshold: {:.4f}'.format(k_thresh), 'Match Threshold: {:.2f}'.format(m_thresh), 'Image Pair: {:06}:{:06}'.format(stem0, stem1), ] out = make_matching_plot_fast(last_frame, frame, kpts0, kpts1, mkpts0, mkpts1, color, text, path=None, show_keypoints=opt.show_keypoints, small_text=small_text) if not opt.no_display: cv2.imshow('SuperGlue matches', out) key = chr(cv2.waitKey(1) & 0xFF) if key == 'q': vs.cleanup() print('Exiting (via q) demo_superglue.py') break elif key == 'n': # set the current frame as anchor last_data = {k + '0': pred[k + '1'] for k in keys}
def getKpts(matching, image0_path, image1_path, n=-1): torch.set_grad_enabled(False) keys = ['keypoints', 'scores', 'descriptors'] # make image0 first, image 1 second if image0_path[-5] == '2': image0_path, image1_path = image1_path, image0_path print(image0_path) frame0 = cv2.imread(image0_path, 0) print(type(frame0), type(mask0)) frame0 = (torch.Tensor(frame0)).to(device) frame0 = torch.mul(frame0, mask0) # apply mask frame_tensor0 = (frame0 / 255.).float()[None, None] last_data = matching.superpoint({'image': frame_tensor0}) last_data = {k + '0': last_data[k] for k in keys} last_data['image0'] = frame_tensor0 frame1 = cv2.imread(image1_path, 0) frame1 = (torch.Tensor(frame1)).to(device) frame1 = torch.mul(frame1, mask1) # apply mask frame1 = torch.rot90(frame1) # rotate image to make images 0,1 similar frame_tensor1 = (frame1 / 255.).float()[None, None] pred = matching({**last_data, 'image1': frame_tensor1}) kpts0 = last_data['keypoints0'][0] kpts1 = pred['keypoints1'][0] matches = pred['matches0'][0] confidence = pred['matching_scores0'][0] # valid_data valid = matches > -1 valid_kpts0 = kpts0[valid].to(device) valid_kpts1 = kpts1[matches[valid]].to(device) valid_conf = confidence[valid].to(device) # print('ind shape before deletion: ', valid_conf.shape) # img_center = torch.Tensor([pic_half_size,pic_half_size]) if len(valid_conf.size()) > 0: # delete out-of-circle points center = torch.Tensor([pic_half_size, pic_half_size]).to(device) len0 = torch.norm(valid_kpts0 - center, dim=1) len1 = torch.norm(valid_kpts1 - center, dim=1) indices0 = (len0 < pic_half_size).nonzero() indices1 = (len1 < pic_half_size).nonzero() indices = torch.unique(torch.cat((indices0, indices1))) valid_kpts0 = torch.squeeze(valid_kpts0[indices]) valid_kpts1 = torch.squeeze(valid_kpts1[indices]) valid_conf = torch.squeeze(valid_conf[indices]) if len(valid_conf.size()) > 0: frame0_brightness = torch.zeros(valid_conf.shape[0]) frame1_brightness = torch.zeros(valid_conf.shape[0]) for k in range(valid_conf.shape[0]): frame0_brightness[k] = frame0[int(valid_kpts0[k, 0]), int(valid_kpts0[k, 1])] / 255 frame1_brightness[k] = frame1[int(valid_kpts1[k, 0]), int(valid_kpts1[k, 1])] / 255 if n != -1: if valid_conf.shape[0] < n: valid_kpts0 = torch.zeros((n, 2)) valid_kpts1 = torch.zeros((n, 2)) valid_conf = torch.zeros(n) frame0_brightness = torch.zeros(n) frame1_brightness = torch.zeros(n) not_skip = 0 # print('here0') else: valid_conf, arr1inds = torch.sort(valid_conf, descending=True) valid_kpts0 = valid_kpts0[arr1inds, :] valid_kpts1 = valid_kpts1[arr1inds, :] frame0_brightness = frame0_brightness[arr1inds] frame1_brightness = frame1_brightness[arr1inds] valid_kpts0 = valid_kpts0[0:n] valid_kpts1 = valid_kpts1[0:n] valid_conf = valid_conf[0:n] frame0_brightness = frame0_brightness[0:n] frame1_brightness = frame1_brightness[0:n] not_skip = 1 # print('here1') # if n==-1, keep all points else: not_skip = 1 # print('here2') else: valid_kpts0 = torch.zeros((1, 2)) valid_kpts1 = torch.zeros((1, 2)) valid_conf = torch.zeros(2) frame0_brightness = torch.zeros(2) frame1_brightness = torch.zeros(2) not_skip = 0 # print('here3') # save image if not_skip: valid_kpts0_pic = valid_kpts0.cpu().numpy() valid_kpts1_pic = valid_kpts1.cpu().numpy() valid_conf_pic = valid_conf.cpu().numpy() frame0_pic = frame0.cpu().numpy() frame1_pic = frame1.cpu().numpy() kpts0_pic = kpts0.cpu().numpy() kpts1_pic = kpts1.cpu().numpy() color = cm.jet(valid_conf_pic) text = [ 'SuperGlue', 'Keypoints: {}:{}'.format(len(kpts0_pic), len(kpts1_pic)), 'Matches: {}'.format(len(valid_kpts0_pic)) ] k_thresh = matching.superpoint.config['keypoint_threshold'] m_thresh = matching.superglue.config['match_threshold'] small_text = [ 'Keypoint Threshold: {:.4f}'.format(k_thresh), 'Match Threshold: {:.2f}'.format(m_thresh), ] out = make_matching_plot_fast( frame0_pic, frame1_pic, kpts0_pic, kpts1_pic, valid_kpts0_pic, valid_kpts1_pic, color, text, path=None, show_keypoints=True, small_text=small_text) else: out = np.zeros((1920, 3850, 3)) return valid_kpts0, valid_kpts1, valid_conf, frame0_brightness, frame1_brightness, not_skip, out
confidence = pred['matching_scores0'][0].cpu().numpy() timer.update('forward') 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)) ] out = make_matching_plot_fast(last_frame, frame, kpts0, kpts1, mkpts0, mkpts1, color, text, add_keypoints=opt.show_keypoints) if not opt.no_display: ds = opt.display_scale out = cv2.resize(out, (int(out.shape[1] * ds), int(out.shape[0] * ds))) cv2.imshow('SuperGlue matches', out) key = chr(cv2.waitKey(1) & 0xFF) if key == 'q': print('Exiting (via q) demo_superglue.py') break elif key == 'n': # set the current frame as reference
] k_thresh = matching.superpoint.config['keypoint_threshold'] m_thresh = matching.superglue.config['match_threshold'] small_text = [ 'Keypoint Threshold: {:.4f}'.format(k_thresh), 'Match Threshold: {:.2f}'.format(m_thresh), 'Image Pair: {:06}:{:06}'.format(stem0, stem1), ] if opt.regionnetvlad: out = make_matching_plot_fast(cv2.cvtColor(last_frame, cv2.COLOR_RGB2GRAY), cv2.cvtColor(frame, cv2.COLOR_RGB2GRAY), kpts0, kpts1, mkpts0, mkpts1, color, text, path=None, show_keypoints=opt.show_keypoints, small_text=small_text) else: out = make_matching_plot_fast(last_frame, frame, kpts0, kpts1, mkpts0, mkpts1, color, text,