Spaces:
Running
Running
File size: 1,532 Bytes
9223079 8320ccc 9223079 2eaeef9 9223079 8320ccc 9223079 8320ccc 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 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 |
import sys
from pathlib import Path
import torch
from hloc import logger
from ..utils.base_model import BaseModel
alike_path = Path(__file__).parent / "../../third_party/ALIKE"
sys.path.append(str(alike_path))
from alike import ALike as Alike_
from alike import configs
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
class Alike(BaseModel):
default_conf = {
"model_name": "alike-t", # 'alike-t', 'alike-s', 'alike-n', 'alike-l'
"use_relu": True,
"multiscale": False,
"max_keypoints": 1000,
"detection_threshold": 0.5,
"top_k": -1,
"sub_pixel": False,
}
required_inputs = ["image"]
def _init(self, conf):
self.net = Alike_(
**configs[conf["model_name"]],
device=device,
top_k=conf["top_k"],
scores_th=conf["detection_threshold"],
n_limit=conf["max_keypoints"],
)
logger.info("Load Alike model done.")
def _forward(self, data):
image = data["image"]
image = image.permute(0, 2, 3, 1).squeeze()
image = image.cpu().numpy() * 255.0
pred = self.net(image, sub_pixel=self.conf["sub_pixel"])
keypoints = pred["keypoints"]
descriptors = pred["descriptors"]
scores = pred["scores"]
return {
"keypoints": torch.from_numpy(keypoints)[None],
"scores": torch.from_numpy(scores)[None],
"descriptors": torch.from_numpy(descriptors.T)[None],
}
|