diff --git a/hloc/extract_features.py b/hloc/extract_features.py index ab9456a8..65b1a098 100644 --- a/hloc/extract_features.py +++ b/hloc/extract_features.py @@ -125,6 +125,34 @@ "resize_max": 1024, }, }, + "dad": { + "output": "feats-dad", + "model": { + "name": "elf-extractor", + "elf_detector": "dad", + "elf_detector_conf": {}, + "elf_descriptor": "xfeat", + "elf_descriptor_conf": {}, + }, + "preprocessing": { + "grayscale": False, + }, + }, + "dedode": { + "output": "feats-dedode", + "model": { + "name": "elf-extractor", + "elf_detector": "dedode", + "elf_detector_conf": { + "detector_weights": "L-C4-v2", + }, + "elf_descriptor": "xfeat", + "elf_descriptor_conf": {}, + }, + "preprocessing": { + "grayscale": False, + }, + }, # Global descriptors "dir": { "output": "global-feats-dir", @@ -240,7 +268,7 @@ def main( overwrite: bool = False, ) -> Path: logger.info( - "Extracting local features with configuration:" f"\n{pprint.pformat(conf)}" + f"Extracting local features with configuration:\n{pprint.pformat(conf)}" ) dataset = ImageDataset(image_dir, conf["preprocessing"], image_list) diff --git a/hloc/extractors/elf-extractor.py b/hloc/extractors/elf-extractor.py new file mode 100644 index 00000000..00cd5118 --- /dev/null +++ b/hloc/extractors/elf-extractor.py @@ -0,0 +1,42 @@ +import torch +from easy_local_features import getExtractor + +from ..utils.base_model import BaseModel + + +class ElfFeatures(BaseModel): + default_conf = { + "model": { + "name": "elf-detectors", + "elf_detector": "xfeat", + "elf_detector_conf": { + "top_k": 512, + }, + "elf_descriptor": "xfeat", + "elf_descriptor_conf": {}, + }, + } + + def _init(self, conf): + print(f"{conf=}") + self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu") + self.detector = getExtractor(conf["elf_detector"], conf["elf_detector_conf"]) + self.detector.to(self.device) + + self.descriptor = getExtractor( + conf["elf_descriptor"], conf["elf_descriptor_conf"] + ) + self.descriptor.to(self.device) + + def _forward(self, data): + kps = self.detector.detect(data["image"])[0] + + desc = self.descriptor.compute(data["image"], kps).detach().squeeze(0).T + return { + "keypoints": [ + kps.cpu(), + ], + "descriptors": [ + desc.cpu(), + ], + } diff --git a/requirements.txt b/requirements.txt index 27d828dc..0d1934d7 100644 --- a/requirements.txt +++ b/requirements.txt @@ -11,3 +11,4 @@ pycolmap>=3.12.6 kornia>=0.6.11 gdown lightglue @ git+https://github.com/cvg/LightGlue +easy-local-features @ git+https://github.com/felipecadar/easy-local-features-baselines