diff --git a/lightglue/__init__.py b/lightglue/__init__.py index 42719c9..95c3654 100644 --- a/lightglue/__init__.py +++ b/lightglue/__init__.py @@ -3,4 +3,5 @@ from .lightglue import LightGlue # noqa from .sift import SIFT # noqa from .superpoint import SuperPoint # noqa +from .dog_hardnet import DoGHardNet # noqa from .utils import match_pair # noqa diff --git a/lightglue/dog_hardnet.py b/lightglue/dog_hardnet.py index e53e97a..8c7e465 100644 --- a/lightglue/dog_hardnet.py +++ b/lightglue/dog_hardnet.py @@ -1,41 +1,19 @@ -import warnings - -import cv2 -import numpy as np import torch from kornia.color import rgb_to_grayscale from kornia.feature import HardNet, LAFDescriptor, laf_from_center_scale_ori -from packaging import version -from .utils import Extractor from .sift import SIFT -try: - import pycolmap -except ImportError: - pycolmap = None - class DoGHardNet(SIFT): - default_conf = { - "nms_radius": 0, # None to disable filtering entirely. - "max_num_keypoints": 2048, - "backend": "opencv", # in {opencv, pycolmap, pycolmap_cpu, pycolmap_cuda} - "detection_threshold": -1, # from COLMAP - "edge_threshold": -1, - "first_octave": -1, # only used by pycolmap, the default of COLMAP - "num_octaves": 4, - "force_num_keypoints": True, - } - required_data_keys = ["image"] - def _init(self, conf): - super()._init(conf) + def __init__(self, **conf): + super().__init__(**conf) self.laf_desc = LAFDescriptor(HardNet(True)).eval() - def _forward(self, data: dict) -> dict: + def forward(self, data: dict) -> dict: image = data["image"] if image.shape[1] == 3: image = rgb_to_grayscale(image) @@ -43,7 +21,10 @@ def _forward(self, data: dict) -> dict: self.laf_desc = self.laf_desc.to(device) self.laf_desc.descriptor = self.laf_desc.descriptor.eval() pred = [] - im_size = data.get("image_size").long() + if "image_size" in data.keys(): + im_size = data.get("image_size").long() + else: + im_size = None for k in range(len(image)): img = image[k] if im_size is not None: @@ -58,8 +39,4 @@ def _forward(self, data: dict) -> dict: p["descriptors"] = self.laf_desc(img[None], lafs).reshape(-1, 128) pred.append(p) pred = {k: torch.stack([p[k] for p in pred], 0).to(device) for k in pred[0]} - return pred - - - - \ No newline at end of file + return pred \ No newline at end of file