Spaces:
Running
Running
File size: 1,061 Bytes
9223079 b7f7f2c 9223079 b7f7f2c 9223079 8004049 9223079 |
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 |
import kornia
from hloc import logger
from ..utils.base_model import BaseModel
class DISK(BaseModel):
default_conf = {
"weights": "depth",
"max_keypoints": None,
"nms_window_size": 5,
"detection_threshold": 0.0,
"pad_if_not_divisible": True,
}
required_inputs = ["image"]
def _init(self, conf):
self.model = kornia.feature.DISK.from_pretrained(conf["weights"])
logger.info(f"Load DISK model done.")
def _forward(self, data):
image = data["image"]
features = self.model(
image,
n=self.conf["max_keypoints"],
window_size=self.conf["nms_window_size"],
score_threshold=self.conf["detection_threshold"],
pad_if_not_divisible=self.conf["pad_if_not_divisible"],
)
return {
"keypoints": [f.keypoints for f in features][0][None],
"scores": [f.detection_scores for f in features][0][None],
"descriptors": [f.descriptors.t() for f in features][0][None],
}
|