diff --git a/README.md b/README.md index 9b6b5f4..117d8e3 100644 --- a/README.md +++ b/README.md @@ -43,7 +43,7 @@ We provide a [demo notebook](demo.ipynb) which shows how to perform feature extr Here is a minimal script to match two images: ```python -from lightglue import LightGlue, SuperPoint, DISK, SIFT, ALIKED +from lightglue import LightGlue, SuperPoint, DISK, SIFT, ALIKED, DoGHardNet from lightglue.utils import load_image, rbd # SuperPoint+LightGlue diff --git a/lightglue/__init__.py b/lightglue/__init__.py index 42719c9..b84d285 100644 --- a/lightglue/__init__.py +++ b/lightglue/__init__.py @@ -1,5 +1,6 @@ from .aliked import ALIKED # noqa from .disk import DISK # noqa +from .dog_hardnet import DoGHardNet # noqa from .lightglue import LightGlue # noqa from .sift import SIFT # noqa from .superpoint import SuperPoint # noqa diff --git a/lightglue/dog_hardnet.py b/lightglue/dog_hardnet.py new file mode 100644 index 0000000..cce307a --- /dev/null +++ b/lightglue/dog_hardnet.py @@ -0,0 +1,41 @@ +import torch +from kornia.color import rgb_to_grayscale +from kornia.feature import HardNet, LAFDescriptor, laf_from_center_scale_ori + +from .sift import SIFT + + +class DoGHardNet(SIFT): + required_data_keys = ["image"] + + def __init__(self, **conf): + super().__init__(**conf) + self.laf_desc = LAFDescriptor(HardNet(True)).eval() + + def forward(self, data: dict) -> dict: + image = data["image"] + if image.shape[1] == 3: + image = rgb_to_grayscale(image) + device = image.device + self.laf_desc = self.laf_desc.to(device) + self.laf_desc.descriptor = self.laf_desc.descriptor.eval() + pred = [] + 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: + w, h = data["image_size"][k] + img = img[:, : h.to(torch.int32), : w.to(torch.int32)] + p = self.extract_single_image(img) + lafs = laf_from_center_scale_ori( + p["keypoints"].reshape(1, -1, 2), + 6.0 * p["scales"].reshape(1, -1, 1, 1), + torch.rad2deg(p["oris"]).reshape(1, -1, 1), + ).to(device) + 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 diff --git a/lightglue/lightglue.py b/lightglue/lightglue.py index fcbb7ee..d437841 100644 --- a/lightglue/lightglue.py +++ b/lightglue/lightglue.py @@ -357,6 +357,11 @@ class LightGlue(nn.Module): "input_dim": 128, "add_scale_ori": True, }, + "doghardnet": { + "weights": "doghardnet_lightglue", + "input_dim": 128, + "add_scale_ori": True, + }, } def __init__(self, features="superpoint", **conf) -> None: