Add Extractor, cleanup SIFT, bugfixes in LightGlue

lightglue/aliked.py

@@ -32,7 +32,6 @@
 # Xiaoming Zhao, Xingming Wu, Weihai Chen, Peter C.Y. Chen, Qingsong Xu, and Zhengguo Li
 # Code from https://github.com/Shiaoming/ALIKED
 
-from types import SimpleNamespace
 from typing import Callable, Optional
 
 import torch
@@ -42,7 +41,7 @@
 from torch.nn.modules.utils import _pair
 from torchvision.models import resnet
 
-from .utils import ImagePreprocessor
+from .utils import Extractor
 
 
 def get_patches(
@@ -609,12 +608,11 @@ def forward(self, x, keypoints):
         return descriptors, offsets
 
 
-class ALIKED(nn.Module):
+class ALIKED(Extractor):
     default_conf = {
         "model_name": "aliked-n16",
         "max_num_keypoints": -1,
         "detection_threshold": 0.2,
-        "force_num_keypoints": False,
         "nms_radius": 2,
     }
 
@@ -630,20 +628,15 @@ class ALIKED(nn.Module):
         "aliked-n32": [16, 32, 64, 128, 128, 3, 32],
     }
     preprocess_conf = {
-        **ImagePreprocessor.default_conf,
         "resize": 1024,
         "grayscale": False,
     }
 
     required_data_keys = ["image"]
 
     def __init__(self, **conf):
-        super().__init__()
-        self.conf = {**self.default_conf, **conf}
-        conf = self.conf = SimpleNamespace(**self.conf)
-        if conf.force_num_keypoints:
-            assert conf.detection_threshold <= 0 and conf.max_num_keypoints > 0
-        # get configurations
+        super().__init__(**conf)  # Update with default configuration.
+        conf = self.conf
         c1, c2, c3, c4, dim, K, M = self.cfgs[conf.model_name]
         conv_types = ["conv", "conv", "dcn", "dcn"]
         conv2D = False
@@ -761,15 +754,3 @@ def forward(self, data: dict) -> dict:
             "descriptors": torch.stack(descriptors),  # B x N x D
             "keypoint_scores": torch.stack(kptscores),  # B x N
         }
-
-    def extract(self, img: torch.Tensor, **conf) -> dict:
-        """Perform extraction with online resizing"""
-        if img.dim() == 3:
-            img = img[None]  # add batch dim
-        assert img.dim() == 4 and img.shape[0] == 1
-        shape = img.shape[-2:][::-1]
-        img, scales = ImagePreprocessor(**{**self.preprocess_conf, **conf})(img)
-        feats = self.forward({"image": img})
-        feats["image_size"] = torch.tensor(shape)[None].to(img).float()
-        feats["keypoints"] = (feats["keypoints"] + 0.5) / scales[None] - 0.5
-        return feats

lightglue/disk.py

@@ -1,13 +1,10 @@
-from types import SimpleNamespace
-
 import kornia
 import torch
-import torch.nn as nn
 
-from .utils import ImagePreprocessor
+from .utils import Extractor
 
 
-class DISK(nn.Module):
+class DISK(Extractor):
     default_conf = {
         "weights": "depth",
         "max_num_keypoints": None,
@@ -18,17 +15,14 @@ class DISK(nn.Module):
     }
 
     preprocess_conf = {
-        **ImagePreprocessor.default_conf,
         "resize": 1024,
         "grayscale": False,
     }
 
     required_data_keys = ["image"]
 
     def __init__(self, **conf) -> None:
-        super().__init__()
-        self.conf = {**self.default_conf, **conf}
-        self.conf = SimpleNamespace(**self.conf)
+        super().__init__(**conf)  # Update with default configuration.
         self.model = kornia.feature.DISK.from_pretrained(self.conf.weights)
 
     def forward(self, data: dict) -> dict:
@@ -57,15 +51,3 @@ def forward(self, data: dict) -> dict:
             "keypoint_scores": scores.to(image).contiguous(),
             "descriptors": descriptors.to(image).contiguous(),
         }
-
-    def extract(self, img: torch.Tensor, **conf) -> dict:
-        """Perform extraction with online resizing"""
-        if img.dim() == 3:
-            img = img[None]  # add batch dim
-        assert img.dim() == 4 and img.shape[0] == 1
-        shape = img.shape[-2:][::-1]
-        img, scales = ImagePreprocessor(**{**self.preprocess_conf, **conf})(img)
-        feats = self.forward({"image": img})
-        feats["image_size"] = torch.tensor(shape)[None].to(img).float()
-        feats["keypoints"] = (feats["keypoints"] + 0.5) / scales[None] - 0.5
-        return feats

lightglue/lightglue.py

@@ -361,11 +361,15 @@ class LightGlue(nn.Module):
 
     def __init__(self, features="superpoint", **conf) -> None:
         super().__init__()
-        self.conf = {**self.default_conf, **conf}
+        self.conf = conf = SimpleNamespace(**{**self.default_conf, **conf})
         if features is not None:
-            assert features in list(self.features.keys())
-            self.conf["weights"], self.conf["input_dim"] = self.features[features]
-        self.conf = conf = SimpleNamespace(**self.conf)
+            if features not in self.features:
+                raise ValueError(
+                    f"Unsupported features: {features} not in "
+                    f"{{{','.join(self.features)}}}"
+                )
+            for k, v in self.features[features].items():
+                setattr(conf, k, v)
 
         if conf.input_dim != conf.descriptor_dim:
             self.input_proj = nn.Linear(conf.input_dim, conf.descriptor_dim, bias=True)
@@ -471,8 +475,12 @@ def _forward(self, data: dict) -> dict:
         kpts1 = normalize_keypoints(kpts1, size1).clone()
 
         if self.conf.add_scale_ori:
-            kpts0 = torch.cat([kpts0, data0["scales"], data0["oris"]], -1)
-            kpts1 = torch.cat([kpts1, data1["scales"], data1["oris"]], -1)
+            kpts0 = torch.cat(
+                [kpts0] + [data0[k].unsqueeze(-1) for k in ("scales", "oris")], -1
+            )
+            kpts1 = torch.cat(
+                [kpts1] + [data1[k].unsqueeze(-1) for k in ("scales", "oris")], -1
+            )
         desc0 = data0["descriptors"].detach().contiguous()
         desc1 = data1["descriptors"].detach().contiguous()