hloc/extractors/alike.py

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],
        }