Vincentqyw
update: wrapper functions
62cd4d6
raw
history blame
992 Bytes
import kornia
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"])
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],
}