MambaglueAachen

archive_old_feats.sh

@@ -0,0 +1,46 @@
+#!/usr/bin/env bash
+# Archive the old open-SuperPoint feature caches (and stale matches/sfm/results that
+# depended on them) instead of deleting, then leave the run dirs clean for the
+# corrected superpoint_mg_aachen re-run.
+set -euo pipefail
+
+REPO=/home/ubuntu/work/MambaGlue/Hierarchical-Localization
+cd "$REPO"
+
+STAMP=$(date +%Y%m%d_%H%M%S)
+ARCHIVE="outputs/_archive_open_sp_${STAMP}"
+mkdir -p "$ARCHIVE"
+echo "Archiving stale artifacts to: $ARCHIVE"
+
+for tag in nc2 nc3 nc4 mg; do
+  src="outputs/aachen_v1.1_${tag}"
+  [[ -d "$src" ]] || { echo "  (skip $tag: $src not found)"; continue; }
+
+  dst="$ARCHIVE/aachen_v1.1_${tag}"
+  mkdir -p "$dst"
+
+  # move anything tied to the old open-SuperPoint features:
+  #   - the 1024/1600 open feature h5s
+  #   - matches, sfm models, results, retrieval pairs built on top of them
+  #   - the local-feature pairs are matcher-independent but cheap; archive too for a clean slate
+  shopt -s nullglob
+  for f in \
+    "$src"/feats-superpoint-open-*.h5 \
+    "$src"/matches-superpoint-*.h5 \
+    "$src"/global-feats-*.h5 \
+    "$src"/pairs-query-netvlad*.txt \
+    "$src"/pairs-db-covis*.txt \
+    "$src"/sfm_superpoint-open+* \
+    "$src"/Aachen-v1.1_hloc_superpoint-open+*; do
+    echo "  mv $f"
+    mv "$f" "$dst"/
+  done
+  shopt -u nullglob
+done
+
+echo "Done. Archived under $ARCHIVE"
+echo "Remaining in run dirs (should be empty or only new artifacts):"
+for tag in nc2 nc3 nc4 mg; do
+  echo "--- outputs/aachen_v1.1_${tag} ---"
+  ls -1 "outputs/aachen_v1.1_${tag}" 2>/dev/null || true
+done

hloc/extract_features.py

@@ -146,6 +146,26 @@
         "model": {"name": "megaloc"},
         "preprocessing": {"resize_max": 1024},
     },
+    "superpoint_open_aachen": {
+        "output": "feats-superpoint-open-n4096-r1024",
+        "model": {
+            "name": "superpoint_open",
+            "nms_radius": 3,
+            "max_num_keypoints": 4096,
+            "detection_threshold": 0.005,
+        },
+        "preprocessing": {"grayscale": True, "resize_max": 1024},
+    },
+    "superpoint_mg_aachen": {
+    "output": "feats-superpoint-mg-n2048-r1600",
+    "model": {
+        "name": "superpoint",          # hloc's native MagicLeap SuperPoint
+        "nms_radius": 3,               # matches training
+        "max_keypoints": 2048,         # test-time budget; see note below
+        "keypoint_threshold": 0.0,     # matches training's detection_threshold: 0.0
+    },
+    "preprocessing": {"grayscale": True, "resize_max": 1600},
+    },
 }
 
 

hloc/extractors/superpoint_open.py

@@ -0,0 +1,64 @@
+"""hloc extractor wrapper for glue-factory's *open* SuperPoint.
+
+Place this file at:  Hierarchical-Localization/hloc/extractors/superpoint_open.py
+
+Why: glue-factory trains on the rpautrat open SuperPoint
+(gluefactory.models.extractors.superpoint_open), but hloc ships the original
+MagicLeap SuperPoint by default. Both are 256-d, but the descriptor
+distributions differ, which biases your Aachen numbers. This wrapper runs the
+exact SuperPoint your LateMambaGlue checkpoint was trained against.
+
+Requires glue-factory importable (`pip install -e .` in the glue-factory repo,
+or uncomment the sys.path block). The open SuperPoint weights auto-download on
+first instantiation, so the first run needs network access.
+"""
+
+import sys
+from pathlib import Path
+
+from omegaconf import OmegaConf
+
+from ..utils.base_model import BaseModel
+
+# --- If glue-factory is NOT pip-installed: ---
+# GLUEFACTORY_ROOT = Path("/abs/path/to/glue-factory")
+# sys.path.append(str(GLUEFACTORY_ROOT))
+
+from gluefactory.models.extractors.superpoint_open import (  # noqa: E402
+    SuperPoint as _GFSuperPoint,
+)
+
+
+class SuperPointOpen(BaseModel):
+    default_conf = {
+        "nms_radius": 3,
+        "max_num_keypoints": 4096,    # localization wants more kpts than the 512/1024 used in training
+        "detection_threshold": 0.005,
+        "remove_borders": 4,
+        "descriptor_dim": 256,
+        "force_num_keypoints": False,  # variable kpt count per image for localization
+    }
+    required_inputs = ["image"]
+
+    def _init(self, conf):
+        gf_conf = OmegaConf.create({
+            "nms_radius": conf["nms_radius"],
+            "max_num_keypoints": conf["max_num_keypoints"],
+            "detection_threshold": conf["detection_threshold"],
+            "remove_borders": conf["remove_borders"],
+            "descriptor_dim": conf["descriptor_dim"],
+            "force_num_keypoints": conf["force_num_keypoints"],
+        })
+        self.net = _GFSuperPoint(gf_conf)
+        self.net.eval()
+
+    def _forward(self, data):
+        # hloc passes a grayscale image [B, 1, H, W] in [0, 1]; glue-factory wants the same.
+        pred = self.net({"image": data["image"]})
+        # glue-factory out: keypoints [B,N,2] (x,y), keypoint_scores [B,N], descriptors [B,N,D]
+        # hloc expects:     keypoints [B,N,2],       scores [B,N],          descriptors [B,D,N]
+        return {
+            "keypoints": pred["keypoints"],
+            "scores": pred["keypoint_scores"],
+            "descriptors": pred["descriptors"].transpose(-1, -2).contiguous(),
+        }

hloc/match_features.py

@@ -85,6 +85,17 @@
         "output": "matches-adalam",
         "model": {"name": "adalam"},
     },
+    "superpoint+latemambaglue": {
+        "output": "matches-superpoint-latemambaglue",
+        "model": {
+            "name": "latemambaglue",
+            "features": "superpoint",
+            "checkpoint": "/home/ubuntu/work/MambaGlue/glue-factory/outputs/training/latemambaglueMDnc3/checkpoint_best.tar",
+            "n_layers": 9,
+            "n_cross_layers": 3,
+            "filter_threshold": 0.1,
+        },
+    },
 }
 
 

hloc/matchers/latemambaglue.py

@@ -0,0 +1,101 @@
+"""hloc matcher wrapper for LateMambaGlue (trained in glue-factory).
+
+Place this file at:  Hierarchical-Localization/hloc/matchers/latemambaglue.py
+
+Requires glue-factory to be importable. Easiest: from the glue-factory repo
+root run `pip install -e .` so `import gluefactory` works. Otherwise uncomment
+the sys.path block below and point it at your glue-factory checkout.
+"""
+
+import sys
+from pathlib import Path
+
+import torch
+from omegaconf import OmegaConf
+
+from ..utils.base_model import BaseModel
+
+# --- If glue-factory is NOT pip-installed, expose it on sys.path: ---
+#GLUEFACTORY_ROOT = Path("/home/ubuntu/work/MambaGlue/glue-factory")
+#sys.path.append(str(GLUEFACTORY_ROOT))
+
+from gluefactory.models.matchers.latemambaglue import (  # noqa: E402
+    LateMambaGlue as _GFLateMambaGlue,
+)
+
+
+class LateMambaGlue(BaseModel):
+    default_conf = {
+        "checkpoint": None,      # REQUIRED: path to glue-factory checkpoint_best.tar
+        "n_layers": 9,
+        "n_cross_layers": 3,     # match the trained variant: nc2/nc3/nc4 -> 2/3/4
+        "filter_threshold": 0.1,
+        "flash": True,
+        "checkpointed": False,   # gradient checkpointing OFF at inference
+        "features": "superpoint",
+    }
+    # hloc feeds these keys to _forward (descriptors are [B, D, N] in hloc)
+    required_inputs = [
+        "image0", "keypoints0", "scores0", "descriptors0",
+        "image1", "keypoints1", "scores1", "descriptors1",
+    ]
+
+    def _init(self, conf):
+        model_conf = OmegaConf.create({
+            "n_layers": conf["n_layers"],
+            "n_cross_layers": conf["n_cross_layers"],
+            "filter_threshold": conf["filter_threshold"],
+            "flash": conf["flash"],
+            "checkpointed": conf["checkpointed"],
+        })
+        self.net = _GFLateMambaGlue(model_conf)
+
+        ckpt = conf["checkpoint"]
+        assert ckpt is not None, "Set the 'checkpoint' field to your .tar file"
+        state = torch.load(ckpt, map_location="cpu")
+        sd = state["model"] if isinstance(state, dict) and "model" in state else state
+
+        # glue-factory saves the whole two_view_pipeline; matcher weights are
+        # prefixed with "matcher." — strip it to get the bare matcher state dict.
+        matcher_sd = {
+            k[len("matcher."):]: v for k, v in sd.items() if k.startswith("matcher.")
+        }
+        if not matcher_sd:            # already a bare matcher checkpoint
+            matcher_sd = sd
+
+        missing, unexpected = self.net.load_state_dict(matcher_sd, strict=False)
+        if missing:
+            print(f"[latemambaglue] {len(missing)} missing keys, e.g. {missing[:3]}")
+        if unexpected:
+            print(f"[latemambaglue] {len(unexpected)} unexpected keys, e.g. {unexpected[:3]}")
+        self.net.eval()
+
+    def _forward(self, data):
+        # hloc descriptors: [B, D, N]  ->  glue-factory wants [B, N, D]
+        desc0 = data["descriptors0"].transpose(-1, -2).contiguous().float()
+        desc1 = data["descriptors1"].transpose(-1, -2).contiguous().float()
+
+        dev = data["keypoints0"].device
+        h0, w0 = data["image0"].shape[-2:]
+        h1, w1 = data["image1"].shape[-2:]
+        size0 = torch.tensor([[w0, h0]], dtype=torch.float32, device=dev)
+        size1 = torch.tensor([[w1, h1]], dtype=torch.float32, device=dev)
+
+        gf_data = {
+            "keypoints0": data["keypoints0"].float(),
+            "keypoints1": data["keypoints1"].float(),
+            "descriptors0": desc0,
+            "descriptors1": desc1,
+            # provide image size in both common conventions so keypoint
+            # normalization is exact regardless of which one your forward reads
+            "image_size0": size0,
+            "image_size1": size1,
+            "view0": {"image_size": size0},
+            "view1": {"image_size": size1},
+        }
+
+        pred = self.net(gf_data)
+        return {
+            "matches0": pred["matches0"],
+            "matching_scores0": pred["matching_scores0"],
+        }

hloc/matchers/mambaglue.py

@@ -0,0 +1,94 @@
+"""hloc matcher wrapper for glue-factory's original MambaGlue.
+
+Place at:  Hierarchical-Localization/hloc/matchers/mambaglue.py
+
+This is the "normal MambaGlue" baseline -- a different class from LateMambaGlue
+(no Context/Interaction split, no n_cross_layers), with its own checkpoint
+(e.g. outputs/training/sp+mg_megadepth/checkpoint_best.tar). I/O handling is the
+same as the LateMambaGlue wrapper since both follow glue-factory's matcher
+convention.
+
+Requires glue-factory importable (`pip install -e .` in the glue-factory repo).
+"""
+
+import sys
+from pathlib import Path
+
+import torch
+from omegaconf import OmegaConf
+
+from ..utils.base_model import BaseModel
+
+# --- If glue-factory is NOT pip-installed: ---
+# GLUEFACTORY_ROOT = Path("/home/ubuntu/work/MambaGlue/glue-factory")
+# sys.path.append(str(GLUEFACTORY_ROOT))
+
+from gluefactory.models.matchers.mambaglue import (  # noqa: E402
+    MambaGlue as _GFMambaGlue,
+)
+
+
+class MambaGlue(BaseModel):
+    default_conf = {
+        "checkpoint": None,       # REQUIRED: path to glue-factory checkpoint_best.tar
+        "filter_threshold": 0.1,
+        "n_layers": None,         # None => use MambaGlue's own default architecture
+        "features": "superpoint",
+    }
+    required_inputs = [
+        "image0", "keypoints0", "scores0", "descriptors0",
+        "image1", "keypoints1", "scores1", "descriptors1",
+    ]
+
+    def _init(self, conf):
+        # keep the override set minimal: only keys certain to be in MambaGlue's
+        # default_conf, so glue-factory's struct-checked merge can't reject them
+        model_conf = {"filter_threshold": conf["filter_threshold"]}
+        if conf.get("n_layers") is not None:
+            model_conf["n_layers"] = conf["n_layers"]
+        self.net = _GFMambaGlue(OmegaConf.create(model_conf))
+
+        ckpt = conf["checkpoint"]
+        assert ckpt is not None, "Set the 'checkpoint' field to your .tar file"
+        state = torch.load(ckpt, map_location="cpu")
+        sd = state["model"] if isinstance(state, dict) and "model" in state else state
+
+        # glue-factory saves the two_view_pipeline; matcher weights prefixed "matcher."
+        matcher_sd = {
+            k[len("matcher."):]: v for k, v in sd.items() if k.startswith("matcher.")
+        }
+        if not matcher_sd:
+            matcher_sd = sd
+
+        missing, unexpected = self.net.load_state_dict(matcher_sd, strict=False)
+        if missing:
+            print(f"[mambaglue] {len(missing)} missing keys, e.g. {missing[:3]}")
+        if unexpected:
+            print(f"[mambaglue] {len(unexpected)} unexpected keys, e.g. {unexpected[:3]}")
+        self.net.eval()
+
+    def _forward(self, data):
+        desc0 = data["descriptors0"].transpose(-1, -2).contiguous().float()
+        desc1 = data["descriptors1"].transpose(-1, -2).contiguous().float()
+
+        dev = data["keypoints0"].device
+        h0, w0 = data["image0"].shape[-2:]
+        h1, w1 = data["image1"].shape[-2:]
+        size0 = torch.tensor([[w0, h0]], dtype=torch.float32, device=dev)
+        size1 = torch.tensor([[w1, h1]], dtype=torch.float32, device=dev)
+
+        gf_data = {
+            "keypoints0": data["keypoints0"].float(),
+            "keypoints1": data["keypoints1"].float(),
+            "descriptors0": desc0,
+            "descriptors1": desc1,
+            "image_size0": size0,
+            "image_size1": size1,
+            "view0": {"image_size": size0},
+            "view1": {"image_size": size1},
+        }
+        pred = self.net(gf_data)
+        return {
+            "matches0": pred["matches0"],
+            "matching_scores0": pred["matching_scores0"],
+        }