Spaces:
Running
Running
| """ | |
| A two-view sparse feature matching pipeline. | |
| This model contains sub-models for each step: | |
| feature extraction, feature matching, outlier filtering, pose estimation. | |
| Each step is optional, and the features or matches can be provided as input. | |
| Default: SuperPoint with nearest neighbor matching. | |
| Convention for the matches: m0[i] is the index of the keypoint in image 1 | |
| that corresponds to the keypoint i in image 0. m0[i] = -1 if i is unmatched. | |
| """ | |
| import numpy as np | |
| import torch | |
| from .. import get_model | |
| from .base_model import BaseModel | |
| def keep_quadrant_kp_subset(keypoints, scores, descs, h, w): | |
| """Keep only keypoints in one of the four quadrant of the image.""" | |
| h2, w2 = h // 2, w // 2 | |
| w_x = np.random.choice([0, w2]) | |
| w_y = np.random.choice([0, h2]) | |
| valid_mask = ((keypoints[..., 0] >= w_x) | |
| & (keypoints[..., 0] < w_x + w2) | |
| & (keypoints[..., 1] >= w_y) | |
| & (keypoints[..., 1] < w_y + h2)) | |
| keypoints = keypoints[valid_mask][None] | |
| scores = scores[valid_mask][None] | |
| descs = descs.permute(0, 2, 1)[valid_mask].t()[None] | |
| return keypoints, scores, descs | |
| def keep_random_kp_subset(keypoints, scores, descs, num_selected): | |
| """Keep a random subset of keypoints.""" | |
| num_kp = keypoints.shape[1] | |
| selected_kp = torch.randperm(num_kp)[:num_selected] | |
| keypoints = keypoints[:, selected_kp] | |
| scores = scores[:, selected_kp] | |
| descs = descs[:, :, selected_kp] | |
| return keypoints, scores, descs | |
| def keep_best_kp_subset(keypoints, scores, descs, num_selected): | |
| """Keep the top num_selected best keypoints.""" | |
| sorted_indices = torch.sort(scores, dim=1)[1] | |
| selected_kp = sorted_indices[:, -num_selected:] | |
| keypoints = torch.gather(keypoints, 1, | |
| selected_kp[:, :, None].repeat(1, 1, 2)) | |
| scores = torch.gather(scores, 1, selected_kp) | |
| descs = torch.gather(descs, 2, | |
| selected_kp[:, None].repeat(1, descs.shape[1], 1)) | |
| return keypoints, scores, descs | |
| class TwoViewPipeline(BaseModel): | |
| default_conf = { | |
| 'extractor': { | |
| 'name': 'superpoint', | |
| 'trainable': False, | |
| }, | |
| 'use_lines': False, | |
| 'use_points': True, | |
| 'randomize_num_kp': False, | |
| 'detector': {'name': None}, | |
| 'descriptor': {'name': None}, | |
| 'matcher': {'name': 'nearest_neighbor_matcher'}, | |
| 'filter': {'name': None}, | |
| 'solver': {'name': None}, | |
| 'ground_truth': { | |
| 'from_pose_depth': False, | |
| 'from_homography': False, | |
| 'th_positive': 3, | |
| 'th_negative': 5, | |
| 'reward_positive': 1, | |
| 'reward_negative': -0.25, | |
| 'is_likelihood_soft': True, | |
| 'p_random_occluders': 0, | |
| 'n_line_sampled_pts': 50, | |
| 'line_perp_dist_th': 5, | |
| 'overlap_th': 0.2, | |
| 'min_visibility_th': 0.5 | |
| }, | |
| } | |
| required_data_keys = ['image0', 'image1'] | |
| strict_conf = False # need to pass new confs to children models | |
| components = [ | |
| 'extractor', 'detector', 'descriptor', 'matcher', 'filter', 'solver'] | |
| def _init(self, conf): | |
| if conf.extractor.name: | |
| self.extractor = get_model(conf.extractor.name)(conf.extractor) | |
| else: | |
| if self.conf.detector.name: | |
| self.detector = get_model(conf.detector.name)(conf.detector) | |
| else: | |
| self.required_data_keys += ['keypoints0', 'keypoints1'] | |
| if self.conf.descriptor.name: | |
| self.descriptor = get_model(conf.descriptor.name)( | |
| conf.descriptor) | |
| else: | |
| self.required_data_keys += ['descriptors0', 'descriptors1'] | |
| if conf.matcher.name: | |
| self.matcher = get_model(conf.matcher.name)(conf.matcher) | |
| else: | |
| self.required_data_keys += ['matches0'] | |
| if conf.filter.name: | |
| self.filter = get_model(conf.filter.name)(conf.filter) | |
| if conf.solver.name: | |
| self.solver = get_model(conf.solver.name)(conf.solver) | |
| def _forward(self, data): | |
| def process_siamese(data, i): | |
| data_i = {k[:-1]: v for k, v in data.items() if k[-1] == i} | |
| if self.conf.extractor.name: | |
| pred_i = self.extractor(data_i) | |
| else: | |
| pred_i = {} | |
| if self.conf.detector.name: | |
| pred_i = self.detector(data_i) | |
| else: | |
| for k in ['keypoints', 'keypoint_scores', 'descriptors', | |
| 'lines', 'line_scores', 'line_descriptors', | |
| 'valid_lines']: | |
| if k in data_i: | |
| pred_i[k] = data_i[k] | |
| if self.conf.descriptor.name: | |
| pred_i = { | |
| **pred_i, **self.descriptor({**data_i, **pred_i})} | |
| return pred_i | |
| pred0 = process_siamese(data, '0') | |
| pred1 = process_siamese(data, '1') | |
| pred = {**{k + '0': v for k, v in pred0.items()}, | |
| **{k + '1': v for k, v in pred1.items()}} | |
| if self.conf.matcher.name: | |
| pred = {**pred, **self.matcher({**data, **pred})} | |
| if self.conf.filter.name: | |
| pred = {**pred, **self.filter({**data, **pred})} | |
| if self.conf.solver.name: | |
| pred = {**pred, **self.solver({**data, **pred})} | |
| return pred | |
| def loss(self, pred, data): | |
| losses = {} | |
| total = 0 | |
| for k in self.components: | |
| if self.conf[k].name: | |
| try: | |
| losses_ = getattr(self, k).loss(pred, {**pred, **data}) | |
| except NotImplementedError: | |
| continue | |
| losses = {**losses, **losses_} | |
| total = losses_['total'] + total | |
| return {**losses, 'total': total} | |
| def metrics(self, pred, data): | |
| metrics = {} | |
| for k in self.components: | |
| if self.conf[k].name: | |
| try: | |
| metrics_ = getattr(self, k).metrics(pred, {**pred, **data}) | |
| except NotImplementedError: | |
| continue | |
| metrics = {**metrics, **metrics_} | |
| return metrics | |