diff --git a/README.md b/README.md index 24af0eb..8f1edd9 100644 --- a/README.md +++ b/README.md @@ -108,18 +108,10 @@ Running these gave me `ScanNet-1500: [34.0, 56.5, 73.9]`, and `Mega-1500: [62.8, ## Fused local correlation kernel -Include the `--extra fused-local-corr` flag as: -```bash -uv sync --extra fused-local-corr -``` -or -```bash -uv pip install romav2[fused-local-corr] -``` -or -```bash -uv add romav2[fused-local-corr] -``` +On Linux the `fused-local-corr` kernel is a core dependency and is installed +automatically by the Setup/Install step above (`uv sync` / `uv pip install -e .`) +— no extra flag is required. On non-Linux platforms it is skipped and RoMa falls +back to the pure-PyTorch local correlation. ## Settings By twiddling with some different settings you may reach better results on your task of interest. diff --git a/pyproject.toml b/pyproject.toml index c17a9e8..baf33e7 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,6 +1,6 @@ [project] name = "romav2" -version = "2.0.0" +version = "2.0.1" description = "RoMa v2: Harder Better Faster Denser Feature Matching" readme = "README.md" authors = [ @@ -8,13 +8,13 @@ authors = [ ] requires-python = ">=3.10" dependencies = [ - "dataclasses>=0.8", "einops>=0.8.1", "pillow>=12.0.0", "rich>=14.2.0", "torch", "torchvision>=0.23.0", "tqdm>=4.67.1", + "fused-local-corr ; sys_platform == 'linux'", ] [build-system] @@ -22,9 +22,6 @@ requires = ["uv_build>=0.8.15,<0.9.0"] build-backend = "uv_build" [project.optional-dependencies] -fused-local-corr = [ - "fused-local-corr ; sys_platform == 'linux'", -] eval = [ "kornia>=0.8.2", "matplotlib>=3.10.7", diff --git a/scripts/batch_match.py b/scripts/batch_match.py index a2f620f..72fd085 100644 --- a/scripts/batch_match.py +++ b/scripts/batch_match.py @@ -744,6 +744,17 @@ def main(): default=True, help='Disable bidirectional matching (only A->B; default is bidirectional)' ) + parser.add_argument( + '--no-compile', + dest='compile', + action='store_false', + default=True, + help='Disable torch.compile. Compiling is enabled by default for ' + 'throughput (match() runs at fixed resolution, so the graph is ' + 'compiled once and reused across all pairs). Disable it when ' + 'profiling/tracing, since torch.compile graph capture hides the ' + 'torch.profiler.record_function annotations.' + ) parser.add_argument( '--calibration-pairs', type=int, @@ -781,10 +792,10 @@ def main(): pairs = optimize_pair_order(pairs) # Initialize model - logger.info(f"Initializing RoMaV2 with setting: {args.setting}") - #Use compile=False for full tracing - #model = RoMaV2(RoMaV2.Cfg(compile=False)) - model = RoMaV2() + logger.info(f"Initializing RoMaV2 with setting: {args.setting} (compile={args.compile})") + # compile=True is the throughput path; pass --no-compile for full tracing + # (torch.compile hides the profiler record_function annotations). + model = RoMaV2(RoMaV2.Cfg(compile=args.compile)) model.apply_setting(args.setting) model.bidirectional = args.bidirectional model.eval() diff --git a/src/romav2/refiner.py b/src/romav2/refiner.py index 8476685..666683d 100644 --- a/src/romav2/refiner.py +++ b/src/romav2/refiner.py @@ -169,8 +169,9 @@ def forward( ) # Corr in other means take a kxk grid around the predicted coordinate in other image f_A_bdhw = f_A.permute(0, 3, 1, 2) - f_B_bdhw = f_BA.permute(0, 3, 1, 2) - d = torch.cat((f_A_bdhw, f_B_bdhw, emb_in_displacement), dim=1) + f_B_bdhw = f_B.permute(0, 3, 1, 2) + f_BA_bdhw = f_BA.permute(0, 3, 1, 2) + d = torch.cat((f_A_bdhw, f_BA_bdhw, emb_in_displacement), dim=1) if self.cfg.local_corr_radius is not None: local_corr = local_correlation( f_A_bdhw, @@ -180,7 +181,6 @@ def forward( scale_factor=scale_factor, ) d = torch.cat((d, local_corr), dim=1) - # d = torch.cat((f_A_bdhw, f_B_bdhw, emb_in_displacement, local_corr), dim=1) if self.cfg.channels_last: d = d.to(memory_format=torch.channels_last) z = self.block1(d) diff --git a/src/romav2/romav2.py b/src/romav2/romav2.py index 1a8963e..b1927be 100644 --- a/src/romav2/romav2.py +++ b/src/romav2/romav2.py @@ -78,7 +78,7 @@ class Cfg: anchor_width: int = 512 anchor_height: int = 512 setting: Setting = "precise" - compile: bool = True + compile: bool = False name: str = "RoMa v2" # settings @@ -97,7 +97,8 @@ def __init__(self, cfg: Cfg | None = None): cfg = RoMaV2.Cfg() weights = torch.hub.load_state_dict_from_url( - "https://github.com/Parskatt/RoMaV2/releases/download/weights/romav2.pt" + "https://github.com/Parskatt/RoMaV2/releases/download/v2.0.1/romav2.0.1.pt", + map_location=device ) self.f = Descriptor(cfg.descriptor) self.matcher = Matcher(cfg.matcher)