Spaces:
Running
Running
File size: 1,551 Bytes
9223079 2eaeef9 9223079 e15a186 9223079 2eaeef9 9223079 ed0584b |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 |
import sys
from pathlib import Path
import torch
from ..utils.base_model import BaseModel
from hloc import logger
sys.path.append(str(Path(__file__).parent / "../../third_party"))
from SuperGluePretrainedNetwork.models import superpoint # noqa E402
# The original keypoint sampling is incorrect. We patch it here but
# we don't fix it upstream to not impact exisiting evaluations.
def sample_descriptors_fix_sampling(keypoints, descriptors, s: int = 8):
"""Interpolate descriptors at keypoint locations"""
b, c, h, w = descriptors.shape
keypoints = (keypoints + 0.5) / (keypoints.new_tensor([w, h]) * s)
keypoints = keypoints * 2 - 1 # normalize to (-1, 1)
descriptors = torch.nn.functional.grid_sample(
descriptors,
keypoints.view(b, 1, -1, 2),
mode="bilinear",
align_corners=False,
)
descriptors = torch.nn.functional.normalize(
descriptors.reshape(b, c, -1), p=2, dim=1
)
return descriptors
class SuperPoint(BaseModel):
default_conf = {
"nms_radius": 4,
"keypoint_threshold": 0.005,
"max_keypoints": -1,
"remove_borders": 4,
"fix_sampling": False,
}
required_inputs = ["image"]
detection_noise = 2.0
def _init(self, conf):
if conf["fix_sampling"]:
superpoint.sample_descriptors = sample_descriptors_fix_sampling
self.net = superpoint.SuperPoint(conf)
logger.info(f"Load SuperPoint model done.")
def _forward(self, data):
return self.net(data, self.conf)
|