File size: 992 Bytes
9223079
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
62cd4d6
 
 
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
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],
        }