From 4e8978b8e107d25dbadd373ff939abe95a3a49dd Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?D=C3=A1vid=20Komorowicz?= Date: Sat, 23 Nov 2024 07:30:33 +0100 Subject: [PATCH] - Changed DeDoDe and LightGlue to kornia implementation - Changed DeDoDe to G variant - Added DeDoDe v2 - Added LightGlue support for DeDoDe --- hloc/extract_features.py | 38 +++++++++++++++++++-- hloc/extractors/dedode.py | 67 +++++--------------------------------- hloc/match_features.py | 32 ++++++++++++++++++ hloc/matchers/lightglue.py | 29 ++++++----------- ui/config.yaml | 26 +++++++++++++-- 5 files changed, 111 insertions(+), 81 deletions(-) diff --git a/hloc/extract_features.py b/hloc/extract_features.py index 62eeb45..fc3c290 100644 --- a/hloc/extract_features.py +++ b/hloc/extract_features.py @@ -259,10 +259,44 @@ "dfactor": 8, }, }, - "dedode": { - "output": "feats-dedode-n5000-r1600", + "dedodeg": { + "output": "feats-dedodeg-n5000-r1600", "model": { "name": "dedode", + "detector_model": "L-upright", + "descriptor_model": "G-upright", + "max_keypoints": 5000, + }, + "preprocessing": { + "grayscale": False, + "force_resize": True, + "resize_max": 1600, + "width": 768, + "height": 768, + "dfactor": 8, + }, + }, + "dedodeg-v2": { + "output": "feats-dedodegv2-n5000-r1600", + "model": { + "name": "dedode", + "detector_model": "L-C4-v2", + "descriptor_model": "G-upright", + "max_keypoints": 5000, + }, + "preprocessing": { + "grayscale": False, + "force_resize": True, + "resize_max": 1600, + "width": 768, + "height": 768, + "dfactor": 8, + }, + }, + "mickey": { + "output": "feats-mickey", + "model": { + "name": "mickey", "max_keypoints": 5000, }, "preprocessing": { diff --git a/hloc/extractors/dedode.py b/hloc/extractors/dedode.py index d6a228d..100d1d3 100644 --- a/hloc/extractors/dedode.py +++ b/hloc/extractors/dedode.py @@ -7,12 +7,7 @@ from hloc import MODEL_REPO_ID, logger from ..utils.base_model import BaseModel - -dedode_path = Path(__file__).parent / "../../third_party/DeDoDe" -sys.path.append(str(dedode_path)) - -from DeDoDe import dedode_descriptor_B, dedode_detector_L -from DeDoDe.utils import to_pixel_coords +from kornia import feature as KF device = torch.device("cuda" if torch.cuda.is_available() else "cpu") @@ -20,8 +15,8 @@ class DeDoDe(BaseModel): default_conf = { "name": "dedode", - "model_detector_name": "dedode_detector_L.pth", - "model_descriptor_name": "dedode_descriptor_B.pth", + "detector_model": "L-upright", + "descriptor_model": "G-upright", "max_keypoints": 2000, "match_threshold": 0.2, "dense": False, # Now fixed to be false @@ -32,34 +27,7 @@ class DeDoDe(BaseModel): # Initialize the line matcher def _init(self, conf): - model_detector_path = self._download_model( - repo_id=MODEL_REPO_ID, - filename="{}/{}".format( - Path(__file__).stem, conf["model_detector_name"] - ), - ) - model_descriptor_path = self._download_model( - repo_id=MODEL_REPO_ID, - filename="{}/{}".format( - Path(__file__).stem, conf["model_descriptor_name"] - ), - ) - logger.info("Loaded DarkFeat model: {}".format(model_detector_path)) - self.normalizer = transforms.Normalize( - mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225] - ) - - # load the model - weights_detector = torch.load(model_detector_path, map_location="cpu") - weights_descriptor = torch.load( - model_descriptor_path, map_location="cpu" - ) - self.detector = dedode_detector_L( - weights=weights_detector, device=device - ) - self.descriptor = dedode_descriptor_B( - weights=weights_descriptor, device=device - ) + self.model = KF.DeDoDe.from_pretrained(detector_weights=conf["detector_model"], descriptor_weights=conf["descriptor_model"]) logger.info("Load DeDoDe model done.") def _forward(self, data): @@ -68,29 +36,12 @@ def _forward(self, data): image shape: N x C x H x W color mode: RGB """ - img0 = self.normalizer(data["image"].squeeze()).float()[None] - H_A, W_A = img0.shape[2:] - - # step 1: detect keypoints - detections_A = None - batch_A = {"image": img0} - if self.conf["dense"]: - detections_A = self.detector.detect_dense(batch_A) - else: - detections_A = self.detector.detect( - batch_A, num_keypoints=self.conf["max_keypoints"] - ) - keypoints_A, P_A = detections_A["keypoints"], detections_A["confidence"] + img0 = data["image"].float() - # step 2: describe keypoints - # dim: 1 x N x 256 - description_A = self.descriptor.describe_keypoints( - batch_A, keypoints_A - )["descriptions"] - keypoints_A = to_pixel_coords(keypoints_A, H_A, W_A) + coords, scores, descriptions = self.model(img0, n=self.conf["max_keypoints"]) return { - "keypoints": keypoints_A, # 1 x N x 2 - "descriptors": description_A.permute(0, 2, 1), # 1 x 256 x N - "scores": P_A, # 1 x N + "keypoints": coords, # 1 x N x 2 + "descriptors": descriptions.permute(0, 2, 1), # 1 x 512 x N + "scores": scores, # 1 x N } diff --git a/hloc/match_features.py b/hloc/match_features.py index 88a44e2..bc9b2ac 100644 --- a/hloc/match_features.py +++ b/hloc/match_features.py @@ -80,6 +80,38 @@ "force_resize": False, }, }, + "dedode-lightglue": { + "output": "matches-dedode-lightglue", + "model": { + "name": "lightglue", + "match_threshold": 0.2, + "width_confidence": 0.99, # for point pruning + "depth_confidence": 0.95, # for early stopping, + "features": "dedodeg", + }, + "preprocessing": { + "grayscale": True, + "resize_max": 1024, + "dfactor": 8, + "force_resize": False, + }, + }, + "dedodev2-lightglue": { + "output": "matches-dedode-lightglue", + "model": { + "name": "lightglue", + "match_threshold": 0.2, + "width_confidence": 0.99, # for point pruning + "depth_confidence": 0.95, # for early stopping, + "features": "dedodeg-v2", + }, + "preprocessing": { + "grayscale": True, + "resize_max": 1024, + "dfactor": 8, + "force_resize": False, + }, + }, "sift-lightglue": { "output": "matches-sift-lightglue", "model": { diff --git a/hloc/matchers/lightglue.py b/hloc/matchers/lightglue.py index 975b554..e778867 100644 --- a/hloc/matchers/lightglue.py +++ b/hloc/matchers/lightglue.py @@ -1,14 +1,9 @@ -import sys -from pathlib import Path +import torch +import kornia.feature as KF from .. import MODEL_REPO_ID, logger from ..utils.base_model import BaseModel -lightglue_path = Path(__file__).parent / "../../third_party/LightGlue" -sys.path.append(str(lightglue_path)) -from lightglue import LightGlue as LG - - class LightGlue(BaseModel): default_conf = { "match_threshold": 0.2, @@ -33,22 +28,17 @@ class LightGlue(BaseModel): ] def _init(self, conf): - logger.info("Loading lightglue model, {}".format(conf["model_name"])) - model_path = self._download_model( - repo_id=MODEL_REPO_ID, - filename="{}/{}".format( - Path(__file__).stem, self.conf["model_name"] - ), - ) - conf["weights"] = str(model_path) - conf["filter_threshold"] = conf["match_threshold"] - self.net = LG(**conf) + features = conf.pop("features") + logger.info("Loading lightglue model, {}".format(features)) + + self.net = KF.LightGlue(features, **conf) logger.info("Load lightglue model done.") def _forward(self, data): input = {} input["image0"] = { - "image": data["image0"], + # "image": data["image0"], + "image_size": torch.tensor(data["image0"].shape[-2:][::-1])[None], "keypoints": data["keypoints0"], "descriptors": data["descriptors0"].permute(0, 2, 1), } @@ -58,7 +48,8 @@ def _forward(self, data): input["image0"] = {**input["image0"], "oris": data["oris0"]} input["image1"] = { - "image": data["image1"], + # "image": data["image1"], + "image_size": torch.tensor(data["image1"].shape[-2:][::-1])[None], "keypoints": data["keypoints1"], "descriptors": data["descriptors1"].permute(0, 2, 1), } diff --git a/ui/config.yaml b/ui/config.yaml index eb6b609..86fac38 100644 --- a/ui/config.yaml +++ b/ui/config.yaml @@ -179,9 +179,31 @@ matcher_zoo: paper: https://arxiv.org/abs/2404.19174 project: null display: false - dedode: + dedodev2: matcher: Dual-Softmax - feature: dedode + feature: dedodeg-v2 + dense: false + info: + name: DeDoDeV2 #dispaly name + source: "3DV 2024" + github: https://github.com/Parskatt/DeDoDe + paper: https://arxiv.org/abs/2308.08479 + project: null + display: true + dedode+lightglue: + matcher: dedode-lightglue + feature: dedodeg + dense: false + info: + name: DeDoDe #dispaly name + source: "3DV 2024" + github: https://github.com/Parskatt/DeDoDe + paper: https://arxiv.org/abs/2308.08479 + project: null + display: true + dedodev2+lightglue: + matcher: dedode-lightglue + feature: dedodeg-v2 dense: false info: name: DeDoDe #dispaly name