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 __init__(self, args, data_loader, config): self.matching = matcher.Matching( config['matcher']).eval().to(device='cuda') self.prev_image = data_loader.get_image() self.cur_image = data_loader.get_image() self.prev_image['left_cuda'] = utils.frame2tensor( self.prev_image['left_gray'], 'cuda') self.cur_image['left_cuda'] = utils.frame2tensor( self.cur_image['left_gray'], 'cuda') self.matches = self.matching({ 'image0': self.prev_image['left_cuda'], 'image1': self.cur_image['left_cuda'] }) self.inlier_threshold = config['inlier_threshold'] self.prev_pose = data_loader.get_true_pose() self.pose = {'true_pose': [self.prev_pose], 'vo_pose': [[0, 0, 0]]} self.R, self.T = np.identity(3), np.zeros((3, 1)) self.optical_flow_model = torch.nn.DataParallel(RAFT(args)) self.optical_flow_model.load_state_dict( torch.load(args.optical_flow_model)) model = self.optical_flow_model.module model.to('cuda').eval() y, x = np.mgrid[0:self.prev_image['left_gray'].shape[0]:1, 0:self.prev_image['left_gray'].shape[1]:1].reshape( 2, -1).astype(int) self.dynamic_model = np.array([x**2, y**2, x * y, x, y, x * 0 + 1]).T config['motion_detection']['x'], config['motion_detection']['y'] = x, y config['motion_detection']['height'], config['motion_detection'][ 'width'] = self.prev_image['left_gray'].shape[0], self.prev_image[ 'left_gray'].shape[1] self.prev_image.update(**self.cur_image)
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
def get_matches(self, image): self.cur_image['left_cuda'] = utils.frame2tensor(image, 'cuda') self.matches = self.matching({ 'image0': self.prev_image['left_cuda'], 'image1': self.cur_image['left_cuda'] }) confidence = self.matches['matching_scores0'][0].detach().cpu().numpy() valid_match = self.matches['matches0'][0].cpu().numpy() > -1 valid_confidence = confidence > 0.85 valid = np.logical_and(valid_match, valid_confidence) self.mkpts0 = self.matches['keypoints0'][0].detach().cpu().numpy( )[valid] self.mkpts1 = self.matches['keypoints1'][0].detach().cpu().numpy()[ self.matches['matches0'][0].cpu().numpy()[valid]] return {'ref_keypoints': self.mkpts0, 'cur_keypoints': self.mkpts1}
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 keypoints = self.anchor_data['keypoints0'][0].cpu().numpy() scores = self.anchor_data['scores0'][0].cpu().numpy() descriptors = self.anchor_data['descriptors0'][0].cpu().numpy() # N x 3 pts = np.hstack((keypoints, scores.reshape(-1, 1))) # N x 256 (256 is descriptor dimension) descriptors = descriptors.T return pts, descriptors
}, '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:
def __getitem__(self, idx): before_getitem_time = time.time() file_name = self.files[idx] image = cv2.imread(file_name, cv2.IMREAD_GRAYSCALE) SuperPoint = self.superpoint width, height = image.shape[:2] corners = np.array([[0, 0], [0, height], [width, 0], [width, height]], dtype=np.float32) warp = np.random.randint(-224, 224, size=(4, 2)).astype(np.float32) # get the corresponding warped image M = cv2.getPerspectiveTransform(corners, corners + warp) warped = cv2.warpPerspective( src=image, M=M, dsize=(image.shape[1], image.shape[0])) # return an image type # extract keypoints of the image pair using SuperPoint # before_superpoint_time = time.time() kp1, descs1 = self.parse_superpoint_result( SuperPoint({'image': frame2tensor(image)})) kp2, descs2 = self.parse_superpoint_result( SuperPoint({'image': frame2tensor(warped)})) print(kp1.shape, descs1.shape) # print('SuperPoint time:', time.time() - before_superpoint_time) # limit the number of keypoints kp1_num = min(self.nfeatures, len(kp1)) kp2_num = min(self.nfeatures, len(kp2)) kp1 = kp1[:kp1_num] kp2 = kp2[:kp2_num] kp1_np = np.array([(kp[0], kp[1]) for kp in kp1]) kp2_np = np.array([(kp[0], kp[1]) for kp in kp2]) # skip this image pair if no keypoints detected in image if len(kp1) < 1 or len(kp2) < 1: return { 'keypoints0': torch.zeros([0, 0, 2], dtype=torch.double), 'keypoints1': torch.zeros([0, 0, 2], dtype=torch.double), 'descriptors0': torch.zeros([0, 2], dtype=torch.double), 'descriptors1': torch.zeros([0, 2], dtype=torch.double), 'image0': image, 'image1': warped, 'file_name': file_name } # confidence of each key point scores1_np = np.array([kp[2] for kp in kp1]) scores2_np = np.array([kp[2] for kp in kp2]) kp1_np = kp1_np[:kp1_num, :] kp2_np = kp2_np[:kp2_num, :] descs1 = descs1[:kp1_num, :] descs2 = descs2[:kp2_num, :] # obtain the matching matrix of the image pair # before_match_time = time.time() matched = self.matcher.match(descs1, descs2) # # print('match time:', time.time() - before_match_time) kp1_projected = cv2.perspectiveTransform(kp1_np.reshape((1, -1, 2)), M)[0, :, :] dists = cdist(kp1_projected, kp2_np) min1 = np.argmin(dists, axis=0) min2 = np.argmin(dists, axis=1) min1v = np.min(dists, axis=1) min1f = min2[min1v < 3] xx = np.where(min2[min1] == np.arange(min1.shape[0]))[0] matches = np.intersect1d(min1f, xx) print(matches) missing1 = np.setdiff1d(np.arange(kp1_np.shape[0]), min1[matches]) missing2 = np.setdiff1d(np.arange(kp2_np.shape[0]), matches) MN = np.concatenate( [min1[matches][np.newaxis, :], matches[np.newaxis, :]]) MN2 = np.concatenate([ missing1[np.newaxis, :], (len(kp2)) * np.ones( (1, len(missing1)), dtype=np.int64) ]) MN3 = np.concatenate([(len(kp1)) * np.ones( (1, len(missing2)), dtype=np.int64), missing2[np.newaxis, :]]) all_matches = np.concatenate([MN, MN2, MN3], axis=1) kp1_np = kp1_np.reshape((1, -1, 2)) kp2_np = kp2_np.reshape((1, -1, 2)) descs1 = np.transpose(descs1 / 256.) descs2 = np.transpose(descs2 / 256.) image = torch.from_numpy(image / 255.).double()[None].cuda() warped = torch.from_numpy(warped / 255.).double()[None].cuda() print('get item time:', time.time() - before_getitem_time) print('kp1_size, kp2_size:', kp1_np.size, kp2_np.size) print('all_matches:', all_matches) return { 'keypoints0': list(kp1_np), 'keypoints1': list(kp2_np), 'descriptors0': list(descs1), 'descriptors1': list(descs2), 'scores0': list(scores1_np), 'scores1': list(scores2_np), 'image0': image, 'image1': warped, 'all_matches': list(all_matches), 'file_name': file_name }
def __getitem__(self, idx): idx = idx % len(self.image_names) # Read image image = cv2.imread( os.path.join(self.image_path, self.image_names[idx]), cv2.IMREAD_GRAYSCALE) image = cv2.resize(image, self.image_size[::-1]) height, width = image.shape[:2] min_size = min(height, width) # Transform image corners = np.array([[0, 0], [0, height], [width, 0], [width, height]], dtype=np.float32) warp = np.random.randint(-min_size / 4, min_size / 4, size=(4, 2)).astype(np.float32) M = cv2.getPerspectiveTransform(corners, corners + warp) image_warped = cv2.warpPerspective(image, M, (width, height)) if self.DEBUG: print(f'Image size: {image.shape} -> {image_warped.shape}') with torch.no_grad(): # Extract keypoints data = frame2tensor(image, self.device) pred0 = self.superpoint({'image': data}) kps0 = pred0['keypoints'][0] desc0 = pred0['descriptors'][0] scores0 = pred0['scores'][0] # filter keypoints idxs = np.argsort( scores0.data.cpu().numpy())[::-1][:self.max_keypoints] scores0 = scores0[idxs.copy()] kps0 = kps0[idxs.copy(), :] desc0 = desc0[:, idxs.copy()] if self.DEBUG: print( f'Original keypoints: {kps0.shape}, descriptors: {desc0.shape}, scores: {scores0.shape}' ) # Transform keypoints kps1 = cv2.perspectiveTransform(kps0.cpu().numpy()[None], M) # Filter keypoints matches = [[], []] kps1_filtered = [] border = self.superpoint.config.get('remove_borders', 4) for i, k in enumerate(kps1.squeeze()): if k[0] < border or k[0] >= width - border: continue if k[1] < border or k[1] >= height - border: continue kps1_filtered.append(k) matches[0].append(i) matches[1].append(len(matches[1])) all_matches = [torch.tensor(ms) for ms in matches] kps1_filtered = array2tensor(np.array(kps1_filtered), self.device) # Compute descriptors & scores data_warped = frame2tensor(image_warped, self.device) desc1, scores1 = self.superpoint.computeDescriptorsAndScores({ 'image': data_warped, 'keypoints': kps1_filtered }) if self.DEBUG: print( f'Transformed keypoints: {kps1_filtered.shape}, descriptor: {desc1[0].shape}, scores: {scores1[0].shape}' ) # Draw keypoints and matches if self.DEBUG: kps0cv = [ cv2.KeyPoint(k[0], k[1], 8) for k in kps0.cpu().numpy().squeeze() ] kps1cv = [ cv2.KeyPoint(k[0], k[1], 8) for k in kps1_filtered.cpu().numpy().squeeze() ] matchescv = [ cv2.DMatch(k0, k1, 0) for k0, k1 in zip(matches[0], matches[1]) ] outimg = None outimg = cv2.drawMatches(image, kps0cv, image_warped, kps1cv, matchescv, outimg) cv2.imwrite('matches.jpg', outimg) outimg = cv2.drawKeypoints(image, kps0cv, outimg) cv2.imwrite('keypoints0.jpg', outimg) outimg = cv2.drawKeypoints(image_warped, kps1cv, outimg) cv2.imwrite('keypoints1.jpg', outimg) return { 'keypoints0': kps0, 'keypoints1': kps1_filtered[0], 'descriptors0': desc0, 'descriptors1': desc1[0], 'scores0': scores0, 'scores1': scores1[0], 'image0': data.squeeze(0), 'image1': data_warped.squeeze(0), 'all_matches': all_matches, 'file_name': self.image_names[idx], 'homography': M }
def __getitem__(self, idx): img1, img2, H = self.pairs[idx] # Read image image = cv2.imread(img1, cv2.IMREAD_GRAYSCALE) height, width = image.shape[:2] min_size = min(height, width) # Extract keypoints data = frame2tensor(image, self.device) pred0 = self.superpoint({'image': data}) kps0 = pred0['keypoints'][0] desc0 = pred0['descriptors'][0] scores0 = pred0['scores'][0] if self.DEBUG: print( f'Original keypoints: {kps0.shape}, descriptors: {desc0.shape}, scores: {scores0.shape}' ) # filter keypoints idxs = np.argsort( scores0.data.cpu().numpy())[::-1][:self.max_keypoints] scores0 = scores0[idxs.copy()] kps0 = kps0[idxs.copy(), :] desc0 = desc0[:, idxs.copy()] image_warped = cv2.imread(img2, cv2.IMREAD_GRAYSCALE) # Extract keypoints data_warped = frame2tensor(image_warped, self.device) pred1 = self.superpoint({'image': data_warped}) kps1 = pred1['keypoints'][0] desc1 = pred1['descriptors'][0] scores1 = pred1['scores'][0] if self.DEBUG: print( f'Original keypoints: {kps1.shape}, descriptors: {desc1.shape}, scores: {scores1.shape}' ) # filter keypoints idxs = np.argsort( scores1.data.cpu().numpy())[::-1][:self.max_keypoints] scores1 = scores1[idxs.copy()] kps1 = kps1[idxs.copy(), :] desc1 = desc1[:, idxs.copy()] # Draw keypoints and matches if self.DEBUG: kps0cv = [ cv2.KeyPoint(k[0], k[1], 8) for k in kps0.cpu().numpy().squeeze() ] kps1cv = [ cv2.KeyPoint(k[0], k[1], 8) for k in kps1.cpu().numpy().squeeze() ] outimg = None outimg = cv2.drawKeypoints(image, kps0cv, outimg) cv2.imwrite('keypoints0.jpg', outimg) outimg = cv2.drawKeypoints(image_warped, kps1cv, outimg) cv2.imwrite('keypoints1.jpg', outimg) return { 'keypoints0': kps0, 'keypoints1': kps1, 'descriptors0': desc0, 'descriptors1': desc1, 'scores0': scores0, 'scores1': scores1, 'image0': data.squeeze(0), 'image1': data_warped.squeeze(0), 'file_name': img2, 'homography': H }
def __getitem__(self, idx): # Read image image = cv2.imread( os.path.join(self.image_path, self.image_names[idx]), cv2.IMREAD_GRAYSCALE) height, width = image.shape[:2] min_size = min(height, width) # Transform image corners = np.array([[0, 0], [0, height], [width, 0], [width, height]], dtype=np.float32) warp = np.random.randint(-min_size / 4, min_size / 4, size=(4, 2)).astype(np.float32) M = cv2.getPerspectiveTransform(corners, corners + warp) image_warped = cv2.warpPerspective(image, M, (width, height)) if self.DEBUG: print(f'Image size: {image.shape} -> {image_warped.shape}') # Extract keypoints data = frame2tensor(image, self.device) pred0 = self.superpoint({'image': data}) kps0 = pred0['keypoints'][0] desc0 = pred0['descriptors'][0] scores0 = pred0['scores'][0] if self.DEBUG: print( f'Original keypoints: {kps0.shape}, descriptors: {desc0.shape}, scores: {scores0.shape}' ) # Extract keypoints data_warped = frame2tensor(image_warped, self.device) pred1 = self.superpoint({'image': data_warped}) kps1 = pred1['keypoints'][0] desc1 = pred1['descriptors'][0] scores1 = pred1['scores'][0] if self.DEBUG: print( f'Original keypoints: {kps1.shape}, descriptors: {desc1.shape}, scores: {scores1.shape}' ) # Draw keypoints and matches if self.DEBUG: kps0cv = [ cv2.KeyPoint(k[0], k[1], 8) for k in kps0.cpu().numpy().squeeze() ] kps1cv = [ cv2.KeyPoint(k[0], k[1], 8) for k in kps1.cpu().numpy().squeeze() ] outimg = None outimg = cv2.drawKeypoints(image, kps0cv, outimg) cv2.imwrite('keypoints0.jpg', outimg) outimg = cv2.drawKeypoints(image_warped, kps1cv, outimg) cv2.imwrite('keypoints1.jpg', outimg) return { 'keypoints0': kps0, 'keypoints1': kps1, 'descriptors0': desc0, 'descriptors1': desc1, 'scores0': scores0, 'scores1': scores1, 'image0': data.squeeze(0), 'image1': data_warped.squeeze(0), 'file_name': self.image_names[idx], 'homography': M }
def __getitem__(self, idx): (image1, depth1, intrinsics1, pose1, bbox1, image2, depth2, intrinsics2, pose2, bbox2) = self.recover_pair(self.dataset[idx]) # SIFT # kp1, descs1 = self.sift.detectAndCompute(cv2.cvtColor(image1, cv2.COLOR_BGR2GRAY), None) # kp2, descs2 = self.sift.detectAndCompute(cv2.cvtColor(image2, cv2.COLOR_BGR2GRAY), None) # SuperPoint SuperPoint = self.superpoint kp1, descs1 = self.parse_superpoint_result( SuperPoint({ 'image': frame2tensor(cv2.cvtColor(image1, cv2.COLOR_BGR2GRAY)) })) kp2, descs2 = self.parse_superpoint_result( SuperPoint({ 'image': frame2tensor(cv2.cvtColor(image2, cv2.COLOR_BGR2GRAY)) })) # im_with_keypoints1 = cv2.drawKeypoints(image1, kp1, np.array([]), (255,0,0)) # im_with_keypoints2 = cv2.drawKeypoints(image2, kp2, np.array([]), (255,0,0)) # cv2.imwrite('match_000_kp1.png', im_with_keypoints1) # cv2.imwrite('match_000_kp2.png', im_with_keypoints2) kp1_num = min(self.nfeatures, len(kp1)) kp2_num = min(self.nfeatures, len(kp2)) if kp1_num < 10 or kp2_num < 10: return { 'keypoints0': torch.zeros([0, 0, 2], dtype=torch.double), 'keypoints1': torch.zeros([0, 0, 2], dtype=torch.double), 'descriptors0': torch.zeros([0, 2], dtype=torch.double), 'descriptors1': torch.zeros([0, 2], dtype=torch.double), 'image0': image1, 'image1': image2, 'file_name': '' } kp1 = kp1[:kp1_num] kp2 = kp2[:kp2_num] descs1 = descs1[:kp1_num, :] descs2 = descs2[:kp2_num, :] # kp1_np = np.array([[kp.pt[0], kp.pt[1]] for kp in kp1]) # kp2_np = np.array([[kp.pt[0], kp.pt[1]] for kp in kp2]) kp1_np = np.array([(kp[0], kp[1]) for kp in kp1]) kp2_np = np.array([(kp[0], kp[1]) for kp in kp2]) KP1 = kp1_np KP2 = kp2_np # scores1_np = np.array([kp.response for kp in kp1]) # scores2_np = np.array([kp.response for kp in kp2]) scores1_np = np.array([kp[2] for kp in kp1]) scores2_np = np.array([kp[2] for kp in kp2]) kp1_np = kp1_np[:kp1_num, :] kp2_np = kp2_np[:kp2_num, :] descs1 = descs1[:kp1_num, :] descs2 = descs2[:kp2_num, :] kp1_np = kp1_np.reshape((1, -1, 2)) kp2_np = kp2_np.reshape((1, -1, 2)) # descs1 = np.transpose(descs1 / 256.) # descs2 = np.transpose(descs2 / 256.) descs1 = np.transpose(descs1) descs2 = np.transpose(descs2) image1_o = image1 image2_o = image2 image1 = torch.from_numpy( cv2.cvtColor(image1, cv2.COLOR_BGR2GRAY) / 255.).double()[None].cuda() image2 = torch.from_numpy( cv2.cvtColor(image2, cv2.COLOR_BGR2GRAY) / 255.).double()[None].cuda() try: all_matches = self.compute_all_matches(KP1, image1_o, depth1, intrinsics1, pose1, bbox1, KP2, image2_o, depth2, intrinsics2, pose2, bbox2) except EmptyTensorError: return { 'keypoints0': torch.zeros([0, 0, 2], dtype=torch.double), 'keypoints1': torch.zeros([0, 0, 2], dtype=torch.double), 'descriptors0': torch.zeros([0, 2], dtype=torch.double), 'descriptors1': torch.zeros([0, 2], dtype=torch.double), 'image0': image1, 'image1': image2, 'file_name': '' } return { 'keypoints0': list(kp1_np), 'keypoints1': list(kp2_np), 'descriptors0': list(descs1), 'descriptors1': list(descs2), 'scores0': list(scores1_np), 'scores1': list(scores2_np), 'image0': image1, 'image1': image2, 'all_matches': all_matches, 'file_name': '' }
def detect(self, img): with torch.no_grad(): result = self.detector({'image': frame2tensor(img.astype('float32')).cuda()}) return result['keypoints'][0].cpu().float().numpy(), \ result['descriptors'][0].cpu().float().numpy(), \ result['scores'][0].cpu().float().numpy()
def __getitem__(self, idx): (image1, depth1, intrinsics1, pose1, bbox1, image2, depth2, intrinsics2, pose2, bbox2) = self.recover_pair(self.dataset[idx]) image1 = preprocess_image( image1, preprocessing=self.preprocessing) # 得到BGR格式图像,不做中心化 image2 = preprocess_image( image2, preprocessing=self.preprocessing) # 得到BGR格式图像,不做中心化 original_image1 = image1 original_image2 = image2 ''' return { 'image1': torch.from_numpy(image1.astype(np.float32)), 'depth1': torch.from_numpy(depth1.astype(np.float32)), 'intrinsics1': torch.from_numpy(intrinsics1.astype(np.float32)), 'pose1': torch.from_numpy(pose1.astype(np.float32)), 'bbox1': torch.from_numpy(bbox1.astype(np.float32)), 'image2': torch.from_numpy(image2.astype(np.float32)), 'depth2': torch.from_numpy(depth2.astype(np.float32)), 'intrinsics2': torch.from_numpy(intrinsics2.astype(np.float32)), 'pose2': torch.from_numpy(pose2.astype(np.float32)), 'bbox2': torch.from_numpy(bbox2.astype(np.float32)) } ''' # 用SuperPoint处理两张图片,得到keypoints、descriptors、scores SuperPoint = self.superpoint image1 = np.transpose(image1, (1, 2, 0)) image1 = Image.fromarray(np.uint8(image1)) image1 = image1.convert('L') image1 = np.array(image1) image2 = np.transpose(image2, (1, 2, 0)) image2 = Image.fromarray(np.uint8(image2)) image2 = image2.convert('L') image2 = np.array(image2) kp1, descs1 = self.parse_superpoint_result( SuperPoint({'image': frame2tensor(image1)})) kp2, descs2 = self.parse_superpoint_result( SuperPoint({'image': frame2tensor(image2)})) # limit the number of keypoints kp1_num = min(self.nfeatures, len(kp1)) kp2_num = min(self.nfeatures, len(kp2)) kp1 = kp1[:kp1_num] kp2 = kp2[:kp2_num] kp1_np = np.array([(kp[0], kp[1]) for kp in kp1]) kp2_np = np.array([(kp[0], kp[1]) for kp in kp2]) # skip this image pair if no keypoints detected in image # 不到10个特征点也跳过此图片对 if len(kp1) < 10 or len(kp2) < 10: return { 'keypoints0': torch.zeros([0, 0, 2], dtype=torch.double), 'keypoints1': torch.zeros([0, 0, 2], dtype=torch.double), 'descriptors0': torch.zeros([0, 2], dtype=torch.double), 'descriptors1': torch.zeros([0, 2], dtype=torch.double), 'image0': image1, 'image1': image2, 'file_name': '' } # confidence of each key point scores1_np = np.array([kp[2] for kp in kp1]) scores2_np = np.array([kp[2] for kp in kp2]) kp1_np = kp1_np[:kp1_num, :] kp2_np = kp2_np[:kp2_num, :] descs1 = descs1[:kp1_num, :] descs2 = descs2[:kp2_num, :] kp1_np = kp1_np.reshape((1, -1, 2)) kp2_np = kp2_np.reshape((1, -1, 2)) descs1 = np.transpose(descs1 / 256.) descs2 = np.transpose(descs2 / 256.) image1 = torch.from_numpy(image1 / 255.).double()[None].cuda() image2 = torch.from_numpy(image2 / 255.).double()[None].cuda() # print(image1.shape, image2.shape, depth1.shape, depth2.shape) # 根据10元组和keypoints,得到所有匹配,按SuperGlue的输入要求返回结果 # image1, depth1, intrinsics1, pose1, bbox1 # image2, depth2, intrinsics2, pose2, bbox2 # depth: (256, 256), intrinsics: (3, 3), pose: (4, 4), bbox: (2) # 例子:all_matches = list(np.array([[0],[0]])) try: all_matches = self.compute_all_matches(kp1_np, original_image1, depth1, intrinsics1, pose1, bbox1, kp2_np, original_image2, depth2, intrinsics2, pose2, bbox2) except EmptyTensorError: return { 'keypoints0': torch.zeros([0, 0, 2], dtype=torch.double), 'keypoints1': torch.zeros([0, 0, 2], dtype=torch.double), 'descriptors0': torch.zeros([0, 2], dtype=torch.double), 'descriptors1': torch.zeros([0, 2], dtype=torch.double), 'image0': image1, 'image1': image2, 'file_name': '' } # print(kp1_np.shape, kp2_np.shape, len(all_matches[0])) return { 'keypoints0': list(kp1_np), 'keypoints1': list(kp2_np), 'descriptors0': list(descs1), 'descriptors1': list(descs2), 'scores0': list(scores1_np), 'scores1': list(scores2_np), 'image0': image1, 'image1': image2, 'all_matches': all_matches, 'file_name': '' } # SuperGlue要的返回值 '''