diff --git a/.github/actions/uv-build/action.yml b/.github/actions/uv-build/action.yml new file mode 100644 index 0000000..93963a7 --- /dev/null +++ b/.github/actions/uv-build/action.yml @@ -0,0 +1,23 @@ +name: uv-build +description: Build the project with uv and run smoke tests +inputs: + python-version: + description: Python version to install + required: true +runs: + using: composite + steps: + - name: Install uv + uses: astral-sh/setup-uv@v6 + - name: Install Python + shell: bash + run: uv python install ${{ inputs.python-version }} + - name: Build + shell: bash + run: uv build + - name: Smoke test (wheel) + shell: bash + run: uv run --isolated --no-project --with dist/*.whl tests/test_smoke.py + - name: Smoke test (source distribution) + shell: bash + run: uv run --isolated --no-project --with dist/*.tar.gz tests/test_smoke.py diff --git a/.github/workflows/build.yml b/.github/workflows/build.yml new file mode 100644 index 0000000..8716c23 --- /dev/null +++ b/.github/workflows/build.yml @@ -0,0 +1,22 @@ +name: Build + +on: + push: + branches: + - main + pull_request: + branches: + - main +jobs: + run: + runs-on: ubuntu-latest + strategy: + matrix: + python-version: ['3.10', '3.11', '3.12', '3.13'] + steps: + - name: Checkout + uses: actions/checkout@v5 + - name: Build and smoke test + uses: ./.github/actions/uv-build + with: + python-version: ${{ matrix.python-version }} diff --git a/.github/workflows/publish.yml b/.github/workflows/publish.yml new file mode 100644 index 0000000..401234a --- /dev/null +++ b/.github/workflows/publish.yml @@ -0,0 +1,41 @@ +name: PyPI Publish + +on: + push: + tags: + # Publish on any tag starting with a `v`, e.g., v0.1.0 + - v* + workflow_dispatch: + release: + types: [published] +jobs: + build: + runs-on: ubuntu-latest + strategy: + matrix: + python-version: ['3.10', '3.11', '3.12', '3.13'] + steps: + - name: Checkout + uses: actions/checkout@v5 + - name: Build and smoke test + uses: ./.github/actions/uv-build + with: + python-version: ${{ matrix.python-version }} + + publish: + runs-on: ubuntu-latest + needs: build + environment: + name: pypi + permissions: + id-token: write + contents: read + steps: + - name: Checkout + uses: actions/checkout@v5 + - name: Build and smoke test + uses: ./.github/actions/uv-build + with: + python-version: '3.12' + - name: Publish + run: uv publish diff --git a/README.md b/README.md index c4d01d0..cc80c18 100644 --- a/README.md +++ b/README.md @@ -17,42 +17,53 @@

## Overview -HardMatch is an extremely difficult image matching benchmark featuring 1,000 hand annotated image pairs. The benchmark is released as part of the [LoMa](https://arxiv.org/abs/2604.04931) paper. +HardMatch is an extremely difficult image matching benchmark featuring 1,000 hand annotated image pairs. The benchmark is released as part of the [LoMa](https://arxiv.org/abs/2604.04931) paper (ECCV 2026, Oral). The best way to visualize the data is through our online [viewer](https://www.davnords.com/loma/hardmatch). ## Updates +- [August 5, 2026] PyPi package released and code made more accesible. - [June 27, 2026] Initial dataset release following ECCV 2026 acceptance. +## Setup/Install +In your python environment (tested on Linux python 3.12): +```bash +uv add hardmatch +``` +or +```bash +pip install hardmatch +``` + ## How to Use ```python from hardmatch import HardMatchBenchmark matcher = YourFancyMatcher() result = HardMatchBenchmark().benchmark(matcher) ``` -We additionally provide an example of the matching API through `demo.py` which uses [SuperPoint](https://github.com/rpautrat/SuperPoint) + [LightGlue](https://github.com/cvg/LightGlue). This defaults to evaluating on the 900 test pairs. There are also 100 validation pairs. To test this demo, just run: +`matcher` needs to satisfy the `BenchmarkMatcher` protocol ([`hardmatch.types.BenchmarkMatcher`](src/hardmatch/types.py)): + +```python +class YourFancyMatcher: + offset = 0.0 # pixel offset subtracted from returned keypoints before scoring + + def match(self, img_A_path: str, img_B_path: str) -> tuple[np.ndarray, np.ndarray]: + """Return (kpts_A, kpts_B): corresponding (N, 2) keypoints in pixel + coordinates of the original image files, one row per match.""" + ... + return kpts_A, kpts_B +``` +We additionally provide two example matchers through `demo.py`: [SuperPoint](https://github.com/rpautrat/SuperPoint) + [LightGlue](https://github.com/cvg/LightGlue) (`SPLG`), and [LoMa](https://github.com/davnords/loma) (`LoMa`). This defaults to evaluating on the 900 test pairs. There are also 100 validation pairs. To try the demo, first clone the repo, and then run: ```bash -uv run demo.py +uv sync --extra baselines +uv run demo.py --matcher loma +# Expected result: mAA_10px: 0.5061 ``` +Note, the results differs a tiny bit from the results in the [paper](https://arxiv.org/abs/2604.04931). This is because after submission we had change the dataset a tiny bit (around 10 pairs) for licensing issues. ## Download Running the benchmark will automatically download the data (660MB). You can also manually download it [here](https://github.com/davnords/storage/releases/download/hardmatch/hardmatch.zip). -## Setup/Install -In your python environment (tested on Linux python 3.12), run: -```bash -uv pip install -e . -``` -or -```bash -uv sync -``` - -## Checklist -- [ ] Provide expected results for SP+LG and LoMa. -- [ ] Remove Kornia dependency and split eval and dataset dependencies. -- [ ] Make into easy-to-use PyPi package. - ## License All our code is MIT license. The pairs are scraped from [WikiMedia Commons](https://commons.wikimedia.org/wiki/Main_Page). As such, each pair has its own license that you can find in the data. They are generally permissive. @@ -60,7 +71,7 @@ All our code is MIT license. The pairs are scraped from [WikiMedia Commons](http Our evaluation technique builds on [WxBS](https://arxiv.org/abs/1504.06603). ## BibTeX -If you find our models useful, please consider citing our paper! +If you find our dataset useful, please consider citing our paper! ``` @inproceedings{nordstrom2026loma, title={LoMa: Local Feature Matching Revisited}, diff --git a/demo.py b/demo.py index 2e9a49d..d05615f 100644 --- a/demo.py +++ b/demo.py @@ -1,4 +1,11 @@ -from baselines import SPLG +import argparse + +from baselines import SPLG, LoMa from hardmatch import HardMatchBenchmark -result = HardMatchBenchmark().benchmark(SPLG()) \ No newline at end of file +parser = argparse.ArgumentParser() +parser.add_argument("--matcher", choices=["splg", "loma"], default="splg") +args = parser.parse_args() + +matcher = SPLG() if args.matcher == "splg" else LoMa() +result = HardMatchBenchmark().benchmark(matcher) diff --git a/pyproject.toml b/pyproject.toml index 12b54b3..538d5a6 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -13,8 +13,25 @@ dependencies = [ "torchvision>=0.23.0", "tqdm>=4.67.1", "einops>=0.8.1", + "numpy", + "pillow", + "opencv-python>=4.12.0.88", +] + +[project.urls] +Homepage = "https://github.com/davnords/HardMatch" +Repository = "https://github.com/davnords/HardMatch" +Paper = "https://arxiv.org/abs/2604.04931" + +[project.optional-dependencies] +baselines = [ + "lomatch>=1.0.0", "kornia", ] +dev = [ + "ruff", + "pyright", +] [build-system] requires = ["uv_build>=0.8.15,<0.9.0"] diff --git a/src/baselines/__init__.py b/src/baselines/__init__.py index 355d8e0..776f9b1 100644 --- a/src/baselines/__init__.py +++ b/src/baselines/__init__.py @@ -1,12 +1,25 @@ +from dataclasses import replace +from typing import Literal + import torch + from .lightglue import LightGlue from .superpoint import SuperPoint, load_image, rbd +device = "cuda" if torch.cuda.is_available() else "cpu" + + class SPLG: - device = "cuda" if torch.cuda.is_available() else "cpu" - superpoint = SuperPoint(max_num_keypoints=4096).to(device).eval() - lightglue = LightGlue(features="superpoint", depth_confidence=-1, filter_threshold=0.0, width_confidence=-1).to(device).eval() - offset = 0.0 # Pixel offset + """SuperPoint + LightGlue baseline matcher. See `hardmatch.types.BenchmarkMatcher`.""" + + offset = 0.0 # Pixel offset + + def __init__(self, device: str = device): + self.device = device + self.superpoint = SuperPoint(max_num_keypoints=4096).to(device).eval() + self.lightglue = LightGlue( + features="superpoint", depth_confidence=-1, filter_threshold=0.0, width_confidence=-1 + ).to(device).eval() def match(self, img_A_path: str, img_B_path: str): # Reference: https://colab.research.google.com/github/cvg/LightGlue/blob/main/demo.ipynb#scrollTo=pjjHazDhC1tU @@ -24,3 +37,27 @@ def match(self, img_A_path: str, img_B_path: str): kpts0, kpts1, matches = feats0["keypoints"], feats1["keypoints"], matches01["matches"] m_kpts0, m_kpts1 = kpts0[matches[..., 0]], kpts1[matches[..., 1]] return m_kpts0.cpu().numpy(), m_kpts1.cpu().numpy() + + +class LoMa: + """LoMa baseline matcher, from https://github.com/davnords/loma (`pip install lomatch`). + See `hardmatch.types.BenchmarkMatcher`. + """ + + offset = 0.0 # Pixel offset + + def __init__( + self, + size: Literal["B", "B128", "L", "G"] = "B", + device: str = device, + num_keypoints: int = 4096, + ): + from loma import LoMa as _LoMa, LoMaB, LoMaB128, LoMaG, LoMaL + + cfg = {"B": LoMaB, "B128": LoMaB128, "L": LoMaL, "G": LoMaG}[size]() + cfg = replace(cfg, num_keypoints=num_keypoints) + self.device = device + self.model = _LoMa(cfg).to(device).eval() + + def match(self, img_A_path: str, img_B_path: str): + return self.model.match(str(img_A_path), str(img_B_path)) diff --git a/src/baselines/lightglue.py b/src/baselines/lightglue.py index f29ca4c..3719ec2 100644 --- a/src/baselines/lightglue.py +++ b/src/baselines/lightglue.py @@ -17,7 +17,6 @@ FLASH_AVAILABLE = True else: FLASH_AVAILABLE = False - torch.backends.cudnn.deterministic = True @@ -27,7 +26,6 @@ else torch.cuda.amp.custom_fwd(cast_inputs=torch.float32) ) - @AMP_CUSTOM_FWD_F32 def normalize_keypoints( kpts: torch.Tensor, size: Optional[torch.Tensor] = None diff --git a/src/baselines/loma/__init__.py b/src/baselines/loma/__init__.py deleted file mode 100644 index e69de29..0000000 diff --git a/src/baselines/loma/descriptor/__init__.py b/src/baselines/loma/descriptor/__init__.py deleted file mode 100644 index 62ff9d0..0000000 --- a/src/baselines/loma/descriptor/__init__.py +++ /dev/null @@ -1 +0,0 @@ -from .dedode import DeDoDeDescriptor as DeDoDeDescriptor diff --git a/src/baselines/loma/descriptor/dedode.py b/src/baselines/loma/descriptor/dedode.py deleted file mode 100644 index dc2c816..0000000 --- a/src/baselines/loma/descriptor/dedode.py +++ /dev/null @@ -1,407 +0,0 @@ -import logging -from dataclasses import dataclass -from typing import Literal - -import numpy as np -import torch -import torch.nn as nn -import torch.nn.functional as F -import torchvision.models as tvm -from PIL import Image - -from ..device import device, amp_dtype -from ..types import Model - -logger = logging.getLogger(__name__) - - -class DeDoDeDescriptor(Model): - @dataclass(frozen=True) - class Cfg: - arch: Literal["dedode_b", "dedode_g"] = "dedode_b" - compile: bool = True - descriptor_dim: int = 256 - hidden_blocks: int = 5 - - def __init__(self, cfg: Cfg) -> None: - super().__init__() - - if cfg.arch == "dedode_b": - encoder, decoder = dedode_descriptor_B( - descriptor_dim=cfg.descriptor_dim, hidden_blocks=cfg.hidden_blocks - ) - elif cfg.arch == "dedode_g": - encoder, decoder = dedode_descriptor_G( - descriptor_dim=cfg.descriptor_dim, hidden_blocks=cfg.hidden_blocks - ) - else: - raise ValueError(f"Architecture {cfg.arch} not supported") - self.cfg = cfg - self.encoder = encoder - self.decoder = decoder - - self.to(device) - if cfg.compile: - logger.info("Compiling DeDoDeDescriptor...") - self.compile() - logger.info(f"{self.name} initialized.") - - @staticmethod - def load_pretrained(arch: Literal["dedode_b", "dedode_g"]): - cfg = DeDoDeDescriptor.Cfg(arch=arch) - if arch == "dedode_b": - weights = torch.hub.load_state_dict_from_url( - "https://github.com/Parskatt/DeDoDe/releases/download/dedode_pretrained_models/dedode_descriptor_B.pth" - ) - elif arch == "dedode_g": - weights = torch.hub.load_state_dict_from_url( - "https://github.com/Parskatt/DeDoDe/releases/download/dedode_pretrained_models/dedode_descriptor_G.pth" - ) - else: - raise ValueError(f"Architecture {arch} not supported") - model = DeDoDeDescriptor(cfg) - missing_keys, unexpected_keys = model.load_state_dict(weights, strict=False) - assert len(unexpected_keys) == 0, ( - f"Unexpected keys when loading pretrained weights: {unexpected_keys}" - ) - return model - - def forward( - self, - images: torch.Tensor, - ): - features, sizes = self.encoder(images) - descriptions = 0 - context = None - scales = self.decoder.scales - for idx, (feature_map, scale) in enumerate(zip(reversed(features), scales)): - delta_description, context = self.decoder( - feature_map, scale=scale, context=context - ) - descriptions = descriptions + delta_description - if idx < len(scales) - 1: - size = sizes[-(idx + 2)] - descriptions = F.interpolate( - descriptions, size=size, mode="bilinear", align_corners=False - ) - context = F.interpolate( - context, size=size, mode="bilinear", align_corners=False - ) - return descriptions - - @torch.inference_mode() - def describe_keypoints( - self, - images: torch.Tensor, - keypoints: torch.Tensor, - ): - self.train(False) - images = images.to(device) - keypoints = keypoints.to(device) - # TODO is this compile error //je ? - description_grid = self(images) - described_keypoints = F.grid_sample( - description_grid.float(), - keypoints[:, None], - mode="bilinear", - align_corners=False, - )[:, :, 0].mT - return {"descriptions": described_keypoints} - - def read_image(self, im_path, H=784, W=784): - return ( - torch.from_numpy( - np.array(Image.open(im_path).convert("RGB").resize((W, H))) / 255.0 - ) - .permute(2, 0, 1) - .float() - .to(device)[None] - ) - - def describe_keypoints_from_path(self, im_path, keypoints, H=784, W=784): - images = self.read_image(im_path, H=H, W=W) - return self.describe_keypoints(images, keypoints) - - -class Decoder(nn.Module): - def __init__( - self, layers, *args, super_resolution=False, descriptor_dim=1, **kwargs - ) -> None: - super().__init__(*args, **kwargs) - self.layers = layers - self.scales = sorted(list(layers.keys()), key=lambda x: int(x), reverse=True) - self.super_resolution = super_resolution - self.descriptor_dim = descriptor_dim - - def forward(self, features, context=None, scale=None): - if context is not None: - features = torch.cat((features, context), dim=1) - stuff = self.layers[scale](features) - logits, context = ( - stuff[:, : self.descriptor_dim], - stuff[:, self.descriptor_dim :], - ) - return logits, context - - -class ConvRefiner(nn.Module): - def __init__( - self, - in_dim: int, - hidden_dim: int, - out_dim: int, - *, - hidden_blocks: int, - amp: bool, - amp_dtype: torch.dtype, - dw: bool = True, - kernel_size: int = 5, - ): - super().__init__() - self.block1 = self.create_block( - in_dim, - hidden_dim, - dw=False, - kernel_size=1, - ) - self.hidden_blocks = nn.Sequential( - *[ - self.create_block( - hidden_dim, - hidden_dim, - dw=dw, - kernel_size=kernel_size, - ) - for hb in range(hidden_blocks) - ] - ) - self.hidden_blocks = self.hidden_blocks - self.out_conv = nn.Conv2d(hidden_dim, out_dim, 1, 1, 0) - self.amp = amp - self.amp_dtype = amp_dtype - - def create_block( - self, - in_dim, - out_dim, - dw=True, - kernel_size=5, - bias=True, - norm_type=nn.BatchNorm2d, - ): - num_groups = 1 if not dw else in_dim - if dw: - assert out_dim % in_dim == 0, ( - "outdim must be divisible by indim for depthwise" - ) - conv1 = nn.Conv2d( - in_dim, - out_dim, - kernel_size=kernel_size, - stride=1, - padding=kernel_size // 2, - groups=num_groups, - bias=bias, - ) - norm = ( - norm_type(out_dim) - if norm_type is nn.BatchNorm2d - else norm_type(num_channels=out_dim) - ) - relu = nn.ReLU(inplace=True) - conv2 = nn.Conv2d(out_dim, out_dim, 1, 1, 0) - return nn.Sequential(conv1, norm, relu, conv2) - - def forward(self, feats): - b, c, hs, ws = feats.shape - with torch.autocast( - device_type=feats.device.type, enabled=self.amp, dtype=self.amp_dtype - ): - x0 = self.block1(feats) - x = self.hidden_blocks(x0) - x = (x + x0) / 1.4 - x = self.out_conv(x) - return x - - -class VGG(nn.Module): - def __init__(self, size="19", amp=False, amp_dtype=amp_dtype) -> None: - super().__init__() - if size == "11": - self.layers = nn.ModuleList(tvm.vgg11_bn().features[:22]) - elif size == "13": - self.layers = nn.ModuleList(tvm.vgg13_bn().features[:28]) - elif size == "19": - self.layers = nn.ModuleList(tvm.vgg19_bn().features[:40]) - # Maxpool layers: 6, 13, 26, 39 - self.amp = amp - self.amp_dtype = amp_dtype - - def forward(self, x, **kwargs): - with torch.autocast( - device_type=x.device.type, enabled=self.amp, dtype=self.amp_dtype - ): - feats = [] - sizes = [] - for layer in self.layers: - if isinstance(layer, nn.MaxPool2d): - feats.append(x) - sizes.append(x.shape[-2:]) - x = layer(x) - return feats, sizes - - -class FrozenDINOv2(nn.Module): - def __init__(self, amp=True, amp_dtype=amp_dtype, dinov2_weights=None): - super().__init__() - if dinov2_weights is None: - dinov2_weights = torch.hub.load_state_dict_from_url( - "https://dl.fbaipublicfiles.com/dinov2/dinov2_vitl14/dinov2_vitl14_pretrain.pth", - map_location="cpu", - ) - from .transformer import vit_large - - vit_kwargs = dict( - img_size=518, - patch_size=14, - init_values=1.0, - ffn_layer="mlp", - block_chunks=0, - ) - dinov2_vitl14 = vit_large(**vit_kwargs).eval() - dinov2_vitl14.load_state_dict(dinov2_weights) - self.amp = amp - self.amp_dtype = amp_dtype - if self.amp: - dinov2_vitl14 = dinov2_vitl14.to(self.amp_dtype) - for param in dinov2_vitl14.parameters(): - param.requires_grad = False - self.dinov2_vitl14 = dinov2_vitl14 - - @torch.compiler.disable - def forward(self, x): - B, C, H, W = x.shape - with torch.inference_mode(): - dinov2_features_16 = self.dinov2_vitl14.forward_features( - x.to(self.amp_dtype) - ) - features_16 = ( - dinov2_features_16["x_norm_patchtokens"] - .permute(0, 2, 1) - .reshape(B, 1024, H // 14, W // 14) - ) - return [features_16.clone()], [ - (H // 14, W // 14) - ] # clone from inference mode to use in autograd - - -class VGG_DINOv2(nn.Module): - def __init__(self, vgg_kwargs=None, dinov2_kwargs=None): - assert vgg_kwargs is not None and dinov2_kwargs is not None, "Input kwargs pls" - super().__init__() - self.vgg = VGG(**vgg_kwargs) - self.frozen_dinov2 = FrozenDINOv2(**dinov2_kwargs) - - def forward(self, x): - feats_vgg, sizes_vgg = self.vgg(x) - feat_dinov2, size_dinov2 = self.frozen_dinov2(x) - return feats_vgg + feat_dinov2, sizes_vgg + size_dinov2 - - -def dedode_descriptor_B(descriptor_dim: int, hidden_blocks: int = 5): - amp = True - conv_refiner = nn.ModuleDict( - { - "8": ConvRefiner( - 512, - 512, - 256 + descriptor_dim, - hidden_blocks=hidden_blocks, - amp=amp, - amp_dtype=amp_dtype, - ), - "4": ConvRefiner( - 256 + 256, - 256, - 128 + descriptor_dim, - hidden_blocks=hidden_blocks, - amp=amp, - amp_dtype=amp_dtype, - ), - "2": ConvRefiner( - 128 + 128, - 64, - 32 + descriptor_dim, - hidden_blocks=hidden_blocks, - amp=amp, - amp_dtype=amp_dtype, - ), - "1": ConvRefiner( - 64 + 32, - 32, - 1 + descriptor_dim, - hidden_blocks=hidden_blocks, - amp=amp, - amp_dtype=amp_dtype, - ), - } - ) - encoder = VGG(size="19", amp=amp, amp_dtype=amp_dtype) - decoder = Decoder(conv_refiner, descriptor_dim=descriptor_dim) - return encoder, decoder - - -def dedode_descriptor_G( - descriptor_dim: int, dinov2_weights: str | None = None, hidden_blocks: int = 5 -): - amp = True - - conv_refiner = nn.ModuleDict( - { - "14": ConvRefiner( - 1024, - 768, - 512 + descriptor_dim, - hidden_blocks=hidden_blocks, - amp=amp, - amp_dtype=amp_dtype, - ), - "8": ConvRefiner( - 512 + 512, - 512, - 256 + descriptor_dim, - hidden_blocks=hidden_blocks, - amp=amp, - amp_dtype=amp_dtype, - ), - "4": ConvRefiner( - 256 + 256, - 256, - 128 + descriptor_dim, - hidden_blocks=hidden_blocks, - amp=amp, - amp_dtype=amp_dtype, - ), - "2": ConvRefiner( - 128 + 128, - 64, - 32 + descriptor_dim, - hidden_blocks=hidden_blocks, - amp=amp, - amp_dtype=amp_dtype, - ), - "1": ConvRefiner( - 64 + 32, - 32, - 1 + descriptor_dim, - hidden_blocks=hidden_blocks, - amp=amp, - amp_dtype=amp_dtype, - ), - } - ) - vgg_kwargs = dict(size="19", amp=amp, amp_dtype=amp_dtype) - dinov2_kwargs = dict(amp=amp, amp_dtype=amp_dtype, dinov2_weights=dinov2_weights) - encoder = VGG_DINOv2(vgg_kwargs=vgg_kwargs, dinov2_kwargs=dinov2_kwargs) - decoder = Decoder(conv_refiner, descriptor_dim=descriptor_dim) - return encoder, decoder diff --git a/src/baselines/loma/descriptor/transformer/__init__.py b/src/baselines/loma/descriptor/transformer/__init__.py deleted file mode 100644 index 4afa910..0000000 --- a/src/baselines/loma/descriptor/transformer/__init__.py +++ /dev/null @@ -1 +0,0 @@ -from .dinov2 import vit_large as vit_large diff --git a/src/baselines/loma/descriptor/transformer/dinov2.py b/src/baselines/loma/descriptor/transformer/dinov2.py deleted file mode 100644 index 4bb7543..0000000 --- a/src/baselines/loma/descriptor/transformer/dinov2.py +++ /dev/null @@ -1,397 +0,0 @@ -# Copyright (c) Meta Platforms, Inc. and affiliates. -# All rights reserved. -# -# This source code is licensed under the license found in the -# LICENSE file in the root directory of this source tree. - -# References: -# https://github.com/facebookresearch/dino/blob/main/vision_transformer.py -# https://github.com/rwightman/pytorch-image-models/tree/master/timm/models/vision_transformer.py - -from functools import partial -import math -from typing import Sequence, Tuple, Union, Callable - -import torch -import torch.nn as nn -from torch.nn.init import trunc_normal_ - -from .layers import Mlp, PatchEmbed, Block, Attention - - -def named_apply( - fn: Callable, module: nn.Module, name="", depth_first=True, include_root=False -) -> nn.Module: - if not depth_first and include_root: - fn(module=module, name=name) - for child_name, child_module in module.named_children(): - child_name = ".".join((name, child_name)) if name else child_name - named_apply( - fn=fn, - module=child_module, - name=child_name, - depth_first=depth_first, - include_root=True, - ) - if depth_first and include_root: - fn(module=module, name=name) - return module - - -class BlockChunk(nn.ModuleList): - def forward(self, x): - for b in self: - x = b(x) - return x - - -class DinoVisionTransformer(nn.Module): - def __init__( - self, - img_size=224, - patch_size=16, - in_chans=3, - embed_dim=768, - depth=12, - num_heads=12, - mlp_ratio=4.0, - qkv_bias=True, - ffn_bias=True, - proj_bias=True, - drop_path_rate=0.0, - drop_path_uniform=False, - init_values=None, # for layerscale: None or 0 => no layerscale - embed_layer=PatchEmbed, - act_layer=nn.GELU, - block_fn=Block, - ffn_layer="mlp", - block_chunks=1, - ): - """ - Args: - img_size (int, tuple): input image size - patch_size (int, tuple): patch size - in_chans (int): number of input channels - embed_dim (int): embedding dimension - depth (int): depth of transformer - num_heads (int): number of attention heads - mlp_ratio (int): ratio of mlp hidden dim to embedding dim - qkv_bias (bool): enable bias for qkv if True - proj_bias (bool): enable bias for proj in attn if True - ffn_bias (bool): enable bias for ffn if True - drop_path_rate (float): stochastic depth rate - drop_path_uniform (bool): apply uniform drop rate across blocks - weight_init (str): weight init scheme - init_values (float): layer-scale init values - embed_layer (nn.Module): patch embedding layer - act_layer (nn.Module): MLP activation layer - block_fn (nn.Module): transformer block class - ffn_layer (str): "mlp", "swiglu", "swiglufused" or "identity" - block_chunks: (int) split block sequence into block_chunks units for FSDP wrap - """ - super().__init__() - norm_layer = partial(nn.LayerNorm, eps=1e-6) - - self.num_features = self.embed_dim = ( - embed_dim # num_features for consistency with other models - ) - self.num_tokens = 1 - self.n_blocks = depth - self.num_heads = num_heads - self.patch_size = patch_size - - self.patch_embed = embed_layer( - img_size=img_size, - patch_size=patch_size, - in_chans=in_chans, - embed_dim=embed_dim, - ) - num_patches = self.patch_embed.num_patches - - self.cls_token = nn.Parameter(torch.zeros(1, 1, embed_dim)) - self.pos_embed = nn.Parameter( - torch.zeros(1, num_patches + self.num_tokens, embed_dim) - ) - - if drop_path_uniform is True: - dpr = [drop_path_rate] * depth - else: - dpr = [ - x.item() for x in torch.linspace(0, drop_path_rate, depth) - ] # stochastic depth decay rule - - if ffn_layer == "mlp": - ffn_layer = Mlp - elif ffn_layer == "identity": - - def f(*args, **kwargs): - return nn.Identity() - - ffn_layer = f - else: - raise NotImplementedError - - blocks_list = [ - block_fn( - dim=embed_dim, - num_heads=num_heads, - mlp_ratio=mlp_ratio, - qkv_bias=qkv_bias, - proj_bias=proj_bias, - ffn_bias=ffn_bias, - drop_path=dpr[i], - norm_layer=norm_layer, - act_layer=act_layer, - ffn_layer=ffn_layer, - init_values=init_values, - ) - for i in range(depth) - ] - if block_chunks > 0: - self.chunked_blocks = True - chunked_blocks = [] - chunksize = depth // block_chunks - for i in range(0, depth, chunksize): - # this is to keep the block index consistent if we chunk the block list - chunked_blocks.append( - [nn.Identity()] * i + blocks_list[i : i + chunksize] - ) - self.blocks = nn.ModuleList([BlockChunk(p) for p in chunked_blocks]) - else: - self.chunked_blocks = False - self.blocks = nn.ModuleList(blocks_list) - - self.norm = norm_layer(embed_dim) - self.head = nn.Identity() - - self.mask_token = nn.Parameter(torch.zeros(1, embed_dim)) - - self.init_weights() - for param in self.parameters(): - param.requires_grad = False - - @property - def device(self): - return self.cls_token.device - - def init_weights(self): - trunc_normal_(self.pos_embed, std=0.02) - nn.init.normal_(self.cls_token, std=1e-6) - named_apply(init_weights_vit_timm, self) - - def interpolate_pos_encoding(self, x, w, h): - previous_dtype = x.dtype - npatch = x.shape[1] - 1 - N = self.pos_embed.shape[1] - 1 - if npatch == N and w == h: - return self.pos_embed - pos_embed = self.pos_embed.float() - class_pos_embed = pos_embed[:, 0].clone() - patch_pos_embed = pos_embed[:, 1:].clone() - dim = x.shape[-1] - w0 = w // self.patch_size - h0 = h // self.patch_size - # we add a small number to avoid floating point error in the interpolation - # see discussion at https://github.com/facebookresearch/dino/issues/8 - w0, h0 = w0 + 0.1, h0 + 0.1 - - patch_pos_embed = nn.functional.interpolate( - patch_pos_embed.reshape( - 1, int(math.sqrt(N)), int(math.sqrt(N)), dim - ).permute(0, 3, 1, 2), - scale_factor=(w0 / math.sqrt(N), h0 / math.sqrt(N)), - mode="bicubic", - ) - - assert ( - int(w0) == patch_pos_embed.shape[-2] - and int(h0) == patch_pos_embed.shape[-1] - ) - patch_pos_embed = patch_pos_embed.permute(0, 2, 3, 1).view(1, -1, dim) - return torch.cat((class_pos_embed.unsqueeze(0), patch_pos_embed), dim=1).to( - previous_dtype - ) - - def prepare_tokens_with_masks(self, x, masks=None): - B, nc, w, h = x.shape - x = self.patch_embed(x) - if masks is not None: - x = torch.where( - masks.unsqueeze(-1), self.mask_token.to(x.dtype).unsqueeze(0), x - ) - - x = torch.cat((self.cls_token.repeat(x.shape[0], 1, 1), x), dim=1) - x = x + self.interpolate_pos_encoding(x, w, h) - - return x.clone() - - def forward_features_list(self, x_list, masks_list): - x = [ - self.prepare_tokens_with_masks(x, masks) - for x, masks in zip(x_list, masks_list) - ] - for blk in self.blocks: - x = blk(x) - - all_x = x - output = [] - for x, masks in zip(all_x, masks_list): - x_norm = self.norm(x) - output.append( - { - "x_norm_clstoken": x_norm[:, 0], - "x_norm_patchtokens": x_norm[:, 1:], - "x_prenorm": x, - "masks": masks, - } - ) - return output - - def forward_features(self, x, masks=None): - if isinstance(x, list): - return self.forward_features_list(x, masks) - - x = self.prepare_tokens_with_masks(x, masks) - - for blk in self.blocks: - x = blk(x) - - x_norm = self.norm(x) - return { - "x_norm_clstoken": x_norm[:, 0], - "x_norm_patchtokens": x_norm[:, 1:], - "x_prenorm": x, - "masks": masks, - } - - def _get_intermediate_layers_not_chunked(self, x, n=1): - x = self.prepare_tokens_with_masks(x) - # If n is an int, take the n last blocks. If it's a list, take them - output, total_block_len = [], len(self.blocks) - blocks_to_take = ( - range(total_block_len - n, total_block_len) if isinstance(n, int) else n - ) - for i, blk in enumerate(self.blocks): - x = blk(x) - if i in blocks_to_take: - output.append(x) - assert len(output) == len(blocks_to_take), ( - f"only {len(output)} / {len(blocks_to_take)} blocks found" - ) - return output - - def _get_intermediate_layers_chunked(self, x, n=1): - x = self.prepare_tokens_with_masks(x) - output, i, total_block_len = [], 0, len(self.blocks[-1]) - # If n is an int, take the n last blocks. If it's a list, take them - blocks_to_take = ( - range(total_block_len - n, total_block_len) if isinstance(n, int) else n - ) - for block_chunk in self.blocks: - for blk in block_chunk[i:]: # Passing the nn.Identity() - x = blk(x) - if i in blocks_to_take: - output.append(x) - i += 1 - assert len(output) == len(blocks_to_take), ( - f"only {len(output)} / {len(blocks_to_take)} blocks found" - ) - return output - - def get_intermediate_layers( - self, - x: torch.Tensor, - n: Union[int, Sequence] = 1, # Layers or n last layers to take - reshape: bool = False, - return_class_token: bool = False, - norm=True, - ) -> Tuple[Union[torch.Tensor, Tuple[torch.Tensor]]]: - if self.chunked_blocks: - outputs = self._get_intermediate_layers_chunked(x, n) - else: - outputs = self._get_intermediate_layers_not_chunked(x, n) - if norm: - outputs = [self.norm(out) for out in outputs] - class_tokens = [out[:, 0] for out in outputs] - outputs = [out[:, 1:] for out in outputs] - if reshape: - B, _, w, h = x.shape - outputs = [ - out.reshape(B, w // self.patch_size, h // self.patch_size, -1) - .permute(0, 3, 1, 2) - .contiguous() - for out in outputs - ] - if return_class_token: - return tuple(zip(outputs, class_tokens)) - return tuple(outputs) - - def forward(self, *args, is_training=False, **kwargs): - ret = self.forward_features(*args, **kwargs) - if is_training: - return ret - else: - return self.head(ret["x_norm_clstoken"]) - - -def init_weights_vit_timm(module: nn.Module, name: str = ""): - """ViT weight initialization, original timm impl (for reproducibility)""" - if isinstance(module, nn.Linear): - trunc_normal_(module.weight, std=0.02) - if module.bias is not None: - nn.init.zeros_(module.bias) - - -def vit_small(patch_size=16, **kwargs): - model = DinoVisionTransformer( - patch_size=patch_size, - embed_dim=384, - depth=12, - num_heads=6, - mlp_ratio=4, - block_fn=partial(Block, attn_class=Attention), - **kwargs, - ) - return model - - -def vit_base(patch_size=16, **kwargs): - model = DinoVisionTransformer( - patch_size=patch_size, - embed_dim=768, - depth=12, - num_heads=12, - mlp_ratio=4, - block_fn=partial(Block, attn_class=Attention), - **kwargs, - ) - return model - - -def vit_large(patch_size=16, **kwargs): - model = DinoVisionTransformer( - patch_size=patch_size, - embed_dim=1024, - depth=24, - num_heads=16, - mlp_ratio=4, - block_fn=partial(Block, attn_class=Attention), - **kwargs, - ) - return model - - -def vit_giant2(patch_size=16, **kwargs): - """ - Close to ViT-giant, with embed-dim 1536 and 24 heads => embed-dim per head 64 - """ - model = DinoVisionTransformer( - patch_size=patch_size, - embed_dim=1536, - depth=40, - num_heads=24, - mlp_ratio=4, - block_fn=partial(Block, attn_class=Attention), - **kwargs, - ) - return model diff --git a/src/baselines/loma/descriptor/transformer/layers/__init__.py b/src/baselines/loma/descriptor/transformer/layers/__init__.py deleted file mode 100644 index 61bbcec..0000000 --- a/src/baselines/loma/descriptor/transformer/layers/__init__.py +++ /dev/null @@ -1,10 +0,0 @@ -# Copyright (c) Meta Platforms, Inc. and affiliates. -# All rights reserved. -# -# This source code is licensed under the license found in the -# LICENSE file in the root directory of this source tree. - -from .mlp import Mlp as Mlp -from .patch_embed import PatchEmbed as PatchEmbed -from .block import Block as Block -from .attention import Attention as Attention diff --git a/src/baselines/loma/descriptor/transformer/layers/attention.py b/src/baselines/loma/descriptor/transformer/layers/attention.py deleted file mode 100644 index 40a788d..0000000 --- a/src/baselines/loma/descriptor/transformer/layers/attention.py +++ /dev/null @@ -1,59 +0,0 @@ -# Copyright (c) Meta Platforms, Inc. and affiliates. -# All rights reserved. -# -# This source code is licensed under the license found in the -# LICENSE file in the root directory of this source tree. - -# References: -# https://github.com/facebookresearch/dino/blob/master/vision_transformer.py -# https://github.com/rwightman/pytorch-image-models/tree/master/timm/models/vision_transformer.py - -import logging - -from torch import Tensor -from torch import nn -import torch.nn.functional as F - - -logger = logging.getLogger("dinov2") - - -class Attention(nn.Module): - def __init__( - self, - dim: int, - num_heads: int = 8, - qkv_bias: bool = False, - proj_bias: bool = True, - attn_drop: float = 0.0, - proj_drop: float = 0.0, - ) -> None: - super().__init__() - self.num_heads = num_heads - head_dim = dim // num_heads - self.scale = head_dim**-0.5 - - self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias) - self.attn_drop = nn.Dropout(attn_drop) - self.proj = nn.Linear(dim, dim, bias=proj_bias) - self.proj_drop = nn.Dropout(proj_drop) - - def forward(self, x: Tensor) -> Tensor: - B, N, C = x.shape - qkv = ( - self.qkv(x) - .reshape(B, N, 3, self.num_heads, C // self.num_heads) - .permute(2, 0, 3, 1, 4) - ) - # q, k, v = qkv[0] * self.scale, qkv[1], qkv[2] - # attn = q @ k.transpose(-2, -1) - # attn = attn.softmax(dim=-1) - # attn = self.attn_drop(attn) - # x = attn @ v - # x = x.transpose(1, 2).reshape(B, N, C) - # almost same //je - x = F.scaled_dot_product_attention(qkv[0], qkv[1], qkv[2]) - x = x.permute(0, 2, 1, 3).reshape(B, N, C) - x = self.proj(x) - x = self.proj_drop(x) - return x diff --git a/src/baselines/loma/descriptor/transformer/layers/block.py b/src/baselines/loma/descriptor/transformer/layers/block.py deleted file mode 100644 index 5bea668..0000000 --- a/src/baselines/loma/descriptor/transformer/layers/block.py +++ /dev/null @@ -1,144 +0,0 @@ -# Copyright (c) Meta Platforms, Inc. and affiliates. -# All rights reserved. -# -# This source code is licensed under the license found in the -# LICENSE file in the root directory of this source tree. - -# References: -# https://github.com/facebookresearch/dino/blob/master/vision_transformer.py -# https://github.com/rwightman/pytorch-image-models/tree/master/timm/layers/patch_embed.py - -import logging -from typing import Callable - -import torch -from torch import nn, Tensor - -from .attention import Attention -from .drop_path import DropPath -from .layer_scale import LayerScale -from .mlp import Mlp - - -logger = logging.getLogger("dinov2") - - -class Block(nn.Module): - def __init__( - self, - dim: int, - num_heads: int, - mlp_ratio: float = 4.0, - qkv_bias: bool = False, - proj_bias: bool = True, - ffn_bias: bool = True, - drop: float = 0.0, - attn_drop: float = 0.0, - init_values=None, - drop_path: float = 0.0, - act_layer: Callable[..., nn.Module] = nn.GELU, - norm_layer: Callable[..., nn.Module] = nn.LayerNorm, - attn_class: Callable[..., nn.Module] = Attention, - ffn_layer: Callable[..., nn.Module] = Mlp, - ) -> None: - super().__init__() - self.norm1 = norm_layer(dim) - self.attn = attn_class( - dim, - num_heads=num_heads, - qkv_bias=qkv_bias, - proj_bias=proj_bias, - attn_drop=attn_drop, - proj_drop=drop, - ) - self.ls1 = ( - LayerScale(dim, init_values=init_values) if init_values else nn.Identity() - ) - self.drop_path1 = DropPath(drop_path) if drop_path > 0.0 else nn.Identity() - - self.norm2 = norm_layer(dim) - mlp_hidden_dim = int(dim * mlp_ratio) - self.mlp = ffn_layer( - in_features=dim, - hidden_features=mlp_hidden_dim, - act_layer=act_layer, - drop=drop, - bias=ffn_bias, - ) - self.ls2 = ( - LayerScale(dim, init_values=init_values) if init_values else nn.Identity() - ) - self.drop_path2 = DropPath(drop_path) if drop_path > 0.0 else nn.Identity() - - self.sample_drop_ratio = drop_path - - def forward(self, x: Tensor) -> Tensor: - def attn_residual_func(x: Tensor) -> Tensor: - return self.ls1(self.attn(self.norm1(x))) - - def ffn_residual_func(x: Tensor) -> Tensor: - return self.ls2(self.mlp(self.norm2(x))) - - if self.training and self.sample_drop_ratio > 0.1: - # the overhead is compensated only for a drop path rate larger than 0.1 - x = drop_add_residual_stochastic_depth( - x, - residual_func=attn_residual_func, - sample_drop_ratio=self.sample_drop_ratio, - ) - x = drop_add_residual_stochastic_depth( - x, - residual_func=ffn_residual_func, - sample_drop_ratio=self.sample_drop_ratio, - ) - elif self.training and self.sample_drop_ratio > 0.0: - x = x + self.drop_path1(attn_residual_func(x)) - x = x + self.drop_path2(ffn_residual_func(x)) - else: - x = x + attn_residual_func(x) - x = x + ffn_residual_func(x) - return x - - -def drop_add_residual_stochastic_depth( - x: Tensor, - residual_func: Callable[[Tensor], Tensor], - sample_drop_ratio: float = 0.0, -) -> Tensor: - # 1) extract subset using permutation - b, n, d = x.shape - sample_subset_size = max(int(b * (1 - sample_drop_ratio)), 1) - brange = (torch.randperm(b, device=x.device))[:sample_subset_size] - x_subset = x[brange] - - # 2) apply residual_func to get residual - residual = residual_func(x_subset) - - x_flat = x.flatten(1) - residual = residual.flatten(1) - - residual_scale_factor = b / sample_subset_size - - # 3) add the residual - x_plus_residual = torch.index_add( - x_flat, 0, brange, residual.to(dtype=x.dtype), alpha=residual_scale_factor - ) - return x_plus_residual.view_as(x) - - -def get_branges_scales(x, sample_drop_ratio=0.0): - b, n, d = x.shape - sample_subset_size = max(int(b * (1 - sample_drop_ratio)), 1) - brange = (torch.randperm(b, device=x.device))[:sample_subset_size] - residual_scale_factor = b / sample_subset_size - return brange, residual_scale_factor - - -def add_residual(x, brange, residual, residual_scale_factor, scaling_vector=None): - assert scaling_vector is None, "Scaling vector is not supported" - x_flat = x.flatten(1) - residual = residual.flatten(1) - x_plus_residual = torch.index_add( - x_flat, 0, brange, residual.to(dtype=x.dtype), alpha=residual_scale_factor - ) - return x_plus_residual diff --git a/src/baselines/loma/descriptor/transformer/layers/drop_path.py b/src/baselines/loma/descriptor/transformer/layers/drop_path.py deleted file mode 100644 index a23ba73..0000000 --- a/src/baselines/loma/descriptor/transformer/layers/drop_path.py +++ /dev/null @@ -1,37 +0,0 @@ -# Copyright (c) Meta Platforms, Inc. and affiliates. -# All rights reserved. -# -# This source code is licensed under the license found in the -# LICENSE file in the root directory of this source tree. - -# References: -# https://github.com/facebookresearch/dino/blob/master/vision_transformer.py -# https://github.com/rwightman/pytorch-image-models/tree/master/timm/layers/drop.py - - -from torch import nn - - -def drop_path(x, drop_prob: float = 0.0, training: bool = False): - if drop_prob == 0.0 or not training: - return x - keep_prob = 1 - drop_prob - shape = (x.shape[0],) + (1,) * ( - x.ndim - 1 - ) # work with diff dim tensors, not just 2D ConvNets - random_tensor = x.new_empty(shape).bernoulli_(keep_prob) - if keep_prob > 0.0: - random_tensor.div_(keep_prob) - output = x * random_tensor - return output - - -class DropPath(nn.Module): - """Drop paths (Stochastic Depth) per sample (when applied in main path of residual blocks).""" - - def __init__(self, drop_prob=None): - super(DropPath, self).__init__() - self.drop_prob = drop_prob - - def forward(self, x): - return drop_path(x, self.drop_prob, self.training) diff --git a/src/baselines/loma/descriptor/transformer/layers/layer_scale.py b/src/baselines/loma/descriptor/transformer/layers/layer_scale.py deleted file mode 100644 index ca5daa5..0000000 --- a/src/baselines/loma/descriptor/transformer/layers/layer_scale.py +++ /dev/null @@ -1,28 +0,0 @@ -# Copyright (c) Meta Platforms, Inc. and affiliates. -# All rights reserved. -# -# This source code is licensed under the license found in the -# LICENSE file in the root directory of this source tree. - -# Modified from: https://github.com/huggingface/pytorch-image-models/blob/main/timm/models/vision_transformer.py#L103-L110 - -from typing import Union - -import torch -from torch import Tensor -from torch import nn - - -class LayerScale(nn.Module): - def __init__( - self, - dim: int, - init_values: Union[float, Tensor] = 1e-5, - inplace: bool = False, - ) -> None: - super().__init__() - self.inplace = inplace - self.gamma = nn.Parameter(init_values * torch.ones(dim)) - - def forward(self, x: Tensor) -> Tensor: - return x.mul_(self.gamma) if self.inplace else x * self.gamma diff --git a/src/baselines/loma/descriptor/transformer/layers/mlp.py b/src/baselines/loma/descriptor/transformer/layers/mlp.py deleted file mode 100644 index 5e4b315..0000000 --- a/src/baselines/loma/descriptor/transformer/layers/mlp.py +++ /dev/null @@ -1,41 +0,0 @@ -# Copyright (c) Meta Platforms, Inc. and affiliates. -# All rights reserved. -# -# This source code is licensed under the license found in the -# LICENSE file in the root directory of this source tree. - -# References: -# https://github.com/facebookresearch/dino/blob/master/vision_transformer.py -# https://github.com/rwightman/pytorch-image-models/tree/master/timm/layers/mlp.py - - -from typing import Callable, Optional - -from torch import Tensor, nn - - -class Mlp(nn.Module): - def __init__( - self, - in_features: int, - hidden_features: Optional[int] = None, - out_features: Optional[int] = None, - act_layer: Callable[..., nn.Module] = nn.GELU, - drop: float = 0.0, - bias: bool = True, - ) -> None: - super().__init__() - out_features = out_features or in_features - hidden_features = hidden_features or in_features - self.fc1 = nn.Linear(in_features, hidden_features, bias=bias) - self.act = act_layer() - self.fc2 = nn.Linear(hidden_features, out_features, bias=bias) - self.drop = nn.Dropout(drop) - - def forward(self, x: Tensor) -> Tensor: - x = self.fc1(x) - x = self.act(x) - x = self.drop(x) - x = self.fc2(x) - x = self.drop(x) - return x diff --git a/src/baselines/loma/descriptor/transformer/layers/patch_embed.py b/src/baselines/loma/descriptor/transformer/layers/patch_embed.py deleted file mode 100644 index e52fae8..0000000 --- a/src/baselines/loma/descriptor/transformer/layers/patch_embed.py +++ /dev/null @@ -1,101 +0,0 @@ -# Copyright (c) Meta Platforms, Inc. and affiliates. -# All rights reserved. -# -# This source code is licensed under the license found in the -# LICENSE file in the root directory of this source tree. - -# References: -# https://github.com/facebookresearch/dino/blob/master/vision_transformer.py -# https://github.com/rwightman/pytorch-image-models/tree/master/timm/layers/patch_embed.py - -from typing import Callable, Optional, Tuple, Union - -from torch import Tensor -import torch.nn as nn - - -def make_2tuple(x): - if isinstance(x, tuple): - assert len(x) == 2 - return x - - assert isinstance(x, int) - return (x, x) - - -class PatchEmbed(nn.Module): - """ - 2D image to patch embedding: (B,C,H,W) -> (B,N,D) - - Args: - img_size: Image size. - patch_size: Patch token size. - in_chans: Number of input image channels. - embed_dim: Number of linear projection output channels. - norm_layer: Normalization layer. - """ - - def __init__( - self, - img_size: Union[int, Tuple[int, int]] = 224, - patch_size: Union[int, Tuple[int, int]] = 16, - in_chans: int = 3, - embed_dim: int = 768, - norm_layer: Optional[Callable] = None, - flatten_embedding: bool = True, - ) -> None: - super().__init__() - - image_HW = make_2tuple(img_size) - patch_HW = make_2tuple(patch_size) - patch_grid_size = ( - image_HW[0] // patch_HW[0], - image_HW[1] // patch_HW[1], - ) - - self.img_size = image_HW - self.patch_size = patch_HW - self.patches_resolution = patch_grid_size - self.num_patches = patch_grid_size[0] * patch_grid_size[1] - - self.in_chans = in_chans - self.embed_dim = embed_dim - - self.flatten_embedding = flatten_embedding - - self.proj = nn.Conv2d( - in_chans, embed_dim, kernel_size=patch_HW, stride=patch_HW - ) - self.norm = norm_layer(embed_dim) if norm_layer else nn.Identity() - - def forward(self, x: Tensor) -> Tensor: - _, _, H, W = x.shape - patch_H, patch_W = self.patch_size - - assert H % patch_H == 0, ( - f"Input image height {H} is not a multiple of patch height {patch_H}" - ) - assert W % patch_W == 0, ( - f"Input image width {W} is not a multiple of patch width: {patch_W}" - ) - - x = self.proj(x) # B C H W - H, W = x.size(2), x.size(3) - x = x.flatten(2).transpose(1, 2) # B HW C - x = self.norm(x) - if not self.flatten_embedding: - x = x.reshape(-1, H, W, self.embed_dim) # B H W C - return x - - def flops(self) -> float: - Ho, Wo = self.patches_resolution - flops = ( - Ho - * Wo - * self.embed_dim - * self.in_chans - * (self.patch_size[0] * self.patch_size[1]) - ) - if self.norm is not None: - flops += Ho * Wo * self.embed_dim - return flops diff --git a/src/baselines/loma/detector/__init__.py b/src/baselines/loma/detector/__init__.py deleted file mode 100644 index 900701b..0000000 --- a/src/baselines/loma/detector/__init__.py +++ /dev/null @@ -1 +0,0 @@ -from .dad import DaD as DaD diff --git a/src/baselines/loma/detector/dad.py b/src/baselines/loma/detector/dad.py deleted file mode 100644 index 952af4d..0000000 --- a/src/baselines/loma/detector/dad.py +++ /dev/null @@ -1,426 +0,0 @@ -from dataclasses import dataclass -from pathlib import Path -from typing import Literal -import numpy as np - -import torch -import torch.nn as nn -import torch.nn.functional as F -import torchvision.models as tvm -import torchvision.transforms as transforms -from PIL import Image -import logging - -from ..io import check_not_i16 -from ..types import Batch, Model -from ..device import device, amp_dtype -from .utils import sample_keypoints - -logger = logging.getLogger(__name__) - - -def _images_from_detector_input( - batch: torch.Tensor | Batch | dict[str, torch.Tensor], -) -> torch.Tensor: - if isinstance(batch, torch.Tensor): - return batch - if isinstance(batch, Batch): - return torch.cat([batch.img_A, batch.img_B], dim=0) - return batch["image"] - - -class DaD(Model): - @dataclass(frozen=True) - class Cfg: - compile: bool = True - remove_borders: bool = False - resize: int = 1024 - nms_size: int = 3 - increase_coverage: bool = False - coverage_pow: int | None = None - coverage_size: int | None = None - subpixel: bool = True - subpixel_temp: float = 0.5 - keep_aspect_ratio: bool = True - coverage_from_sparse: bool = False - arch: Literal["dedode_s"] = "dedode_s" - is_sigmoid: bool = False - - def __init__( - self, - cfg: Cfg | None = None, - ) -> None: - super().__init__() - if cfg is None: - cfg = DaD.Cfg() - weights = torch.hub.load_state_dict_from_url( - "https://github.com/Parskatt/dad/releases/download/v0.1.0/dad.pth", - map_location="cpu", - ) - if cfg.arch == "dedode_s": - encoder, decoder = dedode_detector_S() - else: - raise ValueError(f"Architecture {cfg.arch} not supported") - self.cfg = cfg - self.normalizer = transforms.Normalize( - mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225] - ) - self.encoder = encoder - self.decoder = decoder - self.remove_borders = cfg.remove_borders - self.resize = cfg.resize - self.increase_coverage = cfg.increase_coverage - self.coverage_pow = cfg.coverage_pow - self.coverage_size = cfg.coverage_size - self.nms_size = cfg.nms_size - self.keep_aspect_ratio = cfg.keep_aspect_ratio - self.subpixel = cfg.subpixel - self.subpixel_temp = cfg.subpixel_temp - self.coverage_from_sparse = cfg.coverage_from_sparse - if weights is not None: - self.load_state_dict(weights) - if cfg.compile: - logger.info("Compiling DaD...") - self.compile() - self.to(device) - - @property - def topleft(self): - return 0.5 - - def forward_impl( - self, - images, - ): - images = self.normalizer(images) - features, sizes = self.encoder(images) - logits: torch.Tensor = images[:, :1].new_zeros(()) - context = None - scales = ["8", "4", "2", "1"] - for idx, (feature_map, scale) in enumerate(zip(reversed(features), scales)): - delta_logits, context = self.decoder( - feature_map, context=context, scale=scale - ) - logits = ( - logits + delta_logits.float() - ) # ensure float (need bf16 doesnt have f.interpolate) - if idx < len(scales) - 1: - size = sizes[-(idx + 2)] - logits = F.interpolate( - logits, size=size, mode="bicubic", align_corners=False - ) - context = F.interpolate( - context.float(), size=size, mode="bilinear", align_corners=False - ) - return logits.float() - - def forward( - self, - images: torch.Tensor, - num_keypoints: int, - *, - return_dense_probs: bool = False, - ) -> dict[str, torch.Tensor]: - scoremap = self.forward_impl(images) - B, K, H, W = scoremap.shape - dense_probs = ( - scoremap.reshape(B, K * H * W) - .softmax(dim=-1) - .reshape(B, K, H * W) - .sum(dim=1) - ) - dense_probs = dense_probs.reshape(B, H, W) - keypoints, confidence = sample_keypoints( - dense_probs, - use_nms=True, - nms_size=self.nms_size, - sample_topk=True, - num_samples=num_keypoints, - return_probs=True, - increase_coverage=self.increase_coverage, - coverage_from_sparse=self.coverage_from_sparse, - remove_borders=self.remove_borders, - coverage_pow=float(self.coverage_pow) - if self.coverage_pow is not None - else 0.5, - coverage_size=self.coverage_size if self.coverage_size is not None else 51, - subpixel=self.subpixel, - subpixel_temp=self.subpixel_temp, - scoremap=scoremap.reshape(B, H, W), - ) - result = {"keypoints": keypoints, "keypoint_probs": confidence} - if return_dense_probs: - result["dense_probs"] = dense_probs - return result - - @torch.inference_mode() - def detect( - self, - batch: torch.Tensor | Batch | dict[str, torch.Tensor], - *, - num_keypoints: int, - return_dense_probs: bool = False, - ) -> dict[str, torch.Tensor]: - self.train(False) - images = _images_from_detector_input(batch).to(device) - return self(images, num_keypoints, return_dense_probs=return_dense_probs) - - def load_image(self, im_path, device=device) -> torch.Tensor: - pil_im = Image.open(im_path) - check_not_i16(pil_im) - pil_im = pil_im.convert("RGB") - if self.keep_aspect_ratio: - W, H = pil_im.size - scale = self.resize / max(W, H) - W = int((scale * W) // 8 * 8) - H = int((scale * H) // 8 * 8) - else: - H, W = self.resize, self.resize - pil_im = pil_im.resize((W, H)) - standard_im = np.array(pil_im) / 255.0 - return torch.from_numpy(standard_im).permute(2, 0, 1).float().to(device)[None] - - @torch.inference_mode() - def detect_from_path( - self, - im_path: str | Path, - *, - num_keypoints: int, - return_dense_probs: bool = False, - ) -> dict[str, torch.Tensor]: - return self.detect( - self.load_image(im_path), - num_keypoints=num_keypoints, - return_dense_probs=return_dense_probs, - ) - - def to_pixel_coords( - self, normalized_coords: torch.Tensor, h: int, w: int - ) -> torch.Tensor: - if normalized_coords.shape[-1] != 2: - raise ValueError( - f"Expected shape (..., 2), but got {normalized_coords.shape}" - ) - pixel_coords = torch.stack( - ( - w * (normalized_coords[..., 0] + 1) / 2, - h * (normalized_coords[..., 1] + 1) / 2, - ), - dim=-1, - ) - return pixel_coords - - def to_normalized_coords( - self, pixel_coords: torch.Tensor, h: int, w: int - ) -> torch.Tensor: - if pixel_coords.shape[-1] != 2: - raise ValueError(f"Expected shape (..., 2), but got {pixel_coords.shape}") - normalized_coords = torch.stack( - ( - 2 * (pixel_coords[..., 0]) / w - 1, - 2 * (pixel_coords[..., 1]) / h - 1, - ), - dim=-1, - ) - return normalized_coords - - -class Decoder(nn.Module): - def __init__( - self, layers, *args, super_resolution=False, num_prototypes=1, **kwargs - ) -> None: - super().__init__(*args, **kwargs) - self.layers = layers - self.super_resolution = super_resolution - self.num_prototypes = num_prototypes - - def forward(self, features, context=None, scale=None): - if context is not None: - features = torch.cat((features, context), dim=1) - stuff = self.layers[scale](features) - logits, context = ( - stuff[:, : self.num_prototypes], - stuff[:, self.num_prototypes :], - ) - return logits, context - - -class ConvRefiner(nn.Module): - def __init__( - self, - in_dim=6, - hidden_dim=16, - out_dim=2, - dw=True, - kernel_size=5, - hidden_blocks=5, - amp=True, - residual=False, - amp_dtype=amp_dtype, - ): - super().__init__() - self.block1 = self.create_block( - in_dim, - hidden_dim, - dw=False, - kernel_size=1, - ) - self.hidden_blocks = nn.Sequential( - *[ - self.create_block( - hidden_dim, - hidden_dim, - dw=dw, - kernel_size=kernel_size, - ) - for hb in range(hidden_blocks) - ] - ) - self.out_conv = nn.Conv2d(hidden_dim, out_dim, 1, 1, 0) - self.amp = amp - self.amp_dtype = amp_dtype - self.residual = residual - - def create_block( - self, - in_dim, - out_dim, - dw=True, - kernel_size=5, - bias=True, - norm_type=nn.BatchNorm2d, - ): - num_groups = 1 if not dw else in_dim - if dw: - assert out_dim % in_dim == 0, ( - "outdim must be divisible by indim for depthwise" - ) - conv1 = nn.Conv2d( - in_dim, - out_dim, - kernel_size=kernel_size, - stride=1, - padding=kernel_size // 2, - groups=num_groups, - bias=bias, - ) - norm = ( - norm_type(out_dim) - if norm_type is nn.BatchNorm2d - else norm_type(num_channels=out_dim) - ) - relu = nn.ReLU(inplace=True) - conv2 = nn.Conv2d(out_dim, out_dim, 1, 1, 0) - return nn.Sequential(conv1, norm, relu, conv2) - - def forward(self, feats): - b, c, hs, ws = feats.shape - with torch.autocast( - device_type=feats.device.type, enabled=self.amp, dtype=self.amp_dtype - ): - x0 = self.block1(feats) - x = self.hidden_blocks(x0) - if self.residual: - x = (x + x0) / 1.4 - x = self.out_conv(x) - return x - - -class VGG19(nn.Module): - def __init__(self, amp=False) -> None: - super().__init__() - self.layers = nn.ModuleList(tvm.vgg19_bn().features[:40]) - # Maxpool layers: 6, 13, 26, 39 - self.amp = amp - self.amp_dtype = amp_dtype - - def forward(self, x, **kwargs): - with torch.autocast( - device_type=x.device.type, enabled=self.amp, dtype=self.amp_dtype - ): - feats = [] - sizes = [] - for layer in self.layers: - if isinstance(layer, nn.MaxPool2d): - feats.append(x) - sizes.append(x.shape[-2:]) - x = layer(x) - return feats, sizes - - -class VGG(nn.Module): - def __init__(self, size="19", amp=False, amp_dtype=amp_dtype) -> None: - super().__init__() - if size == "11": - self.layers = nn.ModuleList(tvm.vgg11_bn().features[:22]) - elif size == "13": - self.layers = nn.ModuleList(tvm.vgg13_bn().features[:28]) - elif size == "19": - self.layers = nn.ModuleList(tvm.vgg19_bn().features[:40]) - # Maxpool layers: 6, 13, 26, 39 - self.amp = amp - self.amp_dtype = amp_dtype - - def forward(self, x, **kwargs): - # print(f"device.type: {self.device}") - with torch.autocast( - device_type=x.device.type, enabled=self.amp, dtype=self.amp_dtype - ): - feats = [] - sizes = [] - for layer in self.layers: - if isinstance(layer, nn.MaxPool2d): - feats.append(x) - sizes.append(x.shape[-2:]) - x = layer(x) - return feats, sizes - - -def dedode_detector_S(): - residual = True - hidden_blocks = 3 - amp = True - NUM_PROTOTYPES = 1 - conv_refiner = nn.ModuleDict( - { - "8": ConvRefiner( - 512, - 512, - 256 + NUM_PROTOTYPES, - hidden_blocks=hidden_blocks, - residual=residual, - amp=amp, - amp_dtype=amp_dtype, - ), - "4": ConvRefiner( - 256 + 256, - 256, - 128 + NUM_PROTOTYPES, - hidden_blocks=hidden_blocks, - residual=residual, - amp=amp, - amp_dtype=amp_dtype, - ), - "2": ConvRefiner( - 128 + 128, - 64, - 32 + NUM_PROTOTYPES, - hidden_blocks=hidden_blocks, - residual=residual, - amp=amp, - amp_dtype=amp_dtype, - ), - "1": ConvRefiner( - 64 + 32, - 32, - 1 + NUM_PROTOTYPES, - hidden_blocks=hidden_blocks, - residual=residual, - amp=amp, - amp_dtype=amp_dtype, - ), - } - ) - encoder = VGG(size="11", amp=amp, amp_dtype=amp_dtype) - decoder = Decoder(conv_refiner) - return encoder, decoder diff --git a/src/baselines/loma/detector/utils.py b/src/baselines/loma/detector/utils.py deleted file mode 100644 index 7d648bf..0000000 --- a/src/baselines/loma/detector/utils.py +++ /dev/null @@ -1,213 +0,0 @@ -from __future__ import annotations - -import torch -import torch.nn as nn -import torch.nn.functional as F -from pathlib import Path -from PIL import Image -import numpy as np -import cv2 - -from ..io import tensor_to_pil -from loma.geometry import kde, get_normalized_grid -from loma.types import Detector - - -def extract_patches_from_inds(x: torch.Tensor, inds: torch.Tensor, patch_size: int): - B, H, W = x.shape - B, N = inds.shape - unfolder = nn.Unfold(kernel_size=patch_size, padding=patch_size // 2, stride=1) - unfolded_x: torch.Tensor = unfolder(x[:, None]) # B x K_H * K_W x H * W - patches = torch.gather( - unfolded_x, - dim=2, - index=inds[:, None, :].expand(B, patch_size**2, N), - ) # B x K_H * K_W x N - return patches - - -# @torch.compile() -def sample_keypoints( - keypoint_probs: torch.Tensor, - num_samples=8192, - use_nms=True, - nms_size=1, - sample_topk=True, - increase_coverage=True, - remove_borders=False, - return_probs=False, - coverage_pow=1 / 2, - coverage_size=51, - subpixel=False, - scoremap=None, # required for subpixel - subpixel_temp=0.5, - coverage_from_sparse=False, -): - device = keypoint_probs.device - B, H, W = keypoint_probs.shape - if coverage_from_sparse: - non_coverage_points, non_coverage_probs = sample_keypoints( - keypoint_probs, - num_samples=4 * num_samples, - use_nms=use_nms, - nms_size=nms_size, - sample_topk=sample_topk, - remove_borders=remove_borders, - coverage_pow=coverage_pow, - coverage_size=coverage_size, - subpixel=subpixel, - scoremap=scoremap, - subpixel_temp=subpixel_temp, - increase_coverage=False, - coverage_from_sparse=False, - return_probs=True, - ) - density = kde(non_coverage_points) - balance = 1 / (density + 1) - balance[density < 10] = 1e-7 - p: torch.Tensor = non_coverage_probs * balance - inds = p.argsort(dim=-1, descending=True)[:, :num_samples] - # print(non_coverage_points.shape, inds.shape) - kps = torch.gather( - non_coverage_points, - dim=1, - index=inds[..., None].expand(B, num_samples, 2), - ) - - if return_probs: - return kps, torch.gather(non_coverage_probs, dim=1, index=inds) - return kps - if increase_coverage and not coverage_from_sparse: - weights = ( - -(torch.linspace(-2, 2, steps=coverage_size, device=device) ** 2) - ).exp()[None, None] - # 10000 is just some number for maybe numerical stability, who knows. :), result is invariant anyway - local_density_x = F.conv2d( - (keypoint_probs[:, None] + 1e-6) * 10000, - weights[..., None, :], - padding=(0, coverage_size // 2), - ) - local_density = F.conv2d( - local_density_x, weights[..., None], padding=(coverage_size // 2, 0) - )[:, 0] - keypoint_probs = keypoint_probs * (local_density + 1e-8) ** (-coverage_pow) - grid = get_normalized_grid(B, H, W, overload_device=device).reshape(B, H * W, 2) - if use_nms: - keypoint_probs = keypoint_probs * ( - keypoint_probs - == F.max_pool2d(keypoint_probs, nms_size, stride=1, padding=nms_size // 2) - ) - if remove_borders: - frame = torch.zeros_like(keypoint_probs) - # we hardcode 4px, could do it nicer, but whatever - frame[..., 4:-4, 4:-4] = 1 - keypoint_probs = keypoint_probs * frame - if sample_topk: - inds = torch.topk(keypoint_probs.reshape(B, H * W), k=num_samples).indices - else: - inds = torch.multinomial( - keypoint_probs.reshape(B, H * W), num_samples=num_samples, replacement=False - ) - kps = torch.gather(grid, dim=1, index=inds[..., None].expand(B, num_samples, 2)) - if subpixel: - if scoremap is None: - raise ValueError("scoremap is required when subpixel=True") - offsets = get_normalized_grid( - B, nms_size, nms_size, overload_device=device - ).reshape(B, nms_size**2, 2) # B x K_H x K_W x 2 - offsets[..., 0] = offsets[..., 0] * nms_size / W - offsets[..., 1] = offsets[..., 1] * nms_size / H - keypoint_patch_scores = extract_patches_from_inds(scoremap, inds, nms_size) - keypoint_patch_probs = (keypoint_patch_scores / subpixel_temp).softmax( - dim=1 - ) # B x K_H * K_W x N - keypoint_offsets = torch.einsum("bkn, bkd ->bnd", keypoint_patch_probs, offsets) - kps = kps + keypoint_offsets - if return_probs: - return kps, torch.gather(keypoint_probs.reshape(B, H * W), dim=1, index=inds) - return kps - - -def masked_log_softmax(logits, mask): - masked_logits = torch.full_like(logits, -torch.inf) - masked_logits[mask] = logits[mask] - log_p = masked_logits.log_softmax(dim=-1) - return log_p - - -def masked_softmax(logits, mask): - masked_logits = torch.full_like(logits, -torch.inf) - masked_logits[mask] = logits[mask] - log_p = masked_logits.softmax(dim=-1) - return log_p - - -def cross_entropy(log_p_hat: torch.Tensor, p: torch.Tensor): - return -(log_p_hat * p).sum(dim=-1) - - -def kl_div(p: torch.Tensor, log_p_hat: torch.Tensor): - return cross_entropy(log_p_hat, p) - cross_entropy((p + 1e-12).log(), p) - - -def draw_kpts(im, kpts, radius=2, width=1): - im = np.array(im) - # Convert keypoints to numpy array - kpts_np = kpts.cpu().numpy() - - # Create a copy of the image to draw on - ret = im.copy() - - # Define green color (BGR format in OpenCV) - green_color = (0, 255, 0) - - # Draw green plus signs for each keypoint - for x, y in kpts_np: - # Convert to integer coordinates - x, y = int(x), int(y) - - # Draw horizontal line of the plus sign - cv2.line(ret, (x - radius, y), (x + radius, y), green_color, width) - # Draw vertical line of the plus sign - cv2.line(ret, (x, y - radius), (x, y + radius), green_color, width) - - return ret - - -def visualize_keypoints(img_path, vis_path, detector: Detector, num_keypoints: int): - img_path, vis_path = Path(img_path), Path(vis_path).with_suffix(".png") - img = Image.open(img_path).convert("RGB") - detections = detector.detect_from_path( - img_path, num_keypoints=num_keypoints, return_dense_probs=True - ) - W, H = img.size - kps = detections["keypoints"] - kps = detector.to_pixel_coords(kps, H, W) - (vis_path).parent.mkdir(parents=True, exist_ok=True) - Image.fromarray(draw_kpts(img, kps[0])).save(vis_path) - if detections.get("dense_probs") is not None: - tensor_to_pil(detections["dense_probs"].squeeze().cpu(), autoscale=True).save( - vis_path.as_posix().replace(".png", "_dense_probs.png") - ) - - -def run_qualitative_examples( - *, model: Detector, workspace_path: str | Path, test_num_keypoints: int, step: int -): - workspace_path = Path(workspace_path) - torch.cuda.empty_cache() - for im_path in [ - "assets/0015_A.jpg", - "assets/0015_B.jpg", - "assets/0032_A.jpg", - "assets/0032_B.jpg", - "assets/apprentices.jpg", - "assets/rectangles_and_circles.png", - ]: - visualize_keypoints( - im_path, - workspace_path / "vis" / str(step) / im_path, - model, - num_keypoints=test_num_keypoints, - ) - torch.cuda.empty_cache() diff --git a/src/baselines/loma/device.py b/src/baselines/loma/device.py deleted file mode 100644 index bbd839f..0000000 --- a/src/baselines/loma/device.py +++ /dev/null @@ -1,13 +0,0 @@ -import torch - -if torch.cuda.is_available(): - device = torch.device("cuda") - amp_dtype = ( - torch.bfloat16 if torch.cuda.get_device_capability()[0] >= 8 else torch.float16 - ) -elif torch.backends.mps.is_available(): - device = torch.device("mps") - amp_dtype = torch.float16 -else: - device = torch.device("cpu") - amp_dtype = torch.bfloat16 \ No newline at end of file diff --git a/src/baselines/loma/geometry.py b/src/baselines/loma/geometry.py deleted file mode 100644 index 3a0fa1f..0000000 --- a/src/baselines/loma/geometry.py +++ /dev/null @@ -1,849 +0,0 @@ -from __future__ import annotations -from typing import Literal -import cv2 -import torch -import torch.nn.functional as F -from einops import einsum -import numpy as np - -from .device import device -from .types import Warp, GTSource, Batch - - -def to_homogeneous(x: torch.Tensor) -> torch.Tensor: - return torch.cat((x, torch.ones_like(x[..., :1])), dim=-1) - - -def from_homogeneous(x: torch.Tensor) -> torch.Tensor: - return x[..., :-1] / x[..., -1:] - - -def get_normalized_grid( - B: int, - H: int, - W: int, - overload_device: torch.device | None = None, -) -> torch.Tensor: - if B == 0: - return torch.zeros(0, H, W, 2, device=overload_device or device) - x1_n = torch.meshgrid( - *[ - torch.linspace(-1 + 1 / n, 1 - 1 / n, n, device=overload_device or device) - for n in (B, H, W) - ], - indexing="ij", - ) - x1_n = torch.stack((x1_n[2], x1_n[1]), dim=-1).reshape(B, H, W, 2) - return x1_n - - -def get_pixel_grid( - B: int, - *, - H: int, - W: int, - overload_device: torch.device | None = None, -) -> torch.Tensor: - x1_n = torch.meshgrid( - *[torch.arange(n, device=overload_device or device) + 0.5 for n in (B, H, W)], - indexing="ij", - ) - x1_n = torch.stack((x1_n[2], x1_n[1]), dim=-1).reshape(B, H, W, 2) - return x1_n - - -def to_normalized(x: torch.Tensor, *, H: int, W: int) -> torch.Tensor: - assert x.shape[-1] == 2, "x must have shape (..., 2)" - return torch.stack((2 * x[..., 0] / W, 2 * x[..., 1] / H), dim=-1) - 1 - - -def to_pixel(x: torch.Tensor, *, H: int, W: int) -> torch.Tensor: - assert x.shape[-1] == 2, "x must have shape (..., 2)" - return torch.stack(((x[..., 0] + 1) / 2 * W, (x[..., 1] + 1) / 2 * H), dim=-1) - - -def _pixel_warp_and_depth_from_depth( - *, - pixel_coords_A: torch.Tensor, - depth_A: torch.Tensor, - K_A: torch.Tensor, - K_B: torch.Tensor, - T_AB: torch.Tensor, -) -> tuple[torch.Tensor, torch.Tensor]: - homog_pixel_coords_A = to_homogeneous(pixel_coords_A) - x = einsum( - homog_pixel_coords_A, - K_A.inverse(), - "batch H W pixel, batch calib pixel -> batch H W calib", - ) - x = x * depth_A - x = to_homogeneous(x) - x = einsum(x, T_AB, "batch H W calib, batch world calib -> batch H W world") - x_AB = from_homogeneous(x) - x = einsum(x_AB, K_B, "batch H W world, batch pixel world -> batch H W pixel") - pixel_coords_AB = from_homogeneous(x) - depth = x_AB[..., -1:] - return pixel_coords_AB, depth - - -def warp_and_depth_consistency_from_depths( - *, - depth_A: torch.Tensor, - depth_B: torch.Tensor, - K_A: torch.Tensor, - K_B: torch.Tensor, - T_AB: torch.Tensor, - rel_depth_error_threshold: float, - interp_rtol: float | None = None, - interp_atol: float | None = None, - pos_depth_check_bugfix: bool = True, - local_neighbourhood_size: int = 1, -) -> Warp: - B, H_B, W_B, one = depth_B.shape - assert one == 1, "depth_A and depth_B must have shape (B, H, W, 1)" - B, H_A, W_A, one = depth_A.shape - assert one == 1, "depth_A and depth_B must have shape (B, H, W, 1)" - pixel_coords_A = get_pixel_grid(B, H=H_A, W=W_A, overload_device=device) - pixel_coords_AB, z_AB = _pixel_warp_and_depth_from_depth( - pixel_coords_A=pixel_coords_A, - depth_A=depth_A, - K_A=K_A, - K_B=K_B, - T_AB=T_AB, - ) - - normalized_coords_AB = to_normalized(pixel_coords_AB, H=H_B, W=W_B) - - # Compute covisibility by checking local neighborhood - K = local_neighbourhood_size - K2 = K * K - pad_size = K // 2 - - # Pad and unfold to get all local neighbors of depth_B - # depth_B: (B, H_B, W_B, 1) -> (B, 1, H_B, W_B) - depth_B_bchw = depth_B.permute(0, 3, 1, 2) - pad_depth_B = F.pad( - depth_B_bchw, (pad_size, pad_size, pad_size, pad_size), mode="reflect" - ) - # unfold: (B, 1, H_B+2*pad, W_B+2*pad) -> (B, 1, H_B, W_B, K, K) - unfolded = pad_depth_B.unfold(2, K, 1).unfold(3, K, 1) - # Reshape to (B*K2, H_B, W_B, 1) by stacking neighbors in batch dimension - unfolded = ( - unfolded.reshape(B, 1, H_B, W_B, K2) - .permute(0, 4, 2, 3, 1) - .reshape(B * K2, H_B, W_B, 1) - ) - - # Repeat sampling coordinates and depth values K2 times in batch dimension - normalized_coords_AB_repeated = normalized_coords_AB.repeat_interleave(K2, dim=0) - z_AB_repeated = z_AB.repeat_interleave(K2, dim=0) - depth_A_repeated = depth_A.repeat_interleave(K2, dim=0) - - # Sample from unfolded depth_B at warped coordinates - z_B = bhwc_grid_sample_with_nearest_exact_fallback( - x=unfolded, - grid=normalized_coords_AB_repeated, - rtol=interp_rtol, - atol=interp_atol, - mode="nearest", - align_corners=False, - ) - - # Compute consistency for all neighbors - if pos_depth_check_bugfix: - pos_and_finite_depth = ( - (depth_A_repeated > 0.0) - .logical_and(depth_A_repeated.isfinite()) - .logical_and(z_AB_repeated > 0.0) - .logical_and(z_AB_repeated.isfinite()) - .logical_and(z_B > 0.0) - .logical_and(z_B.isfinite()) - ) - else: - pos_and_finite_depth = ( - (z_AB_repeated > 0.0) - .logical_and(z_AB_repeated.isfinite()) - .logical_and(~z_B.isnan()) - ) - - rel_depth_error = ((z_B - z_AB_repeated) / z_B).abs() - - consistent_depth = (rel_depth_error < rel_depth_error_threshold).logical_and( - pos_and_finite_depth - ) - - # Reshape and aggregate over all neighbors: (B*K2, H_A, W_A, 1) -> (B, K2, H_A, W_A, 1) -> (B, H_A, W_A, 1) - covis_bool = consistent_depth.reshape(B, K2, H_A, W_A, 1).any(dim=1) - rel_depth_error = rel_depth_error.reshape(B, K2, H_A, W_A, 1).amin(dim=1) - pos_and_finite_depth = pos_and_finite_depth.reshape(B, K2, H_A, W_A, 1).any(dim=1) - - covis = covis_bool.float() - pos_depth_and_within_frame = (pos_and_finite_depth).logical_and( - normalized_coords_AB.abs().amax(dim=-1, keepdim=True) < 1.0 - ) - # Always compute valid (equivalent to mode "frame") - valid = pos_depth_and_within_frame.float() - # if rel_depth_error.isnan().any(): - # print("rel_depth_error is nan") - return Warp( - warp=normalized_coords_AB, covis=covis, valid=valid, error=rel_depth_error - ) - - -def warp_and_cycle_consistency_from_depths( - *, - depth_A: torch.Tensor, - depth_B: torch.Tensor, - K_A: torch.Tensor, - K_B: torch.Tensor, - T_AB: torch.Tensor, - cycle_error_threshold: float, - interp_rtol: float | None = None, - interp_atol: float | None = None, - multi_hyp: bool = False, -) -> Warp: - B, H_A, W_A, _ = depth_A.shape - B, H_B, W_B, _ = depth_B.shape - - pixel_coords_A = get_pixel_grid(B, H=H_A, W=W_A, overload_device=device) - pixel_coords_AB, depth_AB = _pixel_warp_and_depth_from_depth( - pixel_coords_A=pixel_coords_A, - depth_A=depth_A, - K_A=K_A, - K_B=K_B, - T_AB=T_AB, - ) - normalized_coords_AB = to_normalized(pixel_coords_AB, H=H_B, W=W_B) - depth_AB_B = bhwc_grid_sample_with_nearest_exact_fallback( - x=depth_B, - grid=normalized_coords_AB, - rtol=interp_rtol, - atol=interp_atol, - mode="bilinear", - align_corners=False, - ) - pos_and_finite_depth = (depth_AB_B > 0.0).logical_and(depth_AB_B.isfinite()) - pos_depth_and_within_frame = (pos_and_finite_depth).logical_and( - normalized_coords_AB.abs().amax(dim=-1, keepdim=True) < 1.0 - ) - valid = pos_depth_and_within_frame.float() - pixel_coords_AB_BA, _ = _pixel_warp_and_depth_from_depth( - pixel_coords_A=pixel_coords_AB, - depth_A=depth_AB_B, - K_A=K_B, - K_B=K_A, - T_AB=T_AB.inverse(), - ) - cycle_pixel_error = (pixel_coords_AB_BA - pixel_coords_A).norm(dim=-1, keepdim=True) - cycle_consistent = (cycle_pixel_error < cycle_error_threshold).logical_and( - pos_and_finite_depth - ) - return Warp( - warp=normalized_coords_AB, - covis=cycle_consistent.float(), - valid=valid, - error=cycle_pixel_error, - ) - - -def bhwc_interpolate( - x: torch.Tensor, - size: tuple[int, int], - mode: str = "bilinear", - align_corners: bool | None = None, - antialias: bool | None = None, -) -> torch.Tensor: - x_chw = x.permute(0, 3, 1, 2) - if align_corners is not None and antialias is not None: - y = F.interpolate( - x_chw, - size=size, - mode=mode, - align_corners=align_corners, - antialias=antialias, - ) - elif align_corners is not None: - y = F.interpolate(x_chw, size=size, mode=mode, align_corners=align_corners) - elif antialias is not None: - y = F.interpolate(x_chw, size=size, mode=mode, antialias=antialias) - else: - y = F.interpolate(x_chw, size=size, mode=mode) - return y.permute(0, 2, 3, 1) - - -def bhwc_interpolate_with_nearest_exact_fallback( - x: torch.Tensor, - size: tuple[int, int], - mode: str = "bilinear", - align_corners: bool | None = None, - rtol: float | None = None, - atol: float | None = None, -) -> torch.Tensor: - x_interp = bhwc_interpolate(x, size, mode, align_corners) - if rtol is None and atol is None: - return x_interp - x_nearest = bhwc_interpolate(x, size, "nearest-exact") - on_boundary = torch.isclose( - x_interp, - x_nearest, - rtol=rtol if rtol is not None else 0.0, - atol=atol if atol is not None else 0.0, - ).logical_not() - x_interp[on_boundary] = x_nearest[on_boundary] - return x_interp - - -def bhwc_grid_sample( - x: torch.Tensor, - grid: torch.Tensor, - mode: str = "bilinear", - align_corners: bool | None = False, -) -> torch.Tensor: - return F.grid_sample( - x.permute(0, 3, 1, 2), grid, mode=mode, align_corners=align_corners - ).permute(0, 2, 3, 1) - - -def bhwc_grid_sample_with_nearest_exact_fallback( - *, - x: torch.Tensor, - grid: torch.Tensor, - mode: str = "bilinear", - align_corners: bool | None = None, - rtol: float | None = None, - atol: float | None = None, -) -> torch.Tensor: - x_interp = F.grid_sample( - x.permute(0, 3, 1, 2), grid, mode=mode, align_corners=align_corners - ).permute(0, 2, 3, 1) - if rtol is None and atol is None: - return x_interp - x_nearest = F.grid_sample( - x.permute(0, 3, 1, 2), grid, mode="nearest", align_corners=align_corners - ).permute(0, 2, 3, 1) - x = x_interp - on_boundary = torch.isclose( - x_interp, - x_nearest, - rtol=rtol if rtol is not None else 0.0, - atol=atol if atol is not None else 0.0, - ).logical_not() - # at motion boundaries we fall back to nearest-exact - # note, this is pointwise - x[on_boundary] = x_nearest[on_boundary] - return x - - -def angle_error_mat(R1, R2): - cos = (np.trace(np.dot(R1.T, R2)) - 1) / 2 - cos = np.clip(cos, -1.0, 1.0) # numercial errors can make it out of bounds - return np.rad2deg(np.abs(np.arccos(cos))) - - -def angle_error_vec(v1, v2): - n = np.linalg.norm(v1) * np.linalg.norm(v2) - return np.rad2deg(np.arccos(np.clip(np.dot(v1, v2) / n, -1.0, 1.0))) - - -def compute_pose_error(R_gt, t_gt, R, t): - error_t = angle_error_vec(t.squeeze(), t_gt) - error_t = np.minimum(error_t, 180 - error_t) # ambiguity of E estimation - error_R = angle_error_mat(R, R_gt) - return error_t, error_R - - -def pose_auc(errors, thresholds): - sort_idx = np.argsort(errors) - errors = np.array(errors.copy())[sort_idx] - recall = (np.arange(len(errors)) + 1) / len(errors) - errors = np.r_[0.0, errors] - recall = np.r_[0.0, recall] - aucs = [] - for t in thresholds: - last_index = np.searchsorted(errors, t) - r = np.r_[recall[:last_index], recall[last_index - 1]] - e = np.r_[errors[:last_index], t] - aucs.append(np.trapezoid(r, x=e) / t) # type: ignore - return aucs - - -def cov_mat_from_cov_params(c: torch.Tensor) -> torch.Tensor: - return prec_mat_from_prec_params(c) - - -def prec_mat_from_prec_params(p: torch.Tensor) -> torch.Tensor: - P = p.new_zeros(p.shape[0], p.shape[1], p.shape[2], 2, 2) - P[..., 0, 0] = p[..., 0] - P[..., 1, 0] = p[..., 1] - P[..., 0, 1] = p[..., 1] - P[..., 1, 1] = p[..., 2] - return P - - -def to_double_angle_rep(v: torch.Tensor) -> torch.Tensor: - angle = torch.atan2(v[..., 1], v[..., 0]) - double_angle_rep = torch.stack((torch.cos(2 * angle), torch.sin(2 * angle)), dim=-1) - return double_angle_rep - - -def prec_mat_to_flow( - P: torch.Tensor, vis_max: float, mode: Literal["smallest", "largest"] = "largest" -) -> torch.Tensor: - vals, vecs = torch.linalg.eigh(P) - if mode == "smallest": - vis_vec = vecs[..., 0] - elif mode == "largest": - vis_vec = vecs[..., -1] - else: - raise ValueError(f"Invalid mode: {mode}") - # select upper half-plane - vis_vec = torch.where(vis_vec[..., 1:2] >= 0, vis_vec, -vis_vec) - double_angle_rep = to_double_angle_rep(vis_vec) - scale = (vals[..., 0] * vals[..., 1]).pow(0.25).clamp(0, vis_max)[..., None] - flow = scale * double_angle_rep - return flow - - -def prec_params_to_flow( - p: torch.Tensor, vis_max: float, mode: Literal["smallest", "largest"] = "largest" -) -> torch.Tensor: - P = prec_mat_from_prec_params(p) - return prec_mat_to_flow(P, vis_max, mode) - - -def flow_forward_backward_error( - flow_AB: torch.Tensor, - flow_BA: torch.Tensor, - *, - interp_rtol: float | None = None, - interp_atol: float | None = None, -) -> torch.Tensor: - B, H, W, _ = flow_AB.shape - grid = get_normalized_grid(B, H, W, overload_device=flow_AB.device) - fwd_bwd = bhwc_grid_sample_with_nearest_exact_fallback( - x=flow_BA, - grid=grid + flow_AB, - rtol=interp_rtol, - atol=interp_atol, - mode="nearest", - # mode="bilinear", - # align_corners=False, - ) - fwd_bwd_error = (flow_AB + fwd_bwd).norm(dim=-1, keepdim=True) - return fwd_bwd_error - - -def warp_and_overlap_from_flows( - *, - flow_AB: torch.Tensor, - flow_BA: torch.Tensor, - interp_rtol: float | None = None, - interp_atol: float | None = None, - error_threshold: float, - # check_local_neighbourhood: bool = False, - local_neighbourhood_size: int = 1, -) -> Warp: - # flow_AB and flow_BA are flows (offsets), convert to warps for processing - B, H, W, _ = flow_AB.shape - grid = get_normalized_grid(B, H, W, overload_device=flow_AB.device) - grid_A = grid - # grid_B = grid - - warp_AB = grid_A + flow_AB - # Check if warp is in normalized coordinates [-1, 1] - - in_frame = warp_AB.abs().le(1.0).all(dim=-1, keepdim=True) - - # Compute covisibility by checking local neighborhood - # if check_local_neighbourhood: - K = local_neighbourhood_size - K2 = K * K - pad_size = K // 2 - - # Pad and unfold to get all local neighbors - # flow_BA: (B, H, W, 2) -> (B, 2, H, W) - flow_BA_bchw = flow_BA.permute(0, 3, 1, 2) - pad_flow_BA = F.pad( - flow_BA_bchw, (pad_size, pad_size, pad_size, pad_size), mode="reflect" - ) - # unfold: (B, 2, H+2*pad, W+2*pad) -> (B, 2, H, W, K, K) - unfolded = pad_flow_BA.unfold(2, K, 1).unfold(3, K, 1) - # Reshape to (B*K2, H, W, 2) by stacking neighbors in batch dimension - unfolded = ( - unfolded.reshape(B, 2, H, W, K2).permute(0, 4, 2, 3, 1).reshape(B * K2, H, W, 2) - ) - - # Repeat flow_AB K2 times in batch dimension (interleaved to match unfolded order) - flow_AB_repeated = flow_AB.repeat_interleave(K2, dim=0) - - # Compute forward-backward error for all neighbors at once - fwd_bwd_error = flow_forward_backward_error( - flow_AB_repeated, unfolded, interp_rtol=interp_rtol, interp_atol=interp_atol - ) - covis = fwd_bwd_error < error_threshold - - # Reshape and take logical_or over all neighbors: (B*K2, H, W, 1) -> (B, K2, H, W, 1) -> (B, H, W, 1) - overlap_covis_bool = covis.reshape(B, K2, H, W, 1).any(dim=1).logical_and(in_frame) - fwd_bwd_error = fwd_bwd_error.reshape(B, K2, H, W, 1).amin(dim=1) * in_frame.float() - # else: - # fwd_bwd_error = flow_forward_backward_error( - # flow_AB, flow_BA, interp_rtol=interp_rtol, interp_atol=interp_atol - # ) - # overlap_covis_bool = (fwd_bwd_error < error_threshold).logical_and(in_frame) - - return Warp( - warp=warp_AB, - covis=overlap_covis_bool.float(), - valid=in_frame.float(), - error=fwd_bwd_error, - ) - - -def compute_gt_warp_from_batch( - batch: Batch, - *, - depth_error_threshold: float, - flow_error_threshold: float, - local_neighbourhood_size: int, -) -> Warp: - assert isinstance(batch.source, list), "source must be a list" - sources: list[GTSource] = batch.source - - is_flow_source = torch.tensor( - [source == "flow" for source in sources], device=device - ) - B, three, H_A, W_A = batch.img_A.shape - B, three, H_B, W_B = batch.img_B.shape - depth_A = batch.depth_A[~is_flow_source] - depth_B = batch.depth_B[~is_flow_source] - K_A = batch.K_A[~is_flow_source] - K_B = batch.K_B[~is_flow_source] - T_AB = batch.T_AB[~is_flow_source] - flow_flow_source_AB = batch.flow_AB[is_flow_source] - flow_flow_source_BA = batch.flow_BA[is_flow_source] - - result_depth = warp_and_depth_consistency_from_depths( - depth_A=depth_A, - depth_B=depth_B, - K_A=K_A, - K_B=K_B, - T_AB=T_AB, - rel_depth_error_threshold=depth_error_threshold, - local_neighbourhood_size=local_neighbourhood_size, - ) - overlap_depth_src = result_depth.covis - warp_depth_src_AB = result_depth.warp - valid_depth_src = result_depth.valid - error_depth_src = result_depth.error - - result_flow = warp_and_overlap_from_flows( - flow_AB=flow_flow_source_AB, - flow_BA=flow_flow_source_BA, - error_threshold=flow_error_threshold, - local_neighbourhood_size=local_neighbourhood_size, - ) - overlap_flow_src = result_flow.covis - warp_flow_src_AB = result_flow.warp - valid_flow_src = result_flow.valid - error_flow_src = result_flow.error - assert overlap_flow_src is not None and overlap_depth_src is not None - assert warp_flow_src_AB is not None and warp_depth_src_AB is not None - assert valid_flow_src is not None and valid_depth_src is not None - assert error_flow_src is not None and error_depth_src is not None - - covis = torch.zeros((B, H_A, W_A, 1), dtype=torch.float32, device=device) - covis[is_flow_source] = overlap_flow_src.float() - covis[~is_flow_source] = overlap_depth_src.float() - - warp = torch.zeros((B, H_A, W_A, 2), device=device, dtype=torch.float32) - warp[is_flow_source] = warp_flow_src_AB - warp[~is_flow_source] = warp_depth_src_AB - warp[~warp.isfinite()] = 0 - - valid = torch.zeros((B, H_A, W_A, 1), device=device, dtype=torch.float32) - valid[is_flow_source] = valid_flow_src - valid[~is_flow_source] = valid_depth_src - - error = torch.zeros((B, H_A, W_A, 1), device=device, dtype=torch.float32) - error[is_flow_source] = error_flow_src - error[~is_flow_source] = error_depth_src - - return Warp(warp=warp, covis=covis, valid=valid, error=error) - - -# Code taken from https://github.com/PruneTruong/DenseMatching/blob/40c29a6b5c35e86b9509e65ab0cd12553d998e5f/validation/utils_pose_estimation.py -# --- GEOMETRY --- -def estimate_pose_cv2_ransac( - kpts0: np.ndarray, - kpts1: np.ndarray, - K0: np.ndarray, - K1: np.ndarray, - norm_thresh: float, - conf: float = 0.99999, -): - import cv2 - - if len(kpts0) < 5: - return None - K0inv = np.linalg.inv(K0[:2, :2]) - K1inv = np.linalg.inv(K1[:2, :2]) - - kpts0 = (K0inv @ (kpts0 - K0[None, :2, 2]).T).T - kpts1 = (K1inv @ (kpts1 - K1[None, :2, 2]).T).T - E, mask = cv2.findEssentialMat( - kpts0, kpts1, np.eye(3), threshold=norm_thresh, prob=conf - ) - - ret = None - if E is not None: - best_num_inliers = 0 - - for _E in np.split(E, len(E) // 3): - n, R, t, _ = cv2.recoverPose(_E, kpts0, kpts1, np.eye(3), 1e9, mask=mask) # type: ignore - if n > best_num_inliers: - best_num_inliers = n - ret = (R, t, mask.ravel() > 0) - assert ret is not None - return ret - - -def compute_pose_inliers_cv2_ransac( - kpts0: np.ndarray, - kpts1: np.ndarray, - K0: np.ndarray, - K1: np.ndarray, - norm_thresh: float, - conf: float = 0.99999, -) -> tuple[float, tuple[np.ndarray, np.ndarray, np.ndarray]] | None: - """Inlier ratio and best pose from the same RANSAC/E hypothesis as `estimate_pose_cv2_ransac`.""" - import cv2 - - if len(kpts0) < 5: - return None - K0inv = np.linalg.inv(K0[:2, :2]) - K1inv = np.linalg.inv(K1[:2, :2]) - - kpts0 = (K0inv @ (kpts0 - K0[None, :2, 2]).T).T - kpts1 = (K1inv @ (kpts1 - K1[None, :2, 2]).T).T - E, mask = cv2.findEssentialMat( - kpts0, kpts1, np.eye(3), threshold=norm_thresh, prob=conf - ) - - if E is None: - return None - best_num_inliers = 0 - ret: tuple[np.ndarray, np.ndarray, np.ndarray] | None = None - - for _E in np.split(E, len(E) // 3): - n, R, t, _ = cv2.recoverPose(_E, kpts0, kpts1, np.eye(3), 1e9, mask=mask) # type: ignore - if n > best_num_inliers: - best_num_inliers = n - ret = (R, t, mask.ravel() > 0) - assert ret is not None - return (best_num_inliers / len(kpts0), ret) - - -def estimate_pose_essential( - kps_A: np.ndarray, - kps_B: np.ndarray, - w_A: int, - h_A: int, - K_A: np.ndarray, - w_B: int, - h_B: int, - K_B: np.ndarray, - th: float, -) -> tuple[np.ndarray, np.ndarray]: - import poselib - - camera1 = { - "model": "PINHOLE", - "width": w_A, - "height": h_A, - "params": K_A[[0, 1, 0, 1], [0, 1, 2, 2]], - } - camera2 = { - "model": "PINHOLE", - "width": w_B, - "height": h_B, - "params": K_B[[0, 1, 0, 1], [0, 1, 2, 2]], - } - - pose, res = poselib.estimate_relative_pose( - kps_A, - kps_B, - camera1, - camera2, - ransac_opt={ - "max_epipolar_error": th, - }, - ) - return pose.R, pose.t - - -def poselib_fundamental(x1, x2, opt): - import poselib - - F, info = poselib.estimate_fundamental(x1, x2, opt, {}) - inl = info["inliers"] - return F, inl - - -def calibrate(x: torch.Tensor, K: torch.Tensor): - # x: ..., 2 - # K: ..., 3, 3 - return to_homogeneous(x) @ K.inverse().mT - - -def estimate_pose_fundamental( - kps_A: np.ndarray, - kps_B: np.ndarray, - w_A: int, - h_A: int, - K_A: np.ndarray, - w_B: int, - h_B: int, - K_B: np.ndarray, - th: float, -) -> tuple[np.ndarray, np.ndarray]: - if len(kps_A) < 8: - return np.eye(3), np.zeros(3) - F, inl = poselib_fundamental( - kps_A, - kps_B, - opt={ - "max_epipolar_error": th, - }, - ) - E: np.ndarray = K_B.T @ F @ K_A - kps_calib_A = from_homogeneous( - calibrate(torch.from_numpy(kps_A).float(), torch.from_numpy(K_A).float()) - ).numpy() - kps_calib_B = from_homogeneous( - calibrate(torch.from_numpy(kps_B).float(), torch.from_numpy(K_B).float()) - ).numpy() - E = E.astype(np.float64) - _, R, t, good = cv2.recoverPose(E, kps_calib_A, kps_calib_B) - t = t[:, 0] - return R, t - - -def compute_relative_pose(R1, t1, R2, t2): - rots = R2 @ (R1.mT) - trans = -rots @ t1 + t2 - return rots, trans - - -def kde(x: torch.Tensor, std: float = 0.1, half: bool = True) -> torch.Tensor: - # use a gaussian kernel to estimate density - if half: - x = x.half() - scores = (-(torch.cdist(x, x) ** 2) / (2 * std**2)).exp() - density = scores.sum(dim=-1) - return density - - -def normalize(x: torch.Tensor, dim: int): - return x / x.norm(dim=dim, keepdim=True) - - -def cosine_similarity(f_A: torch.Tensor, f_B: torch.Tensor) -> torch.Tensor: - # TODO: remove this? - f_A = normalize(f_A, dim=-1) - f_B = normalize(f_B, dim=-1) - res = einsum(f_A, f_B, "B H_A W_A D, B H_B W_B D -> B H_A W_A H_B W_B") - return res - - -@torch.no_grad() -def compute_sparse_mnn_matches( - kpts_A, - kpts_B, - batch, - *, - depth_error_threshold: float, - flow_error_threshold: float, - local_neighbourhood_size: int, - error_threshold: float, - gt_AB: Warp | None = None, - gt_BA: Warp | None = None, -): - """ - Compute mutual nearest neighbor matches. - - For samples with source == "sparse", uses correspondences directly (keypoints should be - [correspondences, detector_keypoints] concatenated, and MNN points to correspondence indices). - - For samples with source == "depth" or "flow", computes MNN using depth/flow warping. - - Args: - gt_AB: Optional precomputed warp from A to B. If provided, skips recomputation. - gt_BA: Optional precomputed warp from B to A. If provided, skips recomputation. - """ - B = kpts_A.shape[0] - device = kpts_A.device - - source = batch.source if isinstance(batch.source, list) else [batch.source] * B - num_corresp = ( - batch.num_corresp - if isinstance(batch.num_corresp, list) - else [batch.num_corresp] * B - ) - - is_sparse = torch.tensor([s == "sparse" for s in source], device=device) - all_mnn = [] - - # Handle sparse samples: MNN points to correspondence indices (idx_A == idx_B) - if is_sparse.any(): - sparse_counts = torch.tensor( - [num_corresp[b] if is_sparse[b] else 0 for b in range(B)], - device=device, - dtype=torch.long, - ) - total_sparse = sparse_counts.sum() - if total_sparse > 0: - batch_idx = torch.repeat_interleave( - torch.arange(B, device=device), sparse_counts - ) - corresp_idx = torch.cat( - [torch.arange(n, device=device) for n in sparse_counts.tolist()] - ) - sparse_mnn = torch.stack([batch_idx, corresp_idx, corresp_idx], dim=1) - all_mnn.append(sparse_mnn) - - # Handle warp-based samples (depth/flow) - if (~is_sparse).any(): - # Use precomputed warps if provided, otherwise compute them - if gt_AB is None: - gt_AB = compute_gt_warp_from_batch( - batch, - depth_error_threshold=depth_error_threshold, - flow_error_threshold=flow_error_threshold, - local_neighbourhood_size=local_neighbourhood_size, - ) - if gt_BA is None: - gt_BA = compute_gt_warp_from_batch( - batch.swap_AB(), - depth_error_threshold=depth_error_threshold, - flow_error_threshold=flow_error_threshold, - local_neighbourhood_size=local_neighbourhood_size, - ) - kpts_A_to_B = bhwc_grid_sample(gt_AB.warp, kpts_A[:, None])[:, 0] # (B, N, 2) - kpts_B_to_A = bhwc_grid_sample(gt_BA.warp, kpts_B[:, None])[:, 0] # (B, N, 2) - D_B = torch.cdist(kpts_A_to_B, kpts_B).nan_to_num(nan=float("inf")) - D_A = torch.cdist(kpts_A, kpts_B_to_A).nan_to_num(nan=float("inf")) - warp_mnn = torch.nonzero( - (D_B == D_B.min(dim=-1, keepdim=True).values) - * (D_A == D_A.min(dim=-2, keepdim=True).values) - * (D_B < error_threshold) - * (D_A < error_threshold) - ) - # Filter to only keep non-sparse samples - if warp_mnn.numel() > 0: - warp_mnn = warp_mnn[~is_sparse[warp_mnn[:, 0]]] - if warp_mnn.numel() > 0: - all_mnn.append(warp_mnn) - - if all_mnn: - return torch.cat(all_mnn, dim=0) - return torch.zeros((0, 3), device=device, dtype=torch.long) \ No newline at end of file diff --git a/src/baselines/loma/io.py b/src/baselines/loma/io.py deleted file mode 100644 index 4905b20..0000000 --- a/src/baselines/loma/io.py +++ /dev/null @@ -1,30 +0,0 @@ -import numpy as np -from PIL import Image -import warnings - - -def numpy_to_pil(x: np.ndarray): - assert x.dtype in [np.float32, np.uint8] - if x.dtype == np.float32: - assert x.min() >= 0.0 and x.max() <= 1.0 - x = (x * 255).astype(np.uint8) - return Image.fromarray(x) - - -def tensor_to_pil(x, autoscale=False): - x = x.clone().detach().cpu().numpy() - if autoscale: - if x.max() == x.min(): - warnings.warn("x max == x min, cant autoscale") - else: - x = (x - x.min()) / (x.max() - x.min()) - - if len(x.shape) == 3: - x = np.transpose(x, (1, 2, 0)) - x = np.clip(x, 0.0, 1.0) - return numpy_to_pil(x) - - -def check_not_i16(pil_img: Image.Image): - if pil_img.mode == "I;16": - raise NotImplementedError("Can't handle 16 bit images") \ No newline at end of file diff --git a/src/baselines/loma/loma.py b/src/baselines/loma/loma.py deleted file mode 100644 index f19df15..0000000 --- a/src/baselines/loma/loma.py +++ /dev/null @@ -1,498 +0,0 @@ -from dataclasses import dataclass -import re -import sys -from typing import Callable, Literal, Tuple -import numpy as np -from PIL import Image - -import torch -import torch.nn.functional as F -import torch.nn as nn - -from .descriptor.dedode import DeDoDeDescriptor -from .detector.dad import DaD -from .device import device, amp_dtype -from .types import Model - -# Reference code from LightGlue: https://github.com/cvg/LightGlue/blob/main/lightglue/lightglue.py -def rotate_half(x: torch.Tensor) -> torch.Tensor: - x = x.unflatten(-1, (-1, 2)) - x1, x2 = x.unbind(dim=-1) - return torch.stack((-x2, x1), dim=-1).flatten(start_dim=-2) - - -def apply_cached_rotary_emb(freqs: torch.Tensor, t: torch.Tensor) -> torch.Tensor: - return (t * freqs[0]) + (rotate_half(t) * freqs[1]) - - -def to_pixel_coords(flow, h1, w1): - flow = torch.stack( - ( - w1 * (flow[..., 0] + 1) / 2, - h1 * (flow[..., 1] + 1) / 2, - ), - dim=-1, - ) - return flow - - -class LearnableFourierPositionalEncoding(nn.Module): - def __init__( - self, - M: int, - dim: int, - F_dim: int | None = None, - *, - gamma: float, - ) -> None: - super().__init__() - F_dim = F_dim if F_dim is not None else dim - self.gamma = gamma - self.Wr = nn.Linear(M, F_dim // 2, bias=False) - nn.init.normal_(self.Wr.weight.data, mean=0, std=self.gamma**-2) - - def forward(self, x: torch.Tensor) -> torch.Tensor: - projected = self.Wr(x) - cosines, sines = torch.cos(projected), torch.sin(projected) - emb = torch.stack([cosines, sines], 0).unsqueeze(-3) - return emb.repeat_interleave(2, dim=-1) - - -class FixedPosEnc(nn.Module): - def __init__( - self, - M: int, - dim: int, - F_dim: int | None = None, - *, - gamma: float, - ) -> None: - super().__init__() - F_dim = F_dim if F_dim is not None else dim - self.gamma = gamma - freqs = torch.randn(F_dim // 2, M).to(device) * self.gamma**-2 - self.Wr = nn.Linear(M, F_dim // 2, bias=False) - with torch.no_grad(): - self.Wr.weight.data = freqs - self.Wr.weight.requires_grad = False - - def forward(self, x: torch.Tensor) -> torch.Tensor: - projected = self.Wr(x) - cosines, sines = torch.cos(projected), torch.sin(projected) - emb = torch.stack([cosines, sines], 0).unsqueeze(-3) - return emb.repeat_interleave(2, dim=-1) - - -class SelfBlock(nn.Module): - def __init__( - self, embed_dim: int, num_heads: int, flash: bool = False, bias: bool = True - ) -> None: - super().__init__() - self.embed_dim = embed_dim - self.num_heads = num_heads - assert self.embed_dim % num_heads == 0 - self.head_dim = self.embed_dim // num_heads - self.Wqkv = nn.Linear(embed_dim, 3 * embed_dim, bias=bias) - self.out_proj = nn.Linear(embed_dim, embed_dim, bias=bias) - self.ffn = nn.Sequential( - nn.Linear(2 * embed_dim, 2 * embed_dim), - nn.LayerNorm(2 * embed_dim, elementwise_affine=True), - nn.GELU(), - nn.Linear(2 * embed_dim, embed_dim), - ) - - def forward(self, x: torch.Tensor, encoding: torch.Tensor) -> torch.Tensor: - qkv = self.Wqkv(x) - qkv = qkv.unflatten(-1, (self.num_heads, -1, 3)).transpose(1, 2) - q, k, v = qkv[..., 0], qkv[..., 1], qkv[..., 2] - # if encoding is not None: - q = apply_cached_rotary_emb(encoding, q) - k = apply_cached_rotary_emb(encoding, k) - context = F.scaled_dot_product_attention(q, k, v) - message = self.out_proj(context.transpose(1, 2).flatten(start_dim=-2)) - return x + self.ffn(torch.cat([x, message], -1)) - - -class CrossBlock(nn.Module): - def __init__( - self, embed_dim: int, num_heads: int, flash: bool = False, bias: bool = True - ) -> None: - super().__init__() - self.heads = num_heads - dim_head = embed_dim // num_heads - self.scale = dim_head**-0.5 - inner_dim = dim_head * num_heads - self.to_qk = nn.Linear(embed_dim, inner_dim, bias=bias) - self.to_v = nn.Linear(embed_dim, inner_dim, bias=bias) - self.to_out = nn.Linear(inner_dim, embed_dim, bias=bias) - self.ffn = nn.Sequential( - nn.Linear(2 * embed_dim, 2 * embed_dim), - nn.LayerNorm(2 * embed_dim, elementwise_affine=True), - nn.GELU(), - nn.Linear(2 * embed_dim, embed_dim), - ) - - def map_(self, func: Callable, x0: torch.Tensor, x1: torch.Tensor): - return func(x0), func(x1) - - def forward( - self, x0: torch.Tensor, x1: torch.Tensor - ) -> tuple[torch.Tensor, torch.Tensor]: - qk0, qk1 = self.map_(self.to_qk, x0, x1) - v0, v1 = self.map_(self.to_v, x0, x1) - qk0, qk1, v0, v1 = map( - lambda t: t.unflatten(-1, (self.heads, -1)).transpose(1, 2), - (qk0, qk1, v0, v1), - ) - m0 = F.scaled_dot_product_attention(qk0, qk1, v1) - m1 = F.scaled_dot_product_attention(qk1, qk0, v0) - m0, m1 = self.map_(lambda t: t.transpose(1, 2).flatten(start_dim=-2), m0, m1) - m0, m1 = self.map_(self.to_out, m0, m1) - x0 = x0 + self.ffn(torch.cat([x0, m0], -1)) - x1 = x1 + self.ffn(torch.cat([x1, m1], -1)) - return x0, x1 - - -class TransformerLayer(nn.Module): - def __init__(self, *args, **kwargs): - super().__init__() - self.self_attn = SelfBlock(*args, **kwargs) - self.cross_attn = CrossBlock(*args, **kwargs) - - def forward(self, desc0, desc1, encoding0, encoding1): - desc0 = self.self_attn(desc0, encoding0) - desc1 = self.self_attn(desc1, encoding1) - return self.cross_attn(desc0, desc1) - - -def log_double_softmax( - sim: torch.Tensor, z0: torch.Tensor, z1: torch.Tensor -) -> torch.Tensor: - b, m, n = sim.shape - scores0 = F.log_softmax(sim, 2) - scores1 = F.log_softmax(sim.transpose(-1, -2).contiguous(), 2).transpose(-1, -2) - scores = sim.new_full((b, m + 1, n + 1), 0) - scores[:, :m, :n] = scores0 + scores1 - scores[:, :-1, -1] = z0.squeeze(-1) - scores[:, -1, :-1] = z1.squeeze(-1) - return scores - - -class MatchAssignment(nn.Module): - def __init__(self, dim: int) -> None: - super().__init__() - self.dim = dim - self.matchability = nn.Linear(dim, 1, bias=True) - self.final_proj = nn.Linear(dim, dim, bias=True) - - def forward(self, desc0: torch.Tensor, desc1: torch.Tensor): - mdesc0, mdesc1 = self.final_proj(desc0), self.final_proj(desc1) - _, _, d = mdesc0.shape - mdesc0, mdesc1 = mdesc0 / d**0.25, mdesc1 / d**0.25 - sim = torch.einsum("bmd,bnd->bmn", mdesc0, mdesc1) - # For training we supervise matchability but for inference we don't use it - # NOTE: This operation, double softmax, takes a lot of current inference time. - if self.training: - z0 = self.matchability(desc0) - z1 = self.matchability(desc1) - scores = log_double_softmax(sim, z0, z1) - else: - scores = F.softmax(sim, dim=2) * F.softmax(sim, dim=1) - return scores, sim - - -def filter_matches(scores: torch.Tensor, th: float): - max0, max1 = scores.max(2), scores.max(1) - m0, m1 = max0.indices, max1.indices - indices0 = torch.arange(m0.shape[1], device=m0.device)[None] - indices1 = torch.arange(m1.shape[1], device=m1.device)[None] - mutual0 = indices0 == m1.gather(1, m0) - mutual1 = indices1 == m0.gather(1, m1) - mscores0 = torch.where(mutual0, max0.values, max0.values.new_tensor(0)) - mscores1 = torch.where(mutual1, mscores0.gather(1, m1), mscores0.new_tensor(0)) - valid0 = mutual0 & (mscores0 > th) - valid1 = mutual1 & valid0.gather(1, m1) - m0 = torch.where(valid0, m0, -1) - m1 = torch.where(valid1, m1, -1) - return m0, m1, mscores0, mscores1 - - -class LoMa(Model): - @dataclass(frozen=True, kw_only=True) - class Cfg: - input_dim: int = 256 - embed_dim: int = 256 - n_layers: int = 9 - num_heads: int = 4 - filter_threshold: float = 0.1 - mp: bool = True - compile: bool = False - descriptor: Literal["dedode_b", "dedode_g"] = "dedode_b" - num_keypoints: int = 2048 - # Positional encoding config - posenc_type: Literal["learnable", "fixed"] = "learnable" - # basically the wavelength of the positional encoding - posenc_gamma: float = 1.0 - # Optional pretrained checkpoint URL. - weights_url: str | None = None - - def __init__(self, cfg: Cfg | None = None) -> None: - super().__init__() - if cfg is None: - cfg = LoMa.Cfg() - self.cfg = cfg - - if cfg.input_dim != cfg.embed_dim: - self.input_proj = nn.Linear(cfg.input_dim, cfg.embed_dim, bias=True) - else: - self.input_proj = nn.Identity() - - head_dim = cfg.embed_dim // cfg.num_heads - if cfg.posenc_type == "learnable": - self.posenc = LearnableFourierPositionalEncoding( - 2, - head_dim, - head_dim, - gamma=cfg.posenc_gamma, - ) - elif cfg.posenc_type == "fixed": - self.posenc = FixedPosEnc( - 2, - head_dim, - head_dim, - gamma=cfg.posenc_gamma, - ) - else: - raise ValueError( - f"Positional encoding type {cfg.posenc_type} not supported" - ) - - self.transformers = nn.ModuleList( - [ - TransformerLayer(cfg.embed_dim, cfg.num_heads) - for _ in range(cfg.n_layers) - ] - ) - self.log_assignment = nn.ModuleList( - [MatchAssignment(cfg.embed_dim) for _ in range(cfg.n_layers)] - ) - - self._detector = DaD(DaD.Cfg(compile=cfg.compile)).eval() - for p in self._detector.parameters(): - p.requires_grad = False - - self._descriptor = DeDoDeDescriptor( - DeDoDeDescriptor.Cfg( - arch=cfg.descriptor, - compile=cfg.compile, - descriptor_dim=cfg.input_dim, - ) - ).eval() - for p in self._descriptor.parameters(): - p.requires_grad = False - - self.to(device) - - if cfg.weights_url is not None: - weights = torch.hub.load_state_dict_from_url( - cfg.weights_url, - map_location=device, - ) - missing_keys, unexpected_keys = self.load_state_dict(weights, strict=False) - if len(unexpected_keys) > 0: - allowed_extra_layer_keys = { - key - for key in unexpected_keys - if self._is_unexpected_extra_layer_key(key, cfg.n_layers) - } - disallowed_keys = sorted( - set(unexpected_keys) - allowed_extra_layer_keys - ) - assert len(disallowed_keys) == 0, ( - "Unexpected keys when loading pretrained weights " - f"(not extra layers beyond n_layers={cfg.n_layers}): {disallowed_keys}" - ) - assert len(missing_keys) == 0, ( - f"Missing keys when loading pretrained weights: {missing_keys}" - ) - - self.eval() - if self.cfg.compile and sys.platform == "linux": - # macos compile slows inference down. - self.compile() - - @staticmethod - def _is_unexpected_extra_layer_key(key: str, n_layers: int) -> bool: - match = re.match(r"^(transformers|log_assignment)\.(\d+)\.", key) - if match is None: - return False - layer_idx = int(match.group(2)) - return layer_idx >= n_layers - - @torch.inference_mode() - def detect(self, batch, num_keypoints: int | None = None) -> dict: - """Detect keypoints using the frozen detector.""" - if num_keypoints is None: - num_keypoints = self.cfg.num_keypoints - return self._detector.detect(batch, num_keypoints=num_keypoints) - - @torch.inference_mode() - def describe(self, batch, keypoints: torch.Tensor) -> dict: - """Describe keypoints using the frozen descriptor.""" - images = torch.cat([batch.img_A, batch.img_B], dim=0) - return self._descriptor.describe_keypoints(images, keypoints) - - def forward( - self, - kpts0: torch.Tensor, - kpts1: torch.Tensor, - desc0: torch.Tensor, - desc1: torch.Tensor, - ) -> dict: - - with torch.autocast( - enabled=self.cfg.mp, dtype=amp_dtype, device_type=device.type - ): - kpts0 = kpts0.to(device) - kpts1 = kpts1.to(device) - desc0 = desc0.to(device).detach().contiguous() - desc1 = desc1.to(device).detach().contiguous() - desc0 = self.input_proj(desc0) - desc1 = self.input_proj(desc1) - encoding0 = self.posenc(kpts0) - encoding1 = self.posenc(kpts1) - scores = None - for i in range(self.cfg.n_layers): - desc0, desc1 = self.transformers[i](desc0, desc1, encoding0, encoding1) - scores, _ = self.log_assignment[i](desc0, desc1) - assert scores is not None - - return { - "scores": scores, - # "all_scores": all_scores, # for training you need to return intermediate scores for layer-wise supervision - } - - @torch.inference_mode() - def detect_and_describe( - self, image: str | torch.Tensor, num_keypoints: int | None = None - ) -> Tuple[torch.Tensor, torch.Tensor, int, int]: - """Returns (keypoints, descriptions, H, W) where H, W are the original image dimensions.""" - if num_keypoints is None: - num_keypoints = self.cfg.num_keypoints - if isinstance(image, str): - keypoints = self._detector.detect_from_path( - image, num_keypoints=num_keypoints - )["keypoints"] - descriptions = self._descriptor.describe_keypoints_from_path( - image, keypoints - )["descriptions"] - w, h = Image.open(image).size - else: - image = image.to(device) - batch = {"image": image} - keypoints = self._detector.detect(batch, num_keypoints=num_keypoints)[ - "keypoints" - ] - descriptions = self._descriptor.describe_keypoints( - batch["image"], keypoints - )["descriptions"] - h, w = image.shape[-2:] - return keypoints, descriptions, h, w - - @torch.inference_mode() - def match( - self, - image_A: str | torch.Tensor, - image_B: str | torch.Tensor, - filter_threshold: float | None = None, - num_keypoints: int | None = None, - ) -> Tuple[np.ndarray, np.ndarray]: - """Match keypoints between two images, returning pixel-coordinate matches. - - Args: - image_A: Path to image A (str) or image tensor (B, C, H, W). - image_B: Path to image B (str) or image tensor (B, C, H, W). - Returns: - Tuple of (kptsA, kptsB) as numpy arrays of pixel coordinates. - """ - keypoints_A, descriptors_A, h1, w1 = self.detect_and_describe( - image_A, num_keypoints - ) - keypoints_B, descriptors_B, h2, w2 = self.detect_and_describe( - image_B, num_keypoints - ) - # return None - if filter_threshold is None: - filter_threshold = self.cfg.filter_threshold - - scores = self(keypoints_A, keypoints_B, descriptors_A, descriptors_B)["scores"] - m0, _, _, _ = filter_matches(scores, filter_threshold) - - valid = m0[0] > -1 - matched_A = keypoints_A[0][torch.where(valid)[0]] - matched_B = keypoints_B[0][m0[0][valid]] - - return to_pixel_coords(matched_A, h1, w1).cpu().numpy(), to_pixel_coords( - matched_B, h2, w2 - ).cpu().numpy() - - -@dataclass(frozen=True, kw_only=True) -class LoMaB128(LoMa.Cfg): - name: Literal["loma_B128"] = "loma_B128" - input_dim: Literal[128] = 128 - embed_dim: Literal[256] = 256 - num_heads: Literal[4] = 4 - descriptor: Literal["dedode_b"] = "dedode_b" - weights_url: str = ( - "https://github.com/davnords/storage/releases/download/loma/loma_B128.pth" - ) - - -@dataclass(frozen=True, kw_only=True) -class LoMaB(LoMa.Cfg): - name: Literal["loma_B"] = "loma_B" - input_dim: Literal[256] = 256 - embed_dim: Literal[256] = 256 - num_heads: Literal[4] = 4 - descriptor: Literal["dedode_g"] = "dedode_g" - weights_url: str = ( - "https://github.com/davnords/storage/releases/download/loma/loma_B.pt" - ) - - -@dataclass(frozen=True, kw_only=True) -class LoMaL(LoMa.Cfg): - name: Literal["loma_L"] = "loma_L" - input_dim: Literal[256] = 256 - embed_dim: Literal[512] = 512 - num_heads: Literal[8] = 8 - descriptor: Literal["dedode_g"] = "dedode_g" - weights_url: str = ( - "https://github.com/davnords/storage/releases/download/loma/loma_L.pth" - ) - - -@dataclass(frozen=True, kw_only=True) -class LoMaG(LoMa.Cfg): - name: Literal["loma_G"] = "loma_G" - input_dim: Literal[256] = 256 - embed_dim: Literal[1024] = 1024 - num_heads: Literal[16] = 16 - descriptor: Literal["dedode_g"] = "dedode_g" - weights_url: str = ( - "https://github.com/davnords/storage/releases/download/loma/loma_G.pth" - ) - -@dataclass(frozen=True, kw_only=True) -class LoMaR(LoMa.Cfg): - name: Literal["loma_R"] = "loma_R" - input_dim: Literal[256] = 256 - embed_dim: Literal[256] = 256 - num_heads: Literal[4] = 4 - descriptor: Literal["dedode_g"] = "dedode_g" - weights_url: str = ( - "https://github.com/davnords/storage/releases/download/loma/loma_R.pth" - ) - -LoMaName = Literal["loma_B128", "loma_B", "loma_L", "loma_G", "loma_R"] \ No newline at end of file diff --git a/src/baselines/loma/types.py b/src/baselines/loma/types.py deleted file mode 100644 index 5f87fa9..0000000 --- a/src/baselines/loma/types.py +++ /dev/null @@ -1,125 +0,0 @@ -from __future__ import annotations -import torch.nn as nn -from dataclasses import dataclass, field -from pathlib import Path -import torch -from typing import Literal, Protocol - - -@dataclass -class Warp: - warp: torch.Tensor - covis: torch.Tensor | None = None - valid: torch.Tensor | None = None - error: torch.Tensor | None = None - - -@dataclass -class Batch: - img_A: torch.Tensor - img_B: torch.Tensor - depth_A: torch.Tensor - depth_B: torch.Tensor - img_A_path: Path - img_B_path: Path - source: GTSource | list[GTSource] - flow_AB: torch.Tensor - flow_BA: torch.Tensor - mask_AB: torch.Tensor - mask_BA: torch.Tensor - quality: Quality | list[Quality] - # Optional sparse pixel correspondences: (N, 4) = (x_A, y_A, x_B, y_B) in pixels. - correspondences_AB: torch.Tensor = field( - default_factory=lambda: torch.zeros((0, 4)) - ) - K_A: torch.Tensor = field(default_factory=lambda: torch.eye(3)) - K_B: torch.Tensor = field(default_factory=lambda: torch.eye(3)) - pose_A: torch.Tensor = field(default_factory=lambda: torch.eye(4)) - pose_B: torch.Tensor = field(default_factory=lambda: torch.eye(4)) - T_AB: torch.Tensor = field(default_factory=lambda: torch.eye(4)) - num_corresp: int | list[int] = 0 - - def to(self, device: torch.device) -> Batch: - # Preserve dynamic batch type (subclasses may add more tensor fields). - data = {} - for k, v in self.__dict__.items(): - if isinstance(v, torch.Tensor): - data[k] = v.to(device) - else: - data[k] = v - return type(self)(**data) - - def swap_AB(self) -> Batch: - """Return a "reversed" batch (B->A) useful for symmetric GT computation.""" - return Batch( - img_A=self.img_B, - img_B=self.img_A, - depth_A=self.depth_B, - depth_B=self.depth_A, - K_A=self.K_B, - K_B=self.K_A, - pose_A=self.pose_B, - pose_B=self.pose_A, - T_AB=torch.linalg.inv(self.T_AB), - img_A_path=self.img_B_path, - img_B_path=self.img_A_path, - source=self.source, - flow_AB=self.flow_BA, - flow_BA=self.flow_AB, - mask_AB=self.mask_BA, - mask_BA=self.mask_AB, - quality=self.quality, - correspondences_AB=self.correspondences_AB[..., [2, 3, 0, 1]], - num_corresp=self.num_corresp, - ) - - @classmethod - def collate(cls, samples: list[Batch]) -> Batch: - keys = samples[0].__dict__.keys() - batch = {} - for k in keys: - if k == "correspondences_AB": - corresps = 2 * torch.rand((len(samples), MAX_NUM_SPARSE_CORRESP, 4)) - 1 - for i, s in enumerate(samples): - corresps[i, : s.num_corresp] = s.correspondences_AB - batch[k] = corresps - elif isinstance(samples[0].__dict__[k], torch.Tensor): - batch[k] = torch.stack([s.__dict__[k] for s in samples]) - else: - batch[k] = [s.__dict__[k] for s in samples] - return Batch(**batch) - - -class Model(nn.Module): - @property - def name(self) -> str: - return self.__class__.__name__ - - -Quality = Literal["high", "low"] -GTSource = Literal["depth", "flow", "sparse"] -DescriptorName = Literal["dinov2_vitl14"] -Normalizer = Literal["imagenet", "inception"] -FineFeaturesType = Literal["vgg19bn"] -MAX_NUM_SPARSE_CORRESP = 64 - - -class Detector(Protocol): - @property - def topleft(self) -> float: ... - - def detect_from_path( - self, - im_path: str | Path, - *, - num_keypoints: int, - return_dense_probs: bool = False, - ) -> dict[str, torch.Tensor]: ... - - def detect( - self, batch: Batch, *, num_keypoints: int, return_dense_probs: bool = False - ) -> dict[str, torch.Tensor]: ... - - def to_pixel_coords( - self, normalized_coords: torch.Tensor, h: int, w: int - ) -> torch.Tensor: ... \ No newline at end of file diff --git a/src/hardmatch/benchmark.py b/src/hardmatch/benchmark.py index 0f70029..3c0c3e4 100644 --- a/src/hardmatch/benchmark.py +++ b/src/hardmatch/benchmark.py @@ -1,33 +1,38 @@ from dataclasses import dataclass from pathlib import Path -import cv2 import torch -from typing import Literal, cast, Tuple +from typing import Literal, cast from tqdm import tqdm import numpy as np -import kornia.geometry as KG from PIL import Image from .types import BenchmarkMatcher, Batch from .dataset import HardMatch from .geometry import to_pixel -from abc import ABC, abstractmethod -class BenchmarkMatcher(ABC): - @abstractmethod - def match(self, img_A_path: str, img_B_path: str) -> Tuple[torch.Tensor, torch.Tensor]: - pass +def _symmetrical_epipolar_distance(pts1: np.ndarray, pts2: np.ndarray, Fm: np.ndarray, eps: float = 1e-8) -> np.ndarray: + """Symmetric epipolar distance, Hartley & Zisserman eq. 11.10. pts1/pts2: (B, N, 2), Fm: (B, 3, 3).""" + ones = np.ones((*pts1.shape[:-1], 1), dtype=pts1.dtype) + pts1_h = np.concatenate([pts1, ones], axis=-1) + pts2_h = np.concatenate([pts2, ones], axis=-1) -def fraction_of_gt_corrs_consisent_with_F(Fm: np.array, gt_corrs: np.array, ths=np.arange(20)): + line1_in_2 = pts1_h @ Fm.transpose(0, 2, 1) # (B, N, 3) + line2_in_1 = pts2_h @ Fm # (B, N, 3) + + numerator = (pts2_h * line1_in_2).sum(-1) ** 2 + denominator_inv = 1.0 / (np.linalg.norm(line1_in_2[..., :2], axis=-1) ** 2) + 1.0 / ( + np.linalg.norm(line2_in_1[..., :2], axis=-1) ** 2 + ) + return np.sqrt(numerator * denominator_inv + eps) + +def fraction_of_gt_corrs_consisent_with_F(Fm: np.ndarray, gt_corrs: np.ndarray, ths=np.arange(20)): assert Fm.shape == (3, 3) assert gt_corrs.shape[1] == 4 - pts_torch = torch.from_numpy(gt_corrs)[None] - F_torch = torch.from_numpy(Fm)[None] - errors = KG.symmetrical_epipolar_distance(pts_torch[:,:,:2], - pts_torch[:,:,2:], - F_torch, squared=False).detach().cpu().numpy().reshape(-1, 1) + errors = _symmetrical_epipolar_distance( + gt_corrs[None, :, :2], gt_corrs[None, :, 2:], Fm[None] + ).reshape(-1, 1) out = [] for th in ths: correct = errors <= th @@ -69,10 +74,15 @@ def benchmark( self, model: BenchmarkMatcher, ): + import cv2 + Fs = [] gt_corrs = [] pair_categories: list[list[str]] = [] - for batch in tqdm(self.dataset, desc="HardMatch Benchmark", total=len(self.dataset)): + pair_results: list[np.ndarray] = [] + ths = np.arange(20) + pbar = tqdm(self.dataset, desc="HardMatch Benchmark", total=len(self.dataset)) + for batch in pbar: batch = cast(Batch, batch) if batch.correspondences_AB.shape[0] < 8: @@ -120,6 +130,11 @@ def benchmark( gt_corrs.append(corrs_pixel) pair_categories.append(categories) + pair_results.append(fraction_of_gt_corrs_consisent_with_F(F, corrs_pixel, ths)) + running_avg_pck = np.stack(pair_results, axis=1).mean(axis=1) + running_mAA_10px = running_avg_pck[:11].mean() + pbar.set_postfix(mAA_10px=f"{running_mAA_10px:.4f}") + avg_pck, thresholds = evaluate_Fs(Fs, gt_corrs) mAA_10px = avg_pck[:11].mean() mAA_5px = avg_pck[:6].mean() diff --git a/src/hardmatch/dataset.py b/src/hardmatch/dataset.py index 31e2d5b..ad4725d 100644 --- a/src/hardmatch/dataset.py +++ b/src/hardmatch/dataset.py @@ -46,6 +46,7 @@ def __init__( if not root.exists(): if cfg.download: download_url = "https://github.com/davnords/storage/releases/download/hardmatch/hardmatch.zip" + root.parent.mkdir(parents=True, exist_ok=True) zip_path = root.parent / "hardmatch.zip" print(f"Downloading HardMatch dataset from {download_url} ...") urllib.request.urlretrieve(download_url, zip_path) @@ -94,7 +95,7 @@ def __init__( raise RuntimeError(f"{type(self).__name__}: found 0 {cfg.split} pairs under {root}") def __len__(self) -> int: - return self.cfg.weight if self.cfg.weight != -1 else len(self.pairs) + return len(self.pairs) def __getitem__(self, idx: int) -> Batch: return self.load_item(idx) diff --git a/src/hardmatch/geometry.py b/src/hardmatch/geometry.py index 419bce8..56d0afc 100644 --- a/src/hardmatch/geometry.py +++ b/src/hardmatch/geometry.py @@ -1,6 +1,5 @@ from __future__ import annotations from typing import Literal -import cv2 import torch import torch.nn.functional as F from einops import einsum @@ -696,6 +695,8 @@ def estimate_pose_fundamental( "max_epipolar_error": th, }, ) + import cv2 + E: np.ndarray = K_B.T @ F @ K_A kps_calib_A = from_homogeneous( calibrate(torch.from_numpy(kps_A).float(), torch.from_numpy(K_A).float()) diff --git a/src/hardmatch/types.py b/src/hardmatch/types.py index ca3b2ee..1606509 100644 --- a/src/hardmatch/types.py +++ b/src/hardmatch/types.py @@ -1,12 +1,9 @@ from __future__ import annotations -import torch.nn as nn from dataclasses import dataclass, field from pathlib import Path import torch from typing import Literal, Protocol import numpy as np -from PIL import Image -import json @dataclass @@ -93,250 +90,37 @@ def collate(cls, samples: list[Batch]) -> Batch: return Batch(**batch) -def _load_cfg(cls, data: dict, *, strict: bool): - annotations = getattr(cls, "__annotations__", {}) - kwargs = {} - for k, v in data.items(): - if k not in annotations: - if strict: - raise ValueError(f"Attribute {k} not found in {cls.__name__}") - else: - print(f"Attribute {k} not found in {cls.__name__}") - continue - annot_type_or_type_name = cls.__dataclass_fields__[k].type - if isinstance(annot_type_or_type_name, str): - annot_type = eval(annot_type_or_type_name) - else: - annot_type = annot_type_or_type_name - if hasattr(annot_type, "__annotations__"): - kwargs[k] = _load_cfg(annot_type, v, strict=strict) - else: - kwargs[k] = v - - return cls(**kwargs) - - -class Model(nn.Module): - @property - def name(self) -> str: - return self.__class__.__name__ - - @classmethod - def from_run( - cls, run_path: str | Path, strict: bool = True, load_avg_weights: bool = True - ) -> Model: - from sigma.device import device - - run_path = Path(run_path) - with open(run_path / "cfg.json", "r") as f: - serialized_cfg = json.load(f)["model"] - - if load_avg_weights: - if not (run_path / "avg_weights.pth").exists(): - raise FileNotFoundError( - f"{load_avg_weights=}, but average weights file not found at {run_path / 'avg_weights.pth'}" - ) - weights = torch.load(run_path / "avg_weights.pth", map_location=device) - weights = { - n.replace("module.", ""): w - for n, w in weights.items() - if n.startswith("module.") - } - else: - weights = torch.load(run_path / "weights.pth", map_location=device) - return cls.from_serialized_cfg_and_weights( - serialized_cfg, weights, strict=strict - ) - - @classmethod - def from_serialized_cfg_and_weights( - cls, serialized_cfg: dict, weights: dict[str, torch.Tensor], strict: bool = True - ) -> Model: - model = cls.from_serialized_cfg(serialized_cfg, strict=strict) - missing, unexpected = model.load_state_dict(weights, strict=strict) - if not strict: - for k in missing: - print(f"Missing key {k}") - for k in unexpected: - print(f"Unexpected key {k}") - return model - - @classmethod - def from_serialized_cfg(cls, serialized_cfg: dict, strict: bool = True) -> Model: - cfg = _load_cfg(cls.Cfg, serialized_cfg, strict=strict) - return cls.from_cfg(cfg) - - @classmethod - def from_cfg(cls, cfg) -> Model: - from .device import device - - return cls(cfg).to(device) - - Quality = Literal["high", "low"] GTSource = Literal["depth", "flow", "sparse"] -HeadType = Literal[ - "dpt-no-pos", - "dptv2-baseline", - "dptv2-flux", - "dptv2-roma", - "dptv2-bilinear", - "dptv2-bilinear-no-align-corners", - "dptv2-no-align-corners", - "dptv2-flux-bilinear", -] -SampleMode = Literal["frame_distance", "overlap"] -ConfidenceMode = Literal["covis", "frame", "positive"] -NormType = Literal["batch"] -RefinersType = Literal["roma-4-pow2"] -MatcherStyle = Literal["romav3"] -DescriptorName = Literal["dinov3_vitl16", "dinov2_vitl14", "mum_vitl16", "flux2"] -Normalizer = Literal["imagenet", "inception"] # Callable[[torch.Tensor], torch.Tensor] -OptimizerName = Literal["adamw"] -ImageLike = torch.Tensor | np.ndarray | str | Path | Image.Image -Setting = Literal[ - "mega1500", "scannet1500", "wxbs", "satast", "base", "precise", "turbo", "fast" -] -FineFeaturesType = Literal["vgg19", "vgg19bn", "flux2"] MAX_NUM_SPARSE_CORRESP = 64 -class Matcher(Protocol): - """Protocol defining the interface for feature matchers.""" +class BenchmarkMatcher(Protocol): + """Interface every matcher passed to `HardMatchBenchmark.benchmark()` must satisfy. - def match( - self, im_A_path: str | Path, im_B_path: str | Path - ) -> tuple[torch.Tensor, torch.Tensor]: - """ - Match features between two images. - """ - ... - - def match_keypoints( - self, - keypoints_A: torch.Tensor, - keypoints_B: torch.Tensor, - warp: torch.Tensor, - certainty: torch.Tensor, - return_tuple: bool = False, - ) -> torch.Tensor: - """ - Match keypoints between two images. - """ - ... - - def to_pixel_coordinates( - self, matches: torch.Tensor, h1: int, w1: int, h2: int, w2: int - ) -> tuple[torch.Tensor, torch.Tensor]: - """ - Convert matches to pixel coordinates. - """ + This is a `typing.Protocol`, so a matcher does not need to inherit from + it -- any object with a matching `match()` method and `offset` attribute + works (structural typing). + Example: + class YourMatcher: + offset = 0.0 -class Detector(Protocol): - """Protocol defining the interface for keypoint detectors.""" + def match(self, img_A_path: str, img_B_path: str) -> tuple[np.ndarray, np.ndarray]: + ... + return kpts_A, kpts_B + """ - @property - def topleft(self) -> float: - """Top-left coordinate offset for keypoint normalization.""" - ... + offset: float + """Pixel offset subtracted from returned keypoints before scoring (0.0 for most matchers).""" - def detect_from_path( - self, - im_path: str | Path, - *, - num_keypoints: int, - return_dense_probs: bool = False, - ) -> dict[str, torch.Tensor]: - """ - Detect keypoints from an image file path. - - Args: - im_path: Path to the image file - num_keypoints: Number of keypoints to detect - return_dense_probs: Whether to return dense probability maps - - Returns: - Dictionary containing: - - "keypoints": torch.Tensor of shape (B, N, 2) with normalized coordinates [-1, 1] - - "keypoint_probs": torch.Tensor of shape (B, N) with confidence scores - - "dense_probs": torch.Tensor of shape (B, H, W), optional if return_dense_probs=True - """ - ... - def detect(self, batch: Batch, *, num_keypoints: int, return_dense_probs: bool = False) -> dict[str, torch.Tensor]: - """ - Detect keypoints from a batch of images. - """ - ... - - def to_pixel_coords( - self, normalized_coords: torch.Tensor, h: int, w: int - ) -> torch.Tensor: - """ - Convert normalized coordinates to pixel coordinates. - - Args: - normalized_coords: Tensor of shape (..., 2) with coordinates in [-1, 1] - h: Image height in pixels - w: Image width in pixels - - Returns: - Tensor of shape (..., 2) with pixel coordinates - """ - ... - -class SparseMatcher(Protocol): - """Protocol defining the interface for sparse matchers, e.g. LightGlue, SuperGlue and DualSoftMaxMatcher.""" - def match( - self, keypoints_A: torch.Tensor, description_A: torch.Tensor, keypoints_B: torch.Tensor, description_B: torch.Tensor - ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: - """ - Match features between two images. - """ - ... - - def to_pixel_coords( - self, normalized_coords: torch.Tensor, h: int, w: int - ) -> torch.Tensor: - """ - Convert normalized coordinates to pixel coordinates. - - Args: - normalized_coords: Tensor of shape (..., 2) with coordinates in [-1, 1] - h: Image height in pixels - w: Image width in pixels + self, img_A_path: str | Path, img_B_path: str | Path + ) -> tuple[np.ndarray, np.ndarray]: + """Match two images given by file path. Returns: - Tensor of shape (..., 2) with pixel coordinates + (kpts_A, kpts_B): corresponding keypoints as (N, 2) arrays in + pixel coordinates of the original image files, one row per match. """ - ... - -class Descriptor(Protocol): - """Protocol defining the interface for feature descriptors, e.g. DeDoDe and SIFT""" - - def describe_keypoints_from_path(self, im_path: str, keypoints: torch.Tensor, H:int=784, W:int=784) -> dict[str, torch.Tensor]: - ... - - -from abc import ABC, abstractmethod -from typing import Tuple -class BenchmarkMatcher(ABC): - @property - @abstractmethod - def name(self) -> str: - pass - - @abstractmethod - def match(self, img_A_path: str, img_B_path: str) -> Tuple[torch.Tensor, torch.Tensor]: - pass - - @property - @abstractmethod - def get_type(self) -> Literal["dense", "sparse"]: - pass - - @property - @abstractmethod - def offset(self) -> float: - pass \ No newline at end of file + ... \ No newline at end of file diff --git a/tests/test_smoke.py b/tests/test_smoke.py new file mode 100644 index 0000000..b5b291e --- /dev/null +++ b/tests/test_smoke.py @@ -0,0 +1,51 @@ +import numpy as np +import torch + +from hardmatch import HardMatchBenchmark +from hardmatch.benchmark import fraction_of_gt_corrs_consisent_with_F +from hardmatch.geometry import to_normalized, to_pixel +from hardmatch.types import Batch + + +def test_public_api_imports() -> None: + assert HardMatchBenchmark is not None + + +def test_pixel_normalization_roundtrip() -> None: + pts = torch.tensor([[10.0, 20.0], [100.0, 50.0]]) + normalized = to_normalized(pts, H=200, W=300) + pixel = to_pixel(normalized, H=200, W=300) + assert torch.allclose(pixel, pts, atol=1e-3) + + +def test_fundamental_matrix_scoring() -> None: + # A perfect correspondence should be epipolar-consistent at every (non-zero) threshold. + F = np.array([[0.0, 0.0, 0.0], [0.0, 0.0, -1.0], [0.0, 1.0, 0.0]]) + gt_corrs = np.array([[10.0, 20.0, 10.0, 20.0]]) + out = fraction_of_gt_corrs_consisent_with_F(F, gt_corrs, ths=np.arange(1, 5)) + assert all(v == 1.0 for v in out) + + +def test_batch_construction() -> None: + batch = Batch( + img_A=torch.zeros(3, 4, 4), + img_B=torch.zeros(3, 4, 4), + depth_A=torch.zeros(4, 4, 1), + depth_B=torch.zeros(4, 4, 1), + img_A_path="a.jpg", + img_B_path="b.jpg", + source="sparse", + flow_AB=torch.zeros(4, 4, 2), + flow_BA=torch.zeros(4, 4, 2), + mask_AB=torch.ones(4, 4, 1, dtype=torch.bool), + mask_BA=torch.ones(4, 4, 1, dtype=torch.bool), + quality="high", + ) + assert batch.img_A.shape == (3, 4, 4) + + +if __name__ == "__main__": + test_public_api_imports() + test_pixel_normalization_roundtrip() + test_fundamental_matrix_scoring() + test_batch_construction() diff --git a/utils.py b/utils.py deleted file mode 100644 index e69de29..0000000