Spaces:
Running
Running
File size: 864 Bytes
b7f7f2c aa46ae9 b7f7f2c |
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 |
import torch
from pathlib import Path
from hloc import logger
from ..utils.base_model import BaseModel
class XFeat(BaseModel):
default_conf = {
"keypoint_threshold": 0.005,
"max_keypoints": -1,
}
required_inputs = ["image"]
def _init(self, conf):
self.net = torch.hub.load(
"verlab/accelerated_features",
"XFeat",
pretrained=True,
top_k=self.conf["max_keypoints"],
)
logger.info(f"Load XFeat(sparse) model done.")
def _forward(self, data):
pred = self.net.detectAndCompute(
data["image"], top_k=self.conf["max_keypoints"]
)[0]
pred = {
"keypoints": pred["keypoints"][None],
"scores": pred["scores"][None],
"descriptors": pred["descriptors"].T[None],
}
return pred
|