diff --git a/.gitmodules b/.gitmodules index 43c0c1d..405046a 100644 --- a/.gitmodules +++ b/.gitmodules @@ -25,9 +25,6 @@ [submodule "third_party/LightGlue"] path = third_party/LightGlue url = https://github.com/cvg/LightGlue.git -[submodule "third_party/DeDoDe"] - path = third_party/DeDoDe - url = https://github.com/Parskatt/DeDoDe.git [submodule "third_party/RoRD"] path = third_party/RoRD url = https://github.com/Vincentqyw/RoRD.git diff --git a/hloc/extract_features.py b/hloc/extract_features.py index 1fc5de9..6431038 100644 --- a/hloc/extract_features.py +++ b/hloc/extract_features.py @@ -287,10 +287,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..4bb52c3 100644 --- a/hloc/extractors/dedode.py +++ b/hloc/extractors/dedode.py @@ -3,25 +3,20 @@ import torch import torchvision.transforms as transforms +from kornia import feature as KF 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 - device = torch.device("cuda" if torch.cuda.is_available() else "cpu") 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,33 +27,9 @@ 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.") @@ -68,29 +39,14 @@ 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:] + img0 = data["image"].float() - # 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"] - - # 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 cb75253..0a66eb2 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, + }, + }, "aliked-lightglue": { "output": "matches-aliked-lightglue", "model": { diff --git a/hloc/matchers/lightglue.py b/hloc/matchers/lightglue.py index 975b554..1dfcbe4 100644 --- a/hloc/matchers/lightglue.py +++ b/hloc/matchers/lightglue.py @@ -1,13 +1,9 @@ -import sys -from pathlib import Path +import kornia.feature as KF +import torch 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 = { @@ -33,22 +29,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 +49,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/third_party/DeDoDe b/third_party/DeDoDe deleted file mode 160000 index 0928deb..0000000 --- a/third_party/DeDoDe +++ /dev/null @@ -1 +0,0 @@ -Subproject commit 0928debd09de6e814c90647a639a7219437bf49a diff --git a/ui/config.yaml b/ui/config.yaml index ffe68f8..8db1610 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