|
""" |
|
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 |
|
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 |
|
|
