update

.gitattributes

@@ -33,3 +33,5 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+weights/ filter=lfs diff=lfs merge=lfs -text
+assets/ filter=lfs diff=lfs merge=lfs -text

app.py

@@ -1,8 +1,140 @@
+import cv2
+import numpy as np
 import gradio as gr
 
-def greet(name):
-    return "Hello " + name + "!!"
+import torch
+import matplotlib.cm as cm
+import sys
 
-demo = gr.Interface(fn=greet, inputs="text", outputs="text")
-demo.launch()
+sys.path.append("src")
 
+from src.utils.plotting import make_matching_figure
+from src.edm import EDM
+from src.config.default import get_cfg_defaults
+from src.utils.misc import lower_config
+
+    
+HEADER = """
+<div align="center">
+    <p>
+        <span style="font-size: 30px; vertical-align: bottom;"> 🎶 EDM: Efficient Deep Feature Matching</span>
+    </p>
+    <p style="margin-top: -15px;">
+        <a href="https://arxiv.org/pdf/2503.05122" target="_blank" style="color: grey;">ArXiv Paper</a>
+        &nbsp;
+        <a href="https://github.com/chicleee/EDM" target="_blank" style="color: grey;">GitHub Repository</a>
+    </p>
+
+</div>
+"""
+
+ABSTRACT = """
+Recent feature matching methods have achieved remarkable performance but lack efficiency consideration. In this paper, we revisit the mainstream detector-free matching pipeline and improve all its stages considering both accuracy and efficiency. We propose an Efficient Deep feature Matching network, EDM. We first adopt a deeper CNN with fewer dimensions to extract multi-level features. Then we present a Correlation Injection Module that conducts feature transformation on high-level deep features, and progressively injects feature correlations from global to local for efficient multi-scale feature aggregation, improving both speed and performance. In the refinement stage, a novel lightweight bidirectional axis-based regression head is designed to directly predict subpixel-level correspondences from latent features, avoiding the significant computational cost of explicitly locating keypoints on high-resolution local feature heatmaps. Moreover, effective selection strategies are introduced to enhance matching accuracy. Extensive experiments show that our EDM achieves competitive matching accuracy on various benchmarks and exhibits excellent efficiency, offering valuable best practices for real-world applications."""
+
+def find_matches(image_0, image_1, conf_thres=0.2, border_rm=2, topk=10000):
+    config = get_cfg_defaults()
+    data_cfg_path = "configs/data/megadepth_test_1500.py"
+    main_cfg_path = "configs/edm/outdoor/edm_base.py"
+    config.merge_from_file(main_cfg_path)
+    config.merge_from_file(data_cfg_path)
+
+    W, H = 832, 832
+    config.EDM.COARSE.MCONF_THR = conf_thres
+    config.EDM.COARSE.BORDER_RM = border_rm
+    config.EDM.COARSE.TOPK = topk
+
+    _config = lower_config(config)
+    matcher = EDM(config=_config["edm"])
+    state_dict = torch.load("weights/edm_outdoor.ckpt")["state_dict"]
+    matcher.load_state_dict(state_dict)
+    matcher = matcher.eval()
+    
+    # Load example images
+    img0_bgr = image_0
+    img1_bgr = image_1
+
+
+    h0, w0 = img0_bgr.shape[:2]
+    h1, w1 = img1_bgr.shape[:2]
+
+    h0_scale = h0 / H
+    w0_scale = w0 / W
+    h1_scale = h1 / H
+    w1_scale = w1 / W
+
+    # For inference
+    img0_raw = cv2.cvtColor(img0_bgr, cv2.COLOR_BGR2GRAY)
+    img1_raw = cv2.cvtColor(img1_bgr, cv2.COLOR_BGR2GRAY)
+    img0_raw = cv2.resize(img0_raw, (W, H))
+    img1_raw = cv2.resize(img1_raw, (W, H))
+
+    img0 = torch.from_numpy(img0_raw)[None][None] / 255.
+    img1 = torch.from_numpy(img1_raw)[None][None]/ 255.
+    batch = {'image0': img0, 'image1': img1}
+
+    # Inference with EDM and get prediction
+    with torch.no_grad():
+        matcher(batch)
+
+    mkpts0 = batch['mkpts0_f'].cpu().numpy()
+    mkpts1 = batch['mkpts1_f'].cpu().numpy()
+    mconf = batch['mconf'].cpu().numpy()
+    
+    mkpts0[:, 0] *= w0_scale
+    mkpts0[:, 1] *= h0_scale
+    mkpts1[:, 0] *= w1_scale
+    mkpts1[:, 1] *= h1_scale
+    
+    color = cm.jet(mconf)
+    # Draw
+    text = [
+        'EDM',
+        'Matches: {}'.format(len(mkpts0)),
+    ]
+    fig = make_matching_figure(img0_bgr, img1_bgr, mkpts0, mkpts1, color, text=text)
+    
+    return fig
+
+
+with gr.Blocks() as demo:
+
+    gr.Markdown(HEADER)
+    with gr.Accordion("Abstract (click to open)", open=False):
+        gr.Image("assets/teaser.jpg")
+        gr.Markdown(ABSTRACT)
+
+    with gr.Row():
+        image_1 = gr.Image()
+        image_2 = gr.Image()
+    with gr.Row():
+        conf_thres = gr.Slider(minimum=0.01, maximum=1, value=0.2, step=0.01, label="Coarse Confidence Threshold")
+        topk = gr.Slider(minimum=100, maximum=10000, value=5000, step=100, label="TopK")
+        border_rm = gr.Slider(minimum=0, maximum=20, value=2, step=1, label="Border Remove (x8)")
+    gr.HTML(
+        """
+        Note: images are actually resized to 832 x 832 for matching.
+        """
+    )
+    with gr.Row():
+        button = gr.Button(value="Find Matches")
+        clear = gr.ClearButton(value="Clear")
+        
+    output = gr.Image()
+    button.click(find_matches, [image_1, image_2, conf_thres, border_rm, topk], output)
+    clear.add([image_1, image_2, output])
+
+    gr.Examples(
+        examples=[
+            ["assets/scannet_sample_images/scene0707_00_15.jpg", "assets/scannet_sample_images/scene0707_00_45.jpg"],
+            ["assets/scannet_sample_images/scene0758_00_165.jpg", "assets/scannet_sample_images/scene0758_00_510.jpg"],
+            ["assets/phototourism_sample_images/london_bridge_19481797_2295892421.jpg", "assets/phototourism_sample_images/london_bridge_49190386_5209386933.jpg"],
+            ["assets/phototourism_sample_images/united_states_capitol_26757027_6717084061.jpg", "assets/phototourism_sample_images/united_states_capitol_98169888_3347710852.jpg"],
+        ],
+        inputs=[image_1, image_2],
+        outputs=[output],
+        fn=find_matches,
+        cache_examples=None,
+    )
+
+if __name__ == "__main__":
+    demo.launch(debug=True)

configs/data/base.py

@@ -0,0 +1,37 @@
+"""
+The data config will be the last one merged into the main config.
+Setups in data configs will override all existed setups!
+"""
+
+from yacs.config import CfgNode as CN
+_CN = CN()
+_CN.DATASET = CN()
+_CN.TRAINER = CN()
+_CN.EDM = CN()
+_CN.EDM.NECK = CN()
+
+# training data config
+_CN.DATASET.TRAIN_DATA_ROOT = None
+_CN.DATASET.TRAIN_POSE_ROOT = None
+_CN.DATASET.TRAIN_NPZ_ROOT = None
+_CN.DATASET.TRAIN_LIST_PATH = None
+_CN.DATASET.TRAIN_INTRINSIC_PATH = None
+# validation set config
+_CN.DATASET.VAL_DATA_ROOT = None
+_CN.DATASET.VAL_POSE_ROOT = None
+_CN.DATASET.VAL_NPZ_ROOT = None
+_CN.DATASET.VAL_LIST_PATH = None
+_CN.DATASET.VAL_INTRINSIC_PATH = None
+
+# testing data config
+_CN.DATASET.TEST_DATA_ROOT = None
+_CN.DATASET.TEST_POSE_ROOT = None
+_CN.DATASET.TEST_NPZ_ROOT = None
+_CN.DATASET.TEST_LIST_PATH = None
+_CN.DATASET.TEST_INTRINSIC_PATH = None
+
+# dataset config
+_CN.DATASET.MIN_OVERLAP_SCORE_TRAIN = 0.4
+_CN.DATASET.MIN_OVERLAP_SCORE_TEST = 0.0  # for both test and val
+
+cfg = _CN

configs/data/megadepth_test_1500.py

@@ -0,0 +1,23 @@
+from configs.data.base import cfg
+
+TEST_BASE_PATH = "assets/megadepth_test_1500_scene_info"
+
+cfg.DATASET.TEST_DATA_SOURCE = "MegaDepth"
+cfg.DATASET.TEST_DATA_ROOT = "data/megadepth/test"
+cfg.DATASET.TEST_NPZ_ROOT = f"{TEST_BASE_PATH}"
+cfg.DATASET.TEST_LIST_PATH = f"{TEST_BASE_PATH}/megadepth_test_1500.txt"
+
+cfg.DATASET.MIN_OVERLAP_SCORE_TEST = 0.0
+
+cfg.EDM.TRAIN_RES_H = 832
+cfg.EDM.TRAIN_RES_W = 832
+cfg.EDM.TEST_RES_H = 832
+cfg.EDM.TEST_RES_W = 832
+
+cfg.EDM.NECK.NPE = [
+    cfg.EDM.TRAIN_RES_H,
+    cfg.EDM.TRAIN_RES_W,
+    cfg.EDM.TEST_RES_H,
+    cfg.EDM.TEST_RES_W,
+]
+

configs/edm/outdoor/edm_base.py

@@ -0,0 +1,17 @@
+from src.config.default import _CN as cfg
+
+cfg.TRAINER.CANONICAL_BS = 4 * 8
+cfg.TRAINER.CANONICAL_LR = 2e-3
+cfg.TRAINER.WARMUP_STEP = int(36800 / cfg.TRAINER.CANONICAL_BS * 3)  # 3 epochs
+cfg.TRAINER.WARMUP_RATIO = 0.1
+cfg.TRAINER.MSLR_MILESTONES = [8, 12, 16, 20, 24]
+cfg.TRAINER.EPI_ERR_THR = 1e-4
+
+cfg.EDM.COARSE.MCONF_THR = 0.05
+cfg.EDM.FINE.SIGMA_THR = 1e-6
+cfg.EDM.COARSE.BORDER_RM = 0
+
+# Top-K should not exceed grid_size = TEST_RES_H / 8 * TEST_RES_W / 8
+# The recommended value is approximately grid_size * 0.35 for Megadepth
+# cfg.EDM.COARSE.TOPK = int(832 / 8 * 832 / 8 * 0.35)  # 3786 for train & LO-RANSAC test
+cfg.EDM.COARSE.TOPK = int(1152 / 8 * 1152 / 8 * 0.35)  # 7258 for test

requirements.txt

@@ -0,0 +1,22 @@
+pytorch==2.0.0
+torchvision==0.15.0
+torchaudio==2.0.0
+opencv-python
+albumentations==0.5.1 --no-binary=imgaug,albumentations
+ray>=1.0.1
+einops==0.3.0
+kornia==0.4.1
+loguru==0.5.3
+yacs>=0.1.8
+tqdm
+autopep8
+pylint
+ipython
+jupyterlab
+matplotlib
+h5py==3.1.0
+pytorch-lightning==1.3.5
+torchmetrics==0.7.0
+joblib>=1.0.1
+pillow==9.5.0
+poselib

src/config/default.py

@@ -0,0 +1,184 @@
+from yacs.config import CfgNode as CN
+
+_CN = CN()
+
+##############  ↓  EDM Pipeline  ↓  ##############
+_CN.EDM = CN()
+_CN.EDM.TRAIN_RES_H = 832
+_CN.EDM.TRAIN_RES_W = 832
+_CN.EDM.TEST_RES_H = 832
+_CN.EDM.TEST_RES_W = 832
+_CN.EDM.LOCAL_RESOLUTION = 8  # coarse matching at 1/8, window_size 8x8 in fine_level
+_CN.EDM.MP = False  # Not use mixed precision (ELoFTR defualt used mixed precision)
+_CN.EDM.HALF = False  # Not FP16
+_CN.EDM.DEPLOY = False # export onnx model
+_CN.EDM.EVAL_TIMES = 1
+
+# 3. Coarse-Matching config
+_CN.EDM.BACKBONE = CN()
+_CN.EDM.BACKBONE.BLOCK_DIMS = [32, 64, 128, 256, 256]  # 1/2 -> 1/32
+
+# 3. Coarse-Matching config
+_CN.EDM.NECK = CN()
+_CN.EDM.NECK.D_MODEL = 256
+_CN.EDM.NECK.NHEAD = 8
+_CN.EDM.NECK.LAYER_NAMES = ["self", "cross"] * 2
+_CN.EDM.NECK.AGG_SIZE0 = 1
+_CN.EDM.NECK.AGG_SIZE1 = 1
+_CN.EDM.NECK.ROPE = True
+_CN.EDM.NECK.NPE = None
+
+# 3. Coarse-Matching config
+_CN.EDM.COARSE = CN()
+_CN.EDM.COARSE.MCONF_THR = 0.2
+_CN.EDM.COARSE.BORDER_RM = 0
+_CN.EDM.COARSE.DSMAX_TEMPERATURE = 0.1
+_CN.EDM.COARSE.TRAIN_PAD_NUM = 32  # training tricks: avoid DDP deadlock
+_CN.EDM.COARSE.TOPK = 2048
+_CN.EDM.COARSE.DS_OPT = True
+
+# 4. EDM-fine module config
+_CN.EDM.FINE = CN()
+_CN.EDM.FINE.DROPRATE = None
+_CN.EDM.FINE.COORD_LENGTH = 16
+_CN.EDM.FINE.BI_DIRECTIONAL_REFINE = True
+_CN.EDM.FINE.SIGMA_THR = 0.0
+_CN.EDM.FINE.SIGMA_SELECTION = True
+
+# 5. EDM Losses
+# -- # coarse-level
+_CN.EDM.LOSS = CN()
+_CN.EDM.LOSS.COARSE_TYPE = "focal"
+_CN.EDM.LOSS.COARSE_WEIGHT = 1.0
+_CN.EDM.LOSS.SPARSE_SPVS = True
+
+# -- - -- # focal loss (coarse)
+_CN.EDM.LOSS.FOCAL_ALPHA = 0.25
+_CN.EDM.LOSS.FOCAL_GAMMA = 2.0
+_CN.EDM.LOSS.POS_WEIGHT = 1.0
+_CN.EDM.LOSS.NEG_WEIGHT = 1.0
+
+
+# -- # fine-level
+_CN.EDM.LOSS.FINE_TYPE = "rle"
+_CN.EDM.LOSS.FINE_WEIGHT = 0.2
+_CN.EDM.LOSS.Q_DISTRIBUTION = "laplace"  # options: ['laplace', 'gaussian']
+
+
+##############  Dataset  ##############
+_CN.DATASET = CN()
+# 1. data config
+# training and validating
+_CN.DATASET.TRAINVAL_DATA_SOURCE = None  # options: ['ScanNet', 'MegaDepth']
+_CN.DATASET.TRAIN_DATA_ROOT = None
+_CN.DATASET.TRAIN_POSE_ROOT = None  # (optional directory for poses)
+_CN.DATASET.TRAIN_NPZ_ROOT = None
+_CN.DATASET.TRAIN_LIST_PATH = None
+_CN.DATASET.TRAIN_INTRINSIC_PATH = None
+_CN.DATASET.VAL_DATA_ROOT = None
+_CN.DATASET.VAL_POSE_ROOT = None  # (optional directory for poses)
+_CN.DATASET.VAL_NPZ_ROOT = None
+_CN.DATASET.VAL_LIST_PATH = (
+    None  # None if val data from all scenes are bundled into a single npz file
+)
+_CN.DATASET.VAL_INTRINSIC_PATH = None
+# testing
+_CN.DATASET.TEST_DATA_SOURCE = None
+_CN.DATASET.TEST_DATA_ROOT = None
+_CN.DATASET.TEST_POSE_ROOT = None  # (optional directory for poses)
+_CN.DATASET.TEST_NPZ_ROOT = None
+_CN.DATASET.TEST_LIST_PATH = (
+    None  # None if test data from all scenes are bundled into a single npz file
+)
+_CN.DATASET.TEST_INTRINSIC_PATH = None
+
+# 2. dataset config
+# general options
+_CN.DATASET.MIN_OVERLAP_SCORE_TRAIN = (
+    0.4  # discard data with overlap_score < min_overlap_score
+)
+_CN.DATASET.MIN_OVERLAP_SCORE_TEST = 0.0
+_CN.DATASET.AUGMENTATION_TYPE = None  # options: [None, 'dark', 'mobile']
+
+# MegaDepth options
+_CN.DATASET.MGDPT_IMG_PAD = True  # pad img to square with size = MGDPT_IMG_RESIZE
+_CN.DATASET.MGDPT_DEPTH_PAD = True  # pad depthmap to square with size = 2000
+_CN.DATASET.MGDPT_DF = 8
+
+
+##############  Trainer  ##############
+_CN.TRAINER = CN()
+_CN.TRAINER.WORLD_SIZE = 1
+_CN.TRAINER.CANONICAL_BS = 4 * 8
+_CN.TRAINER.CANONICAL_LR = 2e-3
+_CN.TRAINER.SCALING = None  # this will be calculated automatically
+_CN.TRAINER.FIND_LR = False  # use learning rate finder from pytorch-lightning
+
+# optimizer
+_CN.TRAINER.OPTIMIZER = "adamw"  # [adam, adamw]
+_CN.TRAINER.TRUE_LR = None  # this will be calculated automatically at runtime
+_CN.TRAINER.ADAM_DECAY = 0.0  # ADAM: for adam
+_CN.TRAINER.ADAMW_DECAY = 0.1
+
+# step-based warm-up
+_CN.TRAINER.WARMUP_TYPE = "linear"  # [linear, constant]
+_CN.TRAINER.WARMUP_RATIO = 0.0
+_CN.TRAINER.WARMUP_STEP = 4800
+
+# learning rate scheduler
+# [MultiStepLR, CosineAnnealing, ExponentialLR]
+_CN.TRAINER.SCHEDULER = "MultiStepLR"
+_CN.TRAINER.SCHEDULER_INTERVAL = "epoch"  # [epoch, step]
+_CN.TRAINER.MSLR_MILESTONES = [3, 6, 9, 12]  # MSLR: MultiStepLR
+_CN.TRAINER.MSLR_GAMMA = 0.5
+_CN.TRAINER.COSA_TMAX = 30  # COSA: CosineAnnealing
+# ELR: ExponentialLR, this value for 'step' interval
+_CN.TRAINER.ELR_GAMMA = 0.999992
+
+# plotting related
+_CN.TRAINER.ENABLE_PLOTTING = False
+_CN.TRAINER.N_VAL_PAIRS_TO_PLOT = 32  # number of val/test paris for plotting
+_CN.TRAINER.PLOT_MODE = "evaluation"  # ['evaluation', 'confidence']
+_CN.TRAINER.PLOT_MATCHES_ALPHA = "dynamic"
+
+# geometric metrics and pose solver
+# recommendation: 5e-4 for ScanNet, 1e-4 for MegaDepth (from SuperGlue)
+_CN.TRAINER.EPI_ERR_THR = 5e-4
+_CN.TRAINER.POSE_GEO_MODEL = "E"  # ['E', 'F', 'H']
+_CN.TRAINER.POSE_ESTIMATION_METHOD = "RANSAC"  # [RANSAC, LO-RANSAC]
+_CN.TRAINER.RANSAC_PIXEL_THR = 0.5
+_CN.TRAINER.RANSAC_CONF = 0.99999
+
+# data sampler for train_dataloader
+_CN.TRAINER.DATA_SAMPLER = (
+    "scene_balance"  # options: ['scene_balance', 'random', 'normal']
+)
+# 'scene_balance' config
+_CN.TRAINER.N_SAMPLES_PER_SUBSET = 200
+# whether sample each scene with replacement or not
+_CN.TRAINER.SB_SUBSET_SAMPLE_REPLACEMENT = True
+# after sampling from scenes, whether shuffle within the epoch or not
+_CN.TRAINER.SB_SUBSET_SHUFFLE = True
+_CN.TRAINER.SB_REPEAT = 1  # repeat N times for training the sampled data
+# 'random' config
+_CN.TRAINER.RDM_REPLACEMENT = True
+_CN.TRAINER.RDM_NUM_SAMPLES = None
+
+# gradient clipping
+_CN.TRAINER.GRADIENT_CLIPPING = 0.5
+
+# reproducibility
+# This seed affects the data sampling. With the same seed, the data sampling is promised
+# to be the same. When resume training from a checkpoint, it's better to use a different
+# seed, otherwise the sampled data will be exactly the same as before resuming, which will
+# cause less unique data items sampled during the entire training.
+# Use of different seed values might affect the final training result, since not all data items
+# are used during training on ScanNet. (60M pairs of images sampled during traing from 230M pairs in total.)
+_CN.TRAINER.SEED = 66
+
+
+def get_cfg_defaults():
+    """Get a yacs CfgNode object with default values for my_project."""
+    # Return a clone so that the defaults will not be altered
+    # This is for the "local variable" use pattern
+    return _CN.clone()

src/edm/__init__.py

@@ -0,0 +1,2 @@
+from .edm import EDM
+

src/edm/backbone/resnet.py

@@ -0,0 +1,116 @@
+import torch.nn as nn
+
+
+def conv1x1(in_planes, out_planes, stride=1, bias=False):
+    """1x1 convolution without padding"""
+    return nn.Conv2d(
+        in_planes, out_planes, kernel_size=1, stride=stride, padding=0, bias=bias
+    )
+
+
+def conv3x3(in_planes, out_planes, stride=1, groups=1, bias=False):
+    """3x3 convolution with padding"""
+    return nn.Conv2d(
+        in_planes,
+        out_planes,
+        kernel_size=3,
+        stride=stride,
+        padding=1,
+        groups=groups,
+        bias=bias,
+    )
+
+
+class BasicBlock(nn.Module):
+    def __init__(self, in_planes, planes, stride=1):
+        super().__init__()
+        self.conv1 = conv3x3(in_planes, planes, stride)
+        self.conv2 = conv3x3(planes, planes)
+        self.bn1 = nn.BatchNorm2d(planes)
+        self.bn2 = nn.BatchNorm2d(planes)
+        self.relu = nn.ReLU(inplace=True)
+
+        if stride == 1:
+            self.downsample = None
+        else:
+            self.downsample = nn.Sequential(
+                conv1x1(in_planes, planes, stride=stride), nn.BatchNorm2d(planes)
+            )
+
+    def forward(self, x):
+        y = x
+        y = self.relu(self.bn1(self.conv1(y)))
+        y = self.bn2(self.conv2(y))
+
+        if self.downsample is not None:
+            x = self.downsample(x)
+
+        return self.relu(x + y)
+
+
+class ResNet18(nn.Module):
+    """
+    Fewer channels
+    """
+
+    def __init__(self, config=None):
+        super().__init__()
+        # Config
+        block_dims = config["backbone"]["block_dims"]
+
+        # Networks
+        self.conv1 = nn.Conv2d(
+            1, block_dims[0], kernel_size=7, stride=2, padding=3, bias=False
+        )
+        self.bn1 = nn.BatchNorm2d(block_dims[0])
+        self.relu = nn.ReLU(inplace=True)
+        self.layer1 = self._make_layer(
+            BasicBlock, block_dims[0], block_dims[0], stride=1
+        )  # 1/2
+        self.layer2 = self._make_layer(
+            BasicBlock, block_dims[0], block_dims[1], stride=2
+        )  # 1/4
+        self.layer3 = self._make_layer(
+            BasicBlock, block_dims[1], block_dims[2], stride=2
+        )  # 1/8
+        self.layer4 = self._make_layer(
+            BasicBlock, block_dims[2], block_dims[3], stride=2
+        )  # 1/16
+        self.layer5 = self._make_layer(
+            BasicBlock, block_dims[3], block_dims[4], stride=2
+        )  # 1/32
+
+        # For fine matching
+        self.fine_conv = nn.Sequential(
+            self._make_layer(
+                BasicBlock, block_dims[2], block_dims[2], stride=1),
+            conv1x1(block_dims[2], block_dims[4]),
+            nn.BatchNorm2d(block_dims[4]),
+        )
+
+        for m in self.modules():
+            if isinstance(m, nn.Conv2d):
+                nn.init.kaiming_normal_(
+                    m.weight, mode="fan_out", nonlinearity="relu")
+            elif isinstance(m, (nn.BatchNorm2d, nn.GroupNorm)):
+                nn.init.constant_(m.weight, 1)
+                nn.init.constant_(m.bias, 0)
+
+    def _make_layer(self, block, in_dim, out_dim, stride=1):
+        layer1 = block(in_dim, out_dim, stride=stride)
+        layer2 = block(out_dim, out_dim, stride=1)
+        layers = (layer1, layer2)
+
+        return nn.Sequential(*layers)
+
+    def forward(self, x):
+        x0 = self.relu(self.bn1(self.conv1(x)))
+        x1 = self.layer1(x0)  # 1/2
+        x2 = self.layer2(x1)  # 1/4
+        x3 = self.layer3(x2)  # 1/8
+        x4 = self.layer4(x3)  # 1/16
+        x5 = self.layer5(x4)  # 1/32
+
+        xf = self.fine_conv(x3)  # 1/8
+
+        return [x3, x4, x5, xf]

src/edm/edm.py

@@ -0,0 +1,204 @@
+from ..utils.misc import detect_NaN
+from .head.fine_matching import FineMatching
+from .head.coarse_matching import CoarseMatching
+from .neck.neck import CIM
+from .backbone.resnet import ResNet18
+from einops.einops import rearrange
+import torch.nn.functional as F
+import torch.nn as nn
+import torch
+torch.set_float32_matmul_precision("highest")  # highest (defualt) high medium
+
+
+class EDM(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        # Misc
+        self.config = config
+        self.local_resolution = self.config["local_resolution"]
+        self.bi_directional_refine = self.config["fine"]["bi_directional_refine"]
+        self.deploy = self.config["deploy"]
+        self.topk = config["coarse"]["topk"]
+
+        # Modules
+        self.backbone = ResNet18(config)
+        self.neck = CIM(config)
+        self.coarse_matching = CoarseMatching(config)
+        self.fine_matching = FineMatching(config)
+
+    def forward(self, data):
+        """
+        Update:
+            data (dict): {
+                'image0': (torch.Tensor): (N, 1, H, W)
+                'image1': (torch.Tensor): (N, 1, H, W)
+                'mask0'(optional) : (torch.Tensor): (N, H, W) '0' indicates a padded position
+                'mask1'(optional) : (torch.Tensor): (N, H, W)
+            }
+        """
+        if self.deploy:
+            image0, image1 = data.split(1, 1)
+            data = {"image0": image0, "image1": image1}
+
+        data.update(
+            {
+                "bs": data["image0"].size(0),
+                "hw0_i": data["image0"].shape[2:],
+                "hw1_i": data["image1"].shape[2:],
+            }
+        )
+
+        # 1. Feature Extraction
+        if data["hw0_i"] == data["hw1_i"]:
+            # faster & better BN convergence
+            feats = self.backbone(
+                torch.cat([data["image0"], data["image1"]], dim=0))
+            f8, f16, f32, f8_fine = feats
+            ms_feats = f8, f16, f32
+            feat_f0, feat_f1 = f8_fine.chunk(2)
+        else:
+            # handle different input shapes
+            # raise ValueError("image0 and image1 should have the same shape.")
+            feats0, feats1 = self.backbone(data["image0"]), self.backbone(
+                data["image1"]
+            )
+            f8_0, f16_0, f32_0, feat_f0 = feats0
+            f8_1, f16_1, f32_1, feat_f1 = feats1
+            ms_feats = f8_0, f16_0, f32_0, f8_1, f16_1, f32_1
+
+        mask_c0 = mask_c1 = None  # mask is useful in training
+        if "mask0" in data:
+            mask_c0, mask_c1 = data["mask0"], data["mask1"]
+
+        # 2.  Feature Interaction & Multi-Scale Fusion
+        feat_c0, feat_c1 = self.neck(ms_feats, mask_c0, mask_c1)
+
+        data.update(
+            {
+                "hw0_c": feat_c0.shape[2:],
+                "hw1_c": feat_c1.shape[2:],
+                "hw0_f": feat_c0.shape[2:] * self.config["local_resolution"],
+                "hw1_f": feat_c1.shape[2:] * self.config["local_resolution"],
+            }
+        )
+        feat_c0 = rearrange(feat_c0, "n c h w -> n (h w) c")
+        feat_c1 = rearrange(feat_c1, "n c h w -> n (h w) c")
+        feat_f0 = rearrange(feat_f0, "n c h w -> n (h w) c")
+        feat_f1 = rearrange(feat_f1, "n c h w -> n (h w) c")
+
+        # detect NaN during mixed precision training
+        if self.config["mp"] and (
+            torch.any(torch.isnan(feat_c0)) or torch.any(torch.isnan(feat_c1))
+        ):
+            detect_NaN(feat_c0, feat_c1)
+
+        # 3. Coarse-Level Matching
+        conf_matrix = self.coarse_matching(
+            feat_c0,
+            feat_c1,
+            data,
+            mask_c0=(
+                mask_c0.view(mask_c0.size(0), -
+                             1) if mask_c0 is not None else mask_c0
+            ),
+            mask_c1=(
+                mask_c1.view(mask_c1.size(0), -
+                             1) if mask_c1 is not None else mask_c1
+            ),
+        )
+
+        if self.deploy:
+            k = self.topk
+            row_max_val, row_max_idx = torch.max(conf_matrix, dim=2)
+            topk_val, topk_idx = torch.topk(row_max_val, k, dim=1)
+
+            b_ids = (
+                torch.arange(conf_matrix.shape[0], device=conf_matrix.device)
+                .unsqueeze(1)
+                .repeat(1, k)
+                .flatten()
+            )
+            i_ids = topk_idx.flatten()
+            j_ids = row_max_idx[b_ids, i_ids].flatten()
+            mconf = conf_matrix[b_ids, i_ids, j_ids]
+   
+            scale = data["hw0_i"][0] / data["hw0_c"][0]
+            scale0 = scale * \
+                data["scale0"][b_ids] if "scale0" in data else scale
+            scale1 = scale * \
+                data["scale1"][b_ids] if "scale1" in data else scale
+            mkpts0_c = (
+                torch.stack(
+                    [
+                        i_ids % data["hw0_c"][1],
+                        torch.div(i_ids, data["hw0_c"][1],
+                                  rounding_mode="floor"),
+                    ],
+                    dim=1,
+                )
+                * scale0
+            )
+            mkpts1_c = (
+                torch.stack(
+                    [
+                        j_ids % data["hw1_c"][1],
+                        torch.div(j_ids, data["hw1_c"][1],
+                                  rounding_mode="floor"),
+                    ],
+                    dim=1,
+                )
+                * scale1
+            )
+
+            data.update(
+                {
+                    "mconf": mconf,
+                    "mkpts0_c": mkpts0_c,
+                    "mkpts1_c": mkpts1_c,
+                    "b_ids": b_ids,
+                    "i_ids": i_ids,
+                    "j_ids": j_ids,
+                }
+            )
+
+        # 4. Fine-Level Matching
+        K0 = data["i_ids"].shape[0] // data["bs"]
+        K1 = data["j_ids"].shape[0] // data["bs"]
+        feat_f0 = feat_f0[data["b_ids"], data["i_ids"]
+                          ].reshape(data["bs"], K0, -1)
+        feat_f1 = feat_f1[data["b_ids"], data["j_ids"]
+                          ].reshape(data["bs"], K1, -1)
+        feat_c0 = feat_c0[data["b_ids"], data["i_ids"]
+                          ].reshape(data["bs"], K0, -1)
+        feat_c1 = feat_c1[data["b_ids"], data["j_ids"]
+                          ].reshape(data["bs"], K1, -1)
+
+        if self.bi_directional_refine:
+            # Bidirectional Refinement
+            offset, score = self.fine_matching(
+                torch.cat([feat_f0, feat_f1], dim=1),
+                torch.cat([feat_f1, feat_f0], dim=1),
+                torch.cat([feat_c0, feat_c1], dim=1),
+                torch.cat([feat_c1, feat_c0], dim=1),
+                data,
+            )
+        else:
+            offset, score = self.fine_matching(
+                feat_f0, feat_f1, feat_c0, feat_c1, data)
+
+        if self.deploy:
+            if self.bi_directional_refine:
+                fine_offset01, fine_offset10 = offset.chunk(2)
+                fine_score01, fine_score10 = score.unsqueeze(dim=1).chunk(2)
+                output = torch.cat(
+                    [mkpts0_c, mkpts1_c, fine_offset01, fine_offset10, fine_score01, fine_score10, mconf.unsqueeze(dim=1)], 1) # [K, 11]
+            else:
+                output = torch.cat(
+                    [mkpts0_c, mkpts1_c, offset, score, mconf.unsqueeze(dim=1)], 1)
+            return output
+
+    def load_state_dict(self, state_dict, *args, **kwargs):
+        for k in list(state_dict.keys()):
+            if k.startswith("matcher."):
+                state_dict[k.replace("matcher.", "", 1)] = state_dict.pop(k)
+        return super().load_state_dict(state_dict, *args, **kwargs)

src/edm/head/coarse_matching.py

@@ -0,0 +1,152 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+INF = 1e9  # -1e4 for fp16matmul
+
+
+class CoarseMatching(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        self.window_size = config["local_resolution"]
+        self.thr = config["coarse"]["mconf_thr"]
+        self.temperature = config["coarse"]["dsmax_temperature"]
+        self.ds_opt = config["coarse"]["ds_opt"]
+        self.pad_num = config["coarse"]["train_pad_num"]
+        self.topk = config["coarse"]["topk"]
+        self.deploy = config["deploy"]
+
+    def forward(self, feat_c0, feat_c1, data, mask_c0=None, mask_c1=None):
+        """
+        Args:
+            feat0 (torch.Tensor): [N, L, C]
+            feat1 (torch.Tensor): [N, S, C]
+            data (dict)
+            mask_c0 (torch.Tensor): [N, L] (optional)
+            mask_c1 (torch.Tensor): [N, S] (optional)
+        Update:
+            data (dict): {
+                'b_ids' (torch.Tensor): [M'],
+                'i_ids' (torch.Tensor): [M'],
+                'j_ids' (torch.Tensor): [M'],
+                'm_bids' (torch.Tensor): [M],
+                'mkpts0_c' (torch.Tensor): [M, 2],
+                'mkpts1_c' (torch.Tensor): [M, 2],
+                'mconf' (torch.Tensor): [M]}
+            NOTE: M' != M during training.
+        """
+        # normalize
+        feat_c0, feat_c1 = map(
+            lambda feat: feat / feat.shape[-1] ** 0.5, [feat_c0, feat_c1]
+        )
+
+        with torch.autocast(enabled=False, device_type="cuda"):
+            sim_matrix = (
+                torch.einsum("nlc,nsc->nls", feat_c0, feat_c1) /
+                self.temperature
+            )
+            del feat_c0, feat_c1
+            if mask_c0 is not None:
+                sim_matrix = sim_matrix.float().masked_fill(
+                    ~(mask_c0[..., None] * mask_c1[:, None]).bool(),
+                    -INF,
+                )
+
+        if not self.training and self.ds_opt:
+            # Alternative implementation of daul-softmax operator for efficient inference
+            sim_matrix = torch.exp(sim_matrix)
+            conf_matrix = F.normalize(sim_matrix, p=1, dim=1) * F.normalize(
+                sim_matrix, p=1, dim=2
+            )
+        else:
+            # Native daul-softmax operator
+            conf_matrix = F.softmax(sim_matrix, 1) * F.softmax(sim_matrix, 2)
+        
+        data.update(
+            {
+                "conf_matrix": conf_matrix,
+            }
+        )
+
+        if not self.deploy:
+            # predict coarse matches from conf_matrix
+            self.coarse_matching_selection(data)
+
+        return conf_matrix # Returning the sim_matrix can be faster, but it may reduce the accuracy. 
+
+    # Static tensor shape for mini-batch inference.
+    @torch.no_grad()
+    def coarse_matching_selection(self, data):
+        """
+        Args:
+            data (dict): with keys ['hw0_i', 'hw1_i', 'hw0_c', 'hw1_c']
+        Returns:
+            coarse_matches (dict): {
+                'b_ids' (torch.Tensor): [M'],
+                'i_ids' (torch.Tensor): [M'],
+                'j_ids' (torch.Tensor): [M'],
+                'm_bids' (torch.Tensor): [M],
+                'mkpts0_c' (torch.Tensor): [M, 2],
+                'mkpts1_c' (torch.Tensor): [M, 2],
+                'mconf' (torch.Tensor): [M]}
+        """
+        conf_matrix = data["conf_matrix"]
+
+        # mutual nearest
+        # mask = (conf_matrix == conf_matrix.max(dim=2, keepdim=True)[0]) \
+        #     * (conf_matrix == conf_matrix.max(dim=1, keepdim=True)[0])
+        # conf_matrix[~mask]=0
+
+        k = self.topk
+        row_max_val, row_max_idx = torch.max(conf_matrix, dim=2)
+
+        # prevent out of range
+        if k == -1 or k > row_max_val.shape[-1]:
+            k = row_max_val.shape[-1]
+
+        topk_val, topk_idx = torch.topk(row_max_val, k)
+        b_ids = (
+            torch.arange(conf_matrix.shape[0], device=conf_matrix.device)
+            .unsqueeze(1)
+            .repeat(1, k)
+            .flatten()
+        )
+        i_ids = topk_idx.flatten()
+        j_ids = row_max_idx[b_ids, i_ids].flatten()
+        mconf = conf_matrix[b_ids, i_ids, j_ids]
+
+        scale = data["hw0_i"][0] / data["hw0_c"][0]
+        scale0 = scale * data["scale0"][b_ids] if "scale0" in data else scale
+        scale1 = scale * data["scale1"][b_ids] if "scale1" in data else scale
+        mkpts0_c = (
+            torch.stack(
+                [
+                    i_ids % data["hw0_c"][1],
+                    torch.div(i_ids, data["hw0_c"][1], rounding_mode="floor"),
+                ],
+                dim=1,
+            )
+            * scale0
+        )
+        mkpts1_c = (
+            torch.stack(
+                [
+                    j_ids % data["hw1_c"][1],
+                    torch.div(j_ids, data["hw1_c"][1], rounding_mode="floor"),
+                ],
+                dim=1,
+            )
+            * scale1
+        )
+
+        data.update(
+            {
+                "mconf": mconf,
+                "mkpts0_c": mkpts0_c,
+                "mkpts1_c": mkpts1_c,
+                "b_ids": b_ids,
+                "i_ids": i_ids,
+                "j_ids": j_ids,
+            }
+        )

src/edm/head/fine_matching.py

@@ -0,0 +1,383 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torch import distributions
+
+
+class Conv1d_BN_Act(nn.Sequential):
+    def __init__(
+        self,
+        a,
+        b,
+        ks=1,
+        stride=1,
+        pad=0,
+        dilation=1,
+        groups=1,
+        bn_weight_init=1,
+        act=None,
+        drop=None,
+    ):
+        super().__init__()
+        self.inp_channel = a
+        self.out_channel = b
+        self.ks = ks
+        self.pad = pad
+        self.stride = stride
+        self.dilation = dilation
+        self.groups = groups
+
+        self.add_module(
+            "c", nn.Conv1d(a, b, ks, stride, pad, dilation, groups, bias=False)
+        )
+        bn = nn.BatchNorm1d(b)
+        nn.init.constant_(bn.weight, bn_weight_init)
+        nn.init.constant_(bn.bias, 0)
+        self.add_module("bn", bn)
+        if act != None:
+            self.add_module("a", act)
+        if drop != None:
+            self.add_module("d", nn.Dropout(drop))
+
+
+class RealNVP(nn.Module):
+    """RealNVP: a flow-based generative model
+
+    `Density estimation using Real NVP
+    arXiv: <https://arxiv.org/abs/1605.08803>`_.
+
+    Code is modified from `the mmpose implementation of RLE
+    <https://github.com/open-mmlab/mmpose/blob/main/mmpose/models/utils/realnvp.py>`_.
+    """
+
+    @staticmethod
+    def get_scale_net(channel):
+        """Get the scale model in a single invertable mapping."""
+        return nn.Sequential(
+            nn.Linear(2, channel),
+            nn.GELU(),
+            nn.Linear(channel, channel),
+            nn.GELU(),
+            nn.Linear(channel, 2),
+            nn.Tanh(),
+        )
+
+    @staticmethod
+    def get_trans_net(channel):
+        """Get the translation model in a single invertable mapping."""
+        return nn.Sequential(
+            nn.Linear(2, channel),
+            nn.GELU(),
+            nn.Linear(channel, channel),
+            nn.GELU(),
+            nn.Linear(channel, 2),
+        )
+
+    @property
+    def prior(self):
+        """The prior distribution."""
+        return distributions.MultivariateNormal(self.loc, self.cov)
+
+    def __init__(self, channel=64):
+        super(RealNVP, self).__init__()
+        self.channel = channel
+        self.register_buffer("loc", torch.zeros(2))
+        self.register_buffer("cov", torch.eye(2))
+        self.register_buffer(
+            "mask", torch.tensor([[0, 1], [1, 0]] * 3, dtype=torch.float32)
+        )
+
+        self.s = torch.nn.ModuleList(
+            [self.get_scale_net(self.channel) for _ in range(len(self.mask))]
+        )
+        self.t = torch.nn.ModuleList(
+            [self.get_trans_net(self.channel) for _ in range(len(self.mask))]
+        )
+        self.init_weights()
+
+    def init_weights(self):
+        """Initialization model weights."""
+        for m in self.modules():
+            if isinstance(m, nn.Linear):
+                nn.init.xavier_uniform_(m.weight, gain=0.01)
+
+    def backward_p(self, x):
+        """Apply mapping form the data space to the latent space and calculate
+        the log determinant of the Jacobian matrix."""
+
+        log_det_jacob, z = x.new_zeros(x.shape[0]), x
+        for i in reversed(range(len(self.t))):
+            z_ = self.mask[i] * z
+            s = self.s[i](z_) * (1 - self.mask[i])  # torch.exp(s): betas
+            t = self.t[i](z_) * (1 - self.mask[i])  # gammas
+            z = (1 - self.mask[i]) * (z - t) * torch.exp(-s) + z_
+            log_det_jacob -= s.sum(dim=1)
+        return z, log_det_jacob
+
+    def log_prob(self, x):
+        """Calculate the log probability of given sample in data space."""
+
+        z, log_det = self.backward_p(x)
+        return self.prior.log_prob(z) + log_det
+
+
+def soft_argmax(x, temperature=1.0):
+    L = x.shape[1]
+    assert L % 2  # L is odd to ensure symmetry
+    idx = torch.arange(0, L, 1, device=x.device).repeat(x.shape[0], 1)
+    scale_x = x / temperature
+    out = F.softmax(scale_x, dim=1) * idx
+    out = out.sum(dim=1, keepdim=True)
+
+    return out
+
+
+class FineMatching(nn.Module):
+    def __init__(self, config, act_layer=nn.GELU):
+        super(FineMatching, self).__init__()
+        self.config = config
+
+        self.block_dims = self.config["backbone"]["block_dims"]
+        self.local_resolution = self.config["local_resolution"]
+        self.drop = self.config["fine"]["droprate"]
+        self.coord_length = self.config["fine"]["coord_length"]
+        self.bi_directional_refine = self.config["fine"]["bi_directional_refine"]
+        self.sigma_selection = self.config["fine"]["sigma_selection"]
+        self.mconf_thr = self.config["coarse"]["mconf_thr"]
+        self.sigma_thr = self.config["fine"]["sigma_thr"]
+        self.border_rm = self.config["coarse"]["border_rm"] * \
+            self.local_resolution
+        self.deploy = self.config["deploy"]
+
+        # network
+        self.query_encoder = nn.Sequential(
+            Conv1d_BN_Act(
+                self.block_dims[-1],
+                self.block_dims[-1],
+                act=act_layer(),
+                drop=self.drop,
+            ),
+            Conv1d_BN_Act(
+                self.block_dims[-1],
+                self.block_dims[-1],
+                act=act_layer(),
+                drop=self.drop,
+            ),
+        )
+
+        self.reference_encoder = nn.Sequential(
+            Conv1d_BN_Act(
+                self.block_dims[-1],
+                self.block_dims[-1],
+                act=act_layer(),
+                drop=self.drop,
+            ),
+            Conv1d_BN_Act(
+                self.block_dims[-1],
+                self.block_dims[-1],
+                act=act_layer(),
+                drop=self.drop,
+            ),
+        )
+
+        self.merge_qr = nn.Sequential(
+            Conv1d_BN_Act(
+                self.block_dims[-1] * 2,
+                self.block_dims[-1] * 2,
+                act=act_layer(),
+                drop=self.drop,
+            ),
+            Conv1d_BN_Act(
+                self.block_dims[-1] * 2,
+                self.block_dims[-1] * 2,
+                act=act_layer(),
+                drop=self.drop,
+            ),
+        )
+
+        self.x_head = nn.Sequential(
+            Conv1d_BN_Act(
+                self.block_dims[-1] * 2,
+                self.block_dims[-1] * 2,
+                act=act_layer(),
+                drop=self.drop,
+            ),
+            nn.Conv1d(self.block_dims[-1] * 2,
+                      self.coord_length + 2, kernel_size=1),
+        )
+
+        self.y_head = nn.Sequential(
+            Conv1d_BN_Act(
+                self.block_dims[-1] * 2,
+                self.block_dims[-1] * 2,
+                act=act_layer(),
+                drop=self.drop,
+            ),
+            nn.Conv1d(self.block_dims[-1] * 2,
+                      self.coord_length + 2, kernel_size=1),
+        )
+
+        self.flow = RealNVP()
+        self.init_params()
+
+    def init_params(self):
+        for m in self.modules():
+            if isinstance(m, nn.Conv1d):
+                nn.init.kaiming_normal_(m.weight)
+                if m.bias is not None:
+                    nn.init.constant_(m.bias, 0)
+
+    def forward(self, feat_f0, feat_f1, feat_c0, feat_c1, data={}):
+        q = self.query_encoder(
+            feat_f0.permute(0, 2, 1).contiguous()
+            + feat_c0.permute(0, 2, 1).contiguous()
+        )
+        r = self.reference_encoder(
+            feat_f1.permute(0, 2, 1).contiguous()
+            + feat_c1.permute(0, 2, 1).contiguous()
+        )
+        out = self.merge_qr(torch.cat([q, r], dim=1))
+
+        x = self.x_head(out).permute(0, 2, 1).contiguous()
+        y = self.y_head(out).permute(0, 2, 1).contiguous()
+
+        if self.bi_directional_refine:
+            x01, x10 = x.chunk(2, dim=1)
+            x01 = x01.reshape(-1, self.coord_length + 2)
+            x10 = x10.reshape(-1, self.coord_length + 2)
+            x_out = torch.cat([x01, x10])
+
+            y01, y10 = y.chunk(2, dim=1)
+            y01 = y01.reshape(-1, self.coord_length + 2)
+            y10 = y10.reshape(-1, self.coord_length + 2)
+            y_out = torch.cat([y01, y10])
+        else:
+            x_out = x.reshape(-1, self.coord_length + 2)
+            y_out = y.reshape(-1, self.coord_length + 2)
+
+        x_cls = x_out[:, : self.coord_length + 1]
+        coord_x = soft_argmax(x_cls) / self.coord_length - \
+            0.5  # range [-0.5, +0.5]
+        x_sigma = x_out[:, -1:].sigmoid()
+
+        y_cls = y_out[:, : self.coord_length + 1]
+        coord_y = soft_argmax(y_cls) / self.coord_length - 0.5
+        y_sigma = y_out[:, -1:].sigmoid()
+
+        coord = torch.cat([coord_x, coord_y], dim=1)
+        sigma = torch.cat([x_sigma, y_sigma], dim=1)
+
+        if data.get("target_uv", None) is not None:
+            gt_uv = data["target_uv"]
+            mask = data["target_uv_weight"].clone()
+
+            if mask.sum() == 0:
+                mask[0] = True
+            mask_coord = coord[mask]
+            mask_gt_uv = gt_uv[mask]
+            mask_sigma = sigma[mask]
+
+            mask_sigma = torch.clamp(mask_sigma, 1e-6, 1 - 1e-6)
+            bar_mu = (mask_coord - mask_gt_uv) / mask_sigma
+
+            log_phi = self.flow.log_prob(bar_mu).unsqueeze(-1)
+            nf_loss = torch.log(mask_sigma) - log_phi
+
+            data.update(
+                {
+                    "pred_coord": coord,
+                    "pred_score": 1.0 - torch.mean(sigma, dim=-1).flatten(),
+                    "mask_coord": mask_coord,
+                    "mask_sigma": mask_sigma,
+                    "nf_loss": nf_loss,
+                }
+            )
+        else:
+            data.update(
+                {
+                    "pred_coord": coord,
+                    "pred_score": 1.0 - torch.mean(sigma, dim=-1).flatten(),
+                }
+            )
+
+        if not self.deploy:
+            self.final_matching_selection(data)
+
+        return data["pred_coord"], data["pred_score"]
+
+    @torch.no_grad()
+    def final_matching_selection(self, data):
+        offset = data["pred_coord"] * self.local_resolution
+
+        if self.bi_directional_refine:
+            fine_offset01, fine_offset10 = torch.clamp(
+                offset, -self.local_resolution / 2, self.local_resolution / 2
+            ).chunk(2)
+        else:
+            fine_offset01 = torch.clamp(
+                offset, -self.local_resolution / 2, self.local_resolution / 2
+            )
+
+        h0, w0 = data["hw0_i"]
+        h1, w1 = data["hw1_i"]
+        scale0 = data["scale0"][data["b_ids"]] if "scale0" in data else 1.0
+        scale1 = data["scale1"][data["b_ids"]] if "scale1" in data else 1.0
+        scale0_w = scale0[:, 0] if "scale0" in data else 1.0
+        scale0_h = scale0[:, 1] if "scale0" in data else 1.0
+        scale1_w = scale1[:, 0] if "scale1" in data else 1.0
+        scale1_h = scale1[:, 1] if "scale1" in data else 1.0
+
+        # Filter by mconf and border
+        mkpts0_f = data["mkpts0_c"]
+        mkpts1_f = data["mkpts1_c"] + fine_offset01 * scale1
+        mask = (
+            (data["mconf"] > self.mconf_thr)
+            & (mkpts0_f[:, 0] >= self.border_rm)
+            & (mkpts0_f[:, 0] <= w0 * scale0_w - self.border_rm)
+            & (mkpts0_f[:, 1] >= self.border_rm)
+            & (mkpts0_f[:, 1] <= h0 * scale0_h - self.border_rm)
+            & (mkpts1_f[:, 0] >= self.border_rm)
+            & (mkpts1_f[:, 0] <= w1 * scale1_w - self.border_rm)
+            & (mkpts1_f[:, 1] >= self.border_rm)
+            & (mkpts1_f[:, 1] <= h1 * scale1_h - self.border_rm)
+        )
+        if self.bi_directional_refine:
+            mkpts0_f_ = data["mkpts0_c"] + fine_offset10 * scale0
+            mkpts1_f_ = data["mkpts1_c"]
+            mask_ = (
+                (data["mconf"] > self.mconf_thr)
+                & (mkpts0_f_[:, 0] >= self.border_rm)
+                & (mkpts0_f_[:, 0] <= w0 * scale0_w - self.border_rm)
+                & (mkpts0_f_[:, 1] >= self.border_rm)
+                & (mkpts0_f_[:, 1] <= h0 * scale0_h - self.border_rm)
+                & (mkpts1_f_[:, 0] >= self.border_rm)
+                & (mkpts1_f_[:, 0] <= w1 * scale1_w - self.border_rm)
+                & (mkpts1_f_[:, 1] >= self.border_rm)
+                & (mkpts1_f_[:, 1] <= h1 * scale1_h - self.border_rm)
+            )
+
+        if self.bi_directional_refine:
+            mkpts0_f = torch.cat([mkpts0_f, mkpts0_f_])
+            mkpts1_f = torch.cat([mkpts1_f, mkpts1_f_])
+            mask = torch.cat([mask, mask_])
+            data["mconf"] = torch.cat([data["mconf"], data["mconf"]])
+            data["b_ids"] = torch.cat([data["b_ids"], data["b_ids"]])
+
+        # Filter by sigma
+        if self.bi_directional_refine and self.sigma_selection:
+            # Retain the more confident matching pair with a smaller sigma (more significant) in the bi-directional matching pairs
+            pred_score01, pred_score10 = data["pred_score"].chunk(2)
+            pred_score_mask = pred_score01 > pred_score10
+            pred_score_mask = torch.cat([pred_score_mask, ~pred_score_mask])
+            pred_score_mask &= data["pred_score"] > self.sigma_thr
+            mask &= pred_score_mask
+
+        data.update(
+            {
+                # "gt_mask": data["mconf"] == 0,
+                "m_bids": data["b_ids"][mask],
+                "mkpts0_f": mkpts0_f[mask],
+                "mkpts1_f": mkpts1_f[mask],
+                "mconf": data["mconf"][mask],
+            }
+        )

src/edm/neck/__init__.py

@@ -0,0 +1 @@
+from .loftr_module import LocalFeatureTransformer

src/edm/neck/loftr_module/__init__.py

@@ -0,0 +1 @@
+from .transformer import LocalFeatureTransformer

src/edm/neck/loftr_module/transformer.py

@@ -0,0 +1,418 @@
+from torch.nn import Module
+import copy
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from einops.einops import rearrange
+import math
+import os
+import sys
+
+sys.path.append(os.path.dirname(os.path.abspath(__file__)))
+
+
+class RoPEPositionEncodingSine(nn.Module):
+    """
+    This is a sinusoidal position encoding that generalized to 2-dimensional images
+    """
+
+    def __init__(self, d_model, max_shape=(128, 128), npe=None, ropefp16=True):
+        """
+        Args:
+            max_shape (tuple): for 1/32 featmap, the max length of 128 corresponds to 4096 pixels
+        """
+        super().__init__()
+
+        i_position = torch.ones(max_shape).cumsum(
+            0).float().unsqueeze(-1)  # [H, 1]
+        j_position = torch.ones(max_shape).cumsum(
+            1).float().unsqueeze(-1)  # [W, 1]
+
+        assert npe is not None
+        train_res_H, train_res_W, test_res_H, test_res_W = (
+            npe[0],
+            npe[1],
+            npe[2],
+            npe[3],
+        )
+        i_position, j_position = (
+            i_position * train_res_H / test_res_H,
+            j_position * train_res_W / test_res_W,
+        )
+
+        div_term = torch.exp(
+            torch.arange(0, d_model // 4, 1).float()
+            * (-math.log(10000.0) / (d_model // 4))
+        )
+        div_term = div_term[None, None, :]  # [1, 1, C//4]
+
+        sin = torch.zeros(
+            *max_shape, d_model // 2, dtype=torch.float16 if ropefp16 else torch.float32
+        )
+        cos = torch.zeros(
+            *max_shape, d_model // 2, dtype=torch.float16 if ropefp16 else torch.float32
+        )
+
+        sin[:, :, 0::2] = (
+            torch.sin(i_position * div_term).half()
+            if ropefp16
+            else torch.sin(i_position * div_term)
+        )
+        sin[:, :, 1::2] = (
+            torch.sin(j_position * div_term).half()
+            if ropefp16
+            else torch.sin(j_position * div_term)
+        )
+        cos[:, :, 0::2] = (
+            torch.cos(i_position * div_term).half()
+            if ropefp16
+            else torch.cos(i_position * div_term)
+        )
+        cos[:, :, 1::2] = (
+            torch.cos(j_position * div_term).half()
+            if ropefp16
+            else torch.cos(j_position * div_term)
+        )
+
+        sin = sin.repeat_interleave(2, dim=-1)
+        cos = cos.repeat_interleave(2, dim=-1)
+
+        self.register_buffer(
+            "sin", sin.unsqueeze(0), persistent=False
+        )  # [1, H, W, C//2]
+        self.register_buffer(
+            "cos", cos.unsqueeze(0), persistent=False
+        )  # [1, H, W, C//2]
+
+    def forward(self, x, ratio=1):
+        """
+        Args:
+            x: [N, H, W, C]
+        """
+        return (x * self.cos[:, : x.size(1), : x.size(2), :]) + (
+            self.rotate_half(x) * self.sin[:, : x.size(1), : x.size(2), :]
+        )
+
+    def rotate_half(self, x):
+        # x = x.unflatten(-1, (-1, 2))
+        a, b, c, d = x.shape
+        x = x.reshape(a, b, c, d // 2, 2)
+
+        x1, x2 = x.unbind(dim=-1)
+        return torch.stack((-x2, x1), dim=-1).flatten(start_dim=-2)
+
+
+"""
+Linear Transformer proposed in "Transformers are RNNs: Fast Autoregressive Transformers with Linear Attention"
+Modified from: https://github.com/idiap/fast-transformers/blob/master/fast_transformers/attention/linear_attention.py
+"""
+
+
+def crop_feature(query, key, value, x_mask, source_mask):
+    mask_h0, mask_w0, mask_h1, mask_w1 = (
+        x_mask[0].sum(-2)[0],
+        x_mask[0].sum(-1)[0],
+        source_mask[0].sum(-2)[0],
+        source_mask[0].sum(-1)[0],
+    )
+    query = query[:, :mask_h0, :mask_w0, :]
+    key = key[:, :mask_h1, :mask_w1, :]
+    value = value[:, :mask_h1, :mask_w1, :]
+    return query, key, value, mask_h0, mask_w0
+
+
+def pad_feature(m, mask_h0, mask_w0, x_mask):
+    bs, hw, nhead, dim = m.shape
+    m = m.view(bs, mask_h0, mask_w0, nhead, dim)
+    if mask_h0 != x_mask.size(-2):
+        m = torch.cat(
+            [
+                m,
+                torch.zeros(
+                    m.size(0),
+                    x_mask.size(-2) - mask_h0,
+                    x_mask.size(-1),
+                    nhead,
+                    dim,
+                    device=m.device,
+                    dtype=m.dtype,
+                ),
+            ],
+            dim=1,
+        )
+    elif mask_w0 != x_mask.size(-1):
+        m = torch.cat(
+            [
+                m,
+                torch.zeros(
+                    m.size(0),
+                    x_mask.size(-2),
+                    x_mask.size(-1) - mask_w0,
+                    nhead,
+                    dim,
+                    device=m.device,
+                    dtype=m.dtype,
+                ),
+            ],
+            dim=2,
+        )
+    return m
+
+
+class Attention(Module):
+    def __init__(self, nhead=8, dim=256, re=False):
+        super().__init__()
+
+        self.nhead = nhead
+        self.dim = dim
+
+    def attention(self, query, key, value, q_mask=None, kv_mask=None):
+        assert (
+            q_mask is None and kv_mask is None
+        ), "Not support generalized attention mask yet."
+        # Scaled Cosine Attention
+        # Refer to "Query-key normalization for transformers" and "https://kexue.fm/archives/9859"
+        query = F.normalize(query, p=2, dim=3)
+        key = F.normalize(key, p=2, dim=3)
+        QK = torch.einsum("nlhd,nshd->nlsh", query, key)
+        s = 20.0
+        A = torch.softmax(s * QK, dim=2)
+
+        out = torch.einsum("nlsh,nshd->nlhd", A, value)
+        return out
+
+    def _forward(self, query, key, value, q_mask=None, kv_mask=None):
+        if q_mask is not None:
+            query, key, value, mask_h0, mask_w0 = crop_feature(
+                query, key, value, q_mask, kv_mask
+            )
+
+        query, key, value = map(
+            lambda x: rearrange(
+                x,
+                "n h w (nhead d) -> n (h w) nhead d",
+                nhead=self.nhead,
+                d=self.dim,
+            ),
+            [query, key, value],
+        )
+
+        m = self.attention(query, key, value, q_mask=None, kv_mask=None)
+
+        if q_mask is not None:
+            m = pad_feature(m, mask_h0, mask_w0, q_mask)
+
+        return m
+
+    def forward(self, query, key, value, q_mask=None, kv_mask=None):
+        """
+        Args:
+            queries: [N, L, H, D]
+            keys: [N, S, H, D]
+            values: [N, S, H, D]
+            q_mask: [N, L]
+            kv_mask: [N, S]
+        Returns:
+            queried_values: (N, L, H, D)
+        """
+        bs = query.size(0)
+        if bs == 1 or q_mask is None:
+            m = self._forward(query, key, value,
+                              q_mask=q_mask, kv_mask=kv_mask)
+        else:  # for faster trainning with padding mask while batch size > 1
+            m_list = []
+            for i in range(bs):
+                m_list.append(
+                    self._forward(
+                        query[i: i + 1],
+                        key[i: i + 1],
+                        value[i: i + 1],
+                        q_mask=q_mask[i: i + 1],
+                        kv_mask=kv_mask[i: i + 1],
+                    )
+                )
+            m = torch.cat(m_list, dim=0)
+        return m
+
+
+class AG_RoPE_EncoderLayer(nn.Module):
+    def __init__(
+        self,
+        d_model,
+        nhead,
+        agg_size0=2,
+        agg_size1=2,
+        rope=False,
+        npe=None,
+    ):
+        super(AG_RoPE_EncoderLayer, self).__init__()
+        self.dim = d_model // nhead
+        self.nhead = nhead
+        self.agg_size0, self.agg_size1 = agg_size0, agg_size1
+        self.rope = rope
+
+        # aggregate and position encoding
+        self.aggregate = (
+            nn.Conv2d(
+                d_model,
+                d_model,
+                kernel_size=agg_size0,
+                padding=0,
+                stride=agg_size0,
+                bias=False,
+                groups=d_model,
+            )
+            if self.agg_size0 != 1
+            else nn.Identity()
+        )
+        self.max_pool = (
+            torch.nn.MaxPool2d(kernel_size=self.agg_size1,
+                               stride=self.agg_size1)
+            if self.agg_size1 != 1
+            else nn.Identity()
+        )
+        self.mask_max_pool = torch.nn.MaxPool2d(kernel_size=2, stride=2)
+        self.rope_pos_enc = RoPEPositionEncodingSine(
+            d_model, max_shape=(128, 128), npe=npe, ropefp16=True
+        )
+
+        # multi-head attention
+        self.q_proj = nn.Linear(d_model, d_model, bias=False)
+        self.k_proj = nn.Linear(d_model, d_model, bias=False)
+        self.v_proj = nn.Linear(d_model, d_model, bias=False)
+        self.attention = Attention(self.nhead, self.dim)
+        self.merge = nn.Linear(d_model, d_model, bias=False)
+
+        # feed-forward network
+        self.mlp = nn.Sequential(
+            nn.Linear(d_model * 2, d_model * 2, bias=False),
+            nn.LeakyReLU(inplace=True),
+            nn.Linear(d_model * 2, d_model, bias=False),
+        )
+
+        # norm
+        self.norm1 = nn.LayerNorm(d_model)
+        self.norm2 = nn.LayerNorm(d_model)
+
+    def forward(self, x, source, x_mask=None, source_mask=None):
+        """
+        Args:
+            x (torch.Tensor): [N, C, H0, W0]
+            source (torch.Tensor): [N, C, H1, W1]
+            x_mask (torch.Tensor): [N, H0, W0] (optional) (L = H0*W0)
+            source_mask (torch.Tensor): [N, H1, W1] (optional) (S = H1*W1)
+        """
+        bs, C, H0, W0 = x.size()
+        H1, W1 = source.size(-2), source.size(-1)
+
+        # Aggragate feature
+        # assert x_mask is None and source_mask is None
+
+        query, source = self.norm1(self.aggregate(x).permute(0, 2, 3, 1)), self.norm1(
+            self.max_pool(source).permute(0, 2, 3, 1)
+        )  # [N, H, W, C]
+        if x_mask is not None:
+            # mask 1/8 to 1/32
+            x_mask, source_mask = map(
+                lambda x: self.mask_max_pool(
+                    self.mask_max_pool(x.float())).bool(),
+                [x_mask, source_mask],
+            )
+        query, key, value = self.q_proj(
+            query), self.k_proj(source), self.v_proj(source)
+
+        # Positional encoding
+        if self.rope:
+            query = self.rope_pos_enc(query)
+            key = self.rope_pos_enc(key)
+
+        # multi-head attention handle padding mask
+        m = self.attention(query, key, value, q_mask=x_mask,
+                           kv_mask=source_mask)
+        m = self.merge(m.reshape(bs, -1, self.nhead * self.dim))  # [N, L, C]
+
+        # Upsample feature
+        m = rearrange(
+            m, "b (h w) c -> b c h w", h=H0 // self.agg_size0, w=W0 // self.agg_size0
+        )  # [N, C, H0, W0]
+
+        if self.agg_size0 != 1:
+            m = torch.nn.functional.interpolate(
+                m, size=(H0, W0), mode="bilinear", align_corners=False
+            )  # [N, C, H0, W0]
+
+        # feed-forward network
+        m = self.mlp(torch.cat([x, m], dim=1).permute(
+            0, 2, 3, 1))  # [N, H0, W0, C]
+        m = self.norm2(m).permute(0, 3, 1, 2)  # [N, C, H0, W0]
+
+        return x + m
+
+
+'''
+Modified from EfficientLoFTR
+'''
+class LocalFeatureTransformer(nn.Module):
+    """A Local Feature Transformer (LoFTR) module."""
+
+    def __init__(self, config):
+        super(LocalFeatureTransformer, self).__init__()
+        self.d_model = config["d_model"]
+        self.nhead = config["nhead"]
+        self.layer_names = config["layer_names"]
+        self.agg_size0, self.agg_size1 = config["agg_size0"], config["agg_size1"]
+        self.rope = config["rope"]
+
+        self_layer = AG_RoPE_EncoderLayer(
+            config["d_model"],
+            config["nhead"],
+            config["agg_size0"],
+            config["agg_size1"],
+            config["rope"],
+            config["npe"],
+        )
+        cross_layer = AG_RoPE_EncoderLayer(
+            config["d_model"],
+            config["nhead"],
+            config["agg_size0"],
+            config["agg_size1"],
+            False,
+            config["npe"],
+        )
+
+        self.layers = nn.ModuleList(
+            [
+                (
+                    copy.deepcopy(self_layer)
+                    if _ == "self"
+                    else copy.deepcopy(cross_layer)
+                )
+                for _ in self.layer_names
+            ]
+        )
+        self._reset_parameters()
+
+    def _reset_parameters(self):
+        for p in self.parameters():
+            if p.dim() > 1:
+                nn.init.xavier_uniform_(p)
+
+    def forward(self, feat0, feat1, mask0=None, mask1=None, data=None):
+        """
+        Args:
+            feat0 (torch.Tensor): [N, C, H, W]
+            feat1 (torch.Tensor): [N, C, H, W]
+            mask0 (torch.Tensor): [N, L] (optional)
+            mask1 (torch.Tensor): [N, S] (optional)
+        """
+        for i, (layer, name) in enumerate(zip(self.layers, self.layer_names)):
+            if name == "self":
+                feat0 = layer(feat0, feat0, mask0, mask0)
+                feat1 = layer(feat1, feat1, mask1, mask1)
+            elif name == "cross":
+                feat0 = layer(feat0, feat1, mask0, mask1)
+                feat1 = layer(feat1, feat0, mask1, mask0)
+            else:
+                raise KeyError
+
+        return feat0, feat1

src/edm/neck/neck.py

@@ -0,0 +1,156 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from .loftr_module.transformer import LocalFeatureTransformer
+
+
+class Conv2d_BN_Act(nn.Sequential):
+    def __init__(
+        self,
+        a,
+        b,
+        ks=1,
+        stride=1,
+        pad=0,
+        dilation=1,
+        groups=1,
+        bn_weight_init=1,
+        act=None,
+        drop=None,
+    ):
+        super().__init__()
+        self.inp_channel = a
+        self.out_channel = b
+        self.ks = ks
+        self.pad = pad
+        self.stride = stride
+        self.dilation = dilation
+        self.groups = groups
+
+        self.add_module(
+            "c", nn.Conv2d(a, b, ks, stride, pad, dilation, groups, bias=False)
+        )
+        bn = nn.BatchNorm2d(b)
+        nn.init.constant_(bn.weight, bn_weight_init)
+        nn.init.constant_(bn.bias, 0)
+        self.add_module("bn", bn)
+        if act != None:
+            self.add_module("a", act)
+        if drop != None:
+            self.add_module("d", nn.Dropout(drop))
+
+
+class CIM(nn.Module):
+    """Feature Aggregation, Correlation Injection Module"""
+
+    def __init__(self, config):
+        super(CIM, self).__init__()
+
+        self.block_dims = config["backbone"]["block_dims"]
+        self.drop = config["fine"]["droprate"]
+
+        self.fc32 = Conv2d_BN_Act(
+            self.block_dims[-1], self.block_dims[-1], 1, drop=self.drop
+        )
+        self.fc16 = Conv2d_BN_Act(
+            self.block_dims[-2], self.block_dims[-1], 1, drop=self.drop
+        )
+        self.fc8 = Conv2d_BN_Act(
+            self.block_dims[-3], self.block_dims[-1], 1, drop=self.drop
+        )
+        self.att32 = Conv2d_BN_Act(
+            self.block_dims[-1],
+            self.block_dims[-1],
+            1,
+            act=nn.Sigmoid(),
+            drop=self.drop,
+        )
+        self.att16 = Conv2d_BN_Act(
+            self.block_dims[-1],
+            self.block_dims[-1],
+            1,
+            act=nn.Sigmoid(),
+            drop=self.drop,
+        )
+        self.dwconv16 = nn.Sequential(
+            Conv2d_BN_Act(
+                self.block_dims[-1],
+                self.block_dims[-1],
+                ks=3,
+                pad=1,
+                groups=self.block_dims[-1],
+                act=nn.GELU(),
+            ),
+            Conv2d_BN_Act(self.block_dims[-1], self.block_dims[-1], 1),
+        )
+        self.dwconv8 = nn.Sequential(
+            Conv2d_BN_Act(
+                self.block_dims[-1],
+                self.block_dims[-1],
+                ks=3,
+                pad=1,
+                groups=self.block_dims[-1],
+                act=nn.GELU(),
+            ),
+            Conv2d_BN_Act(self.block_dims[-1], self.block_dims[-1], 1),
+        )
+
+        self.loftr_32 = LocalFeatureTransformer(config["neck"])
+
+    def forward(self, ms_feats, mask_c0=None, mask_c1=None):
+        if len(ms_feats) == 3:  # same image shape
+            f8, f16, f32 = ms_feats
+            f32 = self.fc32(f32)
+
+            f32_0, f32_1 = f32.chunk(2, dim=0)
+            f32_0, f32_1 = self.loftr_32(f32_0, f32_1, mask_c0, mask_c1)
+            f32 = torch.cat([f32_0, f32_1], dim=0)
+
+            f32_up = F.interpolate(f32, scale_factor=2.0, mode="bilinear")
+            att32_up = F.interpolate(self.att32(
+                f32), scale_factor=2.0, mode="bilinear")
+            f16 = self.fc16(f16)
+            f16 = self.dwconv16(f16 * att32_up + f32_up)
+            f16_up = F.interpolate(f16, scale_factor=2.0, mode="bilinear")
+            att16_up = F.interpolate(self.att16(
+                f16), scale_factor=2.0, mode="bilinear")
+            f8 = self.fc8(f8)
+            f8 = self.dwconv8(f8 * att16_up + f16_up)
+
+            feat_c0, feat_c1 = f8.chunk(2)
+
+        elif len(ms_feats) == 6:  # diffirent image shape
+            f8_0, f16_0, f32_0, f8_1, f16_1, f32_1 = ms_feats
+            f32_0 = self.fc32(f32_0)
+            f32_1 = self.fc32(f32_1)
+
+            f32_0, f32_1 = self.loftr_32(f32_0, f32_1, mask_c0, mask_c1)
+
+            f8, f16, f32 = f8_0, f16_0, f32_0
+            f32_up = F.interpolate(f32, scale_factor=2.0, mode="bilinear")
+            att32_up = F.interpolate(self.att32(
+                f32), scale_factor=2.0, mode="bilinear")
+            f16 = self.fc16(f16)
+            f16 = self.dwconv16(f16 * att32_up + f32_up)
+            f16_up = F.interpolate(f16, scale_factor=2.0, mode="bilinear")
+            att16_up = F.interpolate(self.att16(
+                f16), scale_factor=2.0, mode="bilinear")
+            f8 = self.fc8(f8)
+            f8 = self.dwconv8(f8 * att16_up + f16_up)
+            feat_c0 = f8
+
+            f8, f16, f32 = f8_1, f16_1, f32_1
+            f32_up = F.interpolate(f32, scale_factor=2.0, mode="bilinear")
+            att32_up = F.interpolate(self.att32(
+                f32), scale_factor=2.0, mode="bilinear")
+            f16 = self.fc16(f16)
+            f16 = self.dwconv16(f16 * att32_up + f32_up)
+            f16_up = F.interpolate(f16, scale_factor=2.0, mode="bilinear")
+            att16_up = F.interpolate(self.att16(
+                f16), scale_factor=2.0, mode="bilinear")
+            f8 = self.fc8(f8)
+            f8 = self.dwconv8(f8 * att16_up + f16_up)
+            feat_c1 = f8
+
+        return feat_c0, feat_c1

src/utils/misc.py

@@ -0,0 +1,103 @@
+import os
+import contextlib
+import joblib
+from typing import Union
+from loguru import _Logger, logger
+from itertools import chain
+
+import torch
+from yacs.config import CfgNode as CN
+
+
+
+def lower_config(yacs_cfg):
+    if not isinstance(yacs_cfg, CN):
+        return yacs_cfg
+    return {k.lower(): lower_config(v) for k, v in yacs_cfg.items()}
+
+
+def upper_config(dict_cfg):
+    if not isinstance(dict_cfg, dict):
+        return dict_cfg
+    return {k.upper(): upper_config(v) for k, v in dict_cfg.items()}
+
+
+def log_on(condition, message, level):
+    if condition:
+        assert level in ["INFO", "DEBUG", "WARNING", "ERROR", "CRITICAL"]
+        logger.log(level, message)
+
+
+
+def setup_gpus(gpus: Union[str, int]) -> int:
+    """A temporary fix for pytorch-lighting 1.3.x"""
+    gpus = str(gpus)
+    gpu_ids = []
+
+    if "," not in gpus:
+        n_gpus = int(gpus)
+        return n_gpus if n_gpus != -1 else torch.cuda.device_count()
+    else:
+        gpu_ids = [i.strip() for i in gpus.split(",") if i != ""]
+
+    # setup environment variables
+    visible_devices = os.getenv("CUDA_VISIBLE_DEVICES")
+    if visible_devices is None:
+        os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID"
+        os.environ["CUDA_VISIBLE_DEVICES"] = ",".join(str(i) for i in gpu_ids)
+        visible_devices = os.getenv("CUDA_VISIBLE_DEVICES")
+        logger.warning(
+            f"[Temporary Fix] manually set CUDA_VISIBLE_DEVICES when specifying gpus to use: {visible_devices}"
+        )
+    else:
+        logger.warning(
+            "[Temporary Fix] CUDA_VISIBLE_DEVICES already set by user or the main process."
+        )
+    return len(gpu_ids)
+
+
+def flattenList(x):
+    return list(chain(*x))
+
+
+@contextlib.contextmanager
+def tqdm_joblib(tqdm_object):
+    """Context manager to patch joblib to report into tqdm progress bar given as argument
+
+    Usage:
+        with tqdm_joblib(tqdm(desc="My calculation", total=10)) as progress_bar:
+            Parallel(n_jobs=16)(delayed(sqrt)(i**2) for i in range(10))
+
+    When iterating over a generator, directly use of tqdm is also a solutin (but monitor the task queuing, instead of finishing)
+        ret_vals = Parallel(n_jobs=args.world_size)(
+                    delayed(lambda x: _compute_cov_score(pid, *x))(param)
+                        for param in tqdm(combinations(image_ids, 2),
+                                          desc=f'Computing cov_score of [{pid}]',
+                                          total=len(image_ids)*(len(image_ids)-1)/2))
+    Src: https://stackoverflow.com/a/58936697
+    """
+
+    class TqdmBatchCompletionCallback(joblib.parallel.BatchCompletionCallBack):
+        def __init__(self, *args, **kwargs):
+            super().__init__(*args, **kwargs)
+
+        def __call__(self, *args, **kwargs):
+            tqdm_object.update(n=self.batch_size)
+            return super().__call__(*args, **kwargs)
+
+    old_batch_callback = joblib.parallel.BatchCompletionCallBack
+    joblib.parallel.BatchCompletionCallBack = TqdmBatchCompletionCallback
+    try:
+        yield tqdm_object
+    finally:
+        joblib.parallel.BatchCompletionCallBack = old_batch_callback
+        tqdm_object.close()
+
+
+def detect_NaN(feat_0, feat_1):
+    logger.info(f"NaN detected in feature")
+    logger.info(
+        f"#NaN in feat_0: {torch.isnan(feat_0).int().sum()}, #NaN in feat_1: {torch.isnan(feat_1).int().sum()}"
+    )
+    feat_0[torch.isnan(feat_0)] = 0
+    feat_1[torch.isnan(feat_1)] = 0

src/utils/plotting.py

@@ -0,0 +1,219 @@
+import bisect
+import numpy as np
+import matplotlib.pyplot as plt
+import matplotlib
+import cv2
+
+def _compute_conf_thresh(data):
+    dataset_name = data["dataset_name"][0].lower()
+    if dataset_name == "scannet":
+        thr = 5e-4
+    elif dataset_name == "megadepth":
+        thr = 1e-4
+    else:
+        raise ValueError(f"Unknown dataset: {dataset_name}")
+    return thr
+
+
+# --- VISUALIZATION --- #
+
+
+def make_matching_figure(
+    img0,
+    img1,
+    mkpts0,
+    mkpts1,
+    color,
+    kpts0=None,
+    kpts1=None,
+    text=[],
+    path=None,
+):
+    """
+    使用OpenCV绘制匹配点可视化图像
+    
+    参数:
+        img0: 第一张图像 (BGR格式)
+        img1: 第二张图像 (BGR格式)
+        mkpts0: 第一张图像中的匹配点 (Nx2数组)
+        mkpts1: 第二张图像中的匹配点 (Nx2数组)
+        color: 每个匹配点的颜色
+        kpts0: 第一张图像中的所有关键点 (可选)
+        kpts1: 第二张图像中的所有关键点 (可选)
+        text: 要添加的文本 (可选)
+        path: 保存图像的路径 (可选)
+    
+    返回:
+        绘制好的OpenCV图像 (BGR格式)
+    """
+    # 确保匹配点数量一致
+    assert mkpts0.shape[0] == mkpts1.shape[0], \
+        f"mkpts0: {mkpts0.shape[0]} v.s. mkpts1: {mkpts1.shape[0]}"
+    
+    # 确保图像有相同的高度，如果不同则调整
+    h0, w0 = img0.shape[:2]
+    h1, w1 = img1.shape[:2]
+    max_height = max(h0, h1)
+    
+    # 创建画布，两张图像并排显示
+    canvas = np.ones((max_height, w0 + w1, 3), dtype=np.uint8) * 255
+    
+    # 将图像放置到画布上
+    canvas[:h0, :w0] = img0
+    canvas[:h1, w0:w0+w1] = img1
+    
+    # 绘制所有关键点（如果提供）
+    if kpts0 is not None and kpts1 is not None:
+        for (x, y) in kpts0.astype(np.int32):
+            cv2.circle(canvas, (x, y), 1, (255, 255, 255), -1)
+        
+        for (x, y) in kpts1.astype(np.int32):
+            cv2.circle(canvas, (x + w0, y), 1, (255, 255, 255), -1)
+    
+    # 绘制匹配点和连接线
+    if mkpts0.shape[0] > 0 and mkpts1.shape[0] > 0:
+        # 转换为整数坐标
+        mkpts0_int = mkpts0.astype(np.int32)
+        mkpts1_int = mkpts1.astype(np.int32)
+        
+        # 绘制连接线
+        for i in range(len(mkpts0_int)):
+            x0, y0 = mkpts0_int[i]
+            x1, y1 = mkpts1_int[i]
+            # 第二张图的x坐标需要加上第一张图的宽度
+            x1 += w0
+            
+            # 将颜色从0-1范围转换为0-255
+            line_color = tuple(int(c * 255) for c in color[i][:3])
+            # 转换为BGR格式（因为OpenCV使用BGR）
+            # line_color = (line_color[2], line_color[1], line_color[0])
+            
+            cv2.line(canvas, (x0, y0), (x1, y1), line_color, 1)
+        
+        # 绘制匹配点
+        for i in range(len(mkpts0_int)):
+            x0, y0 = mkpts0_int[i]
+            x1, y1 = mkpts1_int[i]
+            x1 += w0
+            
+            pt_color = tuple(int(c * 255) for c in color[i][:3])
+            # pt_color = (pt_color[2], pt_color[1], pt_color[0])
+            
+            cv2.circle(canvas, (x0, y0), 2, pt_color, -1)
+            cv2.circle(canvas, (x1, y1), 2, pt_color, -1)
+    
+    # 添加文本
+    if text:
+        # 确定文本颜色（基于图像亮度）
+        roi = img0[:100, :200] if h0 > 100 and w0 > 200 else img0
+        brightness = cv2.cvtColor(roi, cv2.COLOR_BGR2GRAY).mean()
+        text_color = (0, 0, 0) if brightness > 200 else (255, 255, 255)
+        
+        # 绘制文本
+        y_pos = 30
+        for i, line in enumerate(text):
+            cv2.putText(
+                canvas, line, (10, y_pos + i * 30),
+                cv2.FONT_HERSHEY_SIMPLEX, 0.8, text_color, 2
+            )
+    
+    # 保存图像（如果指定了路径）
+    if path:
+        cv2.imwrite(path, canvas)
+    
+    return canvas
+
+def _make_evaluation_figure(data, b_id, alpha="dynamic"):
+    b_mask = data["m_bids"] == b_id
+    conf_thr = _compute_conf_thresh(data)
+
+    img0 = (data["image0"][b_id][0].cpu().numpy()
+            * 255).round().astype(np.int32)
+    img1 = (data["image1"][b_id][0].cpu().numpy()
+            * 255).round().astype(np.int32)
+    kpts0 = data["mkpts0_f"][b_mask].clone().detach().cpu().numpy()
+    kpts1 = data["mkpts1_f"][b_mask].clone().detach().cpu().numpy()
+
+    # for megadepth, we visualize matches on the resized image
+    if "scale0" in data:
+        kpts0 = kpts0 / data["scale0"][b_id].cpu().numpy()[[1, 0]]
+        kpts1 = kpts1 / data["scale1"][b_id].cpu().numpy()[[1, 0]]
+
+    epi_errs = data["epi_errs"][b_mask].cpu().numpy()
+    correct_mask = epi_errs < conf_thr
+    precision = np.mean(correct_mask) if len(correct_mask) > 0 else 0
+    n_correct = np.sum(correct_mask)
+    n_gt_matches = int(data["conf_matrix_gt"][b_id].sum().cpu())
+    recall = 0 if n_gt_matches == 0 else n_correct / (n_gt_matches)
+    # recall might be larger than 1, since the calculation of conf_matrix_gt
+    # uses groundtruth depths and camera poses, but epipolar distance is used here.
+
+    # matching info
+    if alpha == "dynamic":
+        alpha = dynamic_alpha(len(correct_mask))
+    color = error_colormap(epi_errs, conf_thr, alpha=alpha)
+    
+    text = [
+        f"#Matches {len(kpts0)}",
+        f"Precision({conf_thr:.2e}) ({100 * precision:.1f}%): {n_correct}/{len(kpts0)}",
+        f"Recall({conf_thr:.2e}) ({100 * recall:.1f}%): {n_correct}/{n_gt_matches}",
+    ]
+
+    # make the figure
+    figure = make_matching_figure(img0, img1, kpts0, kpts1, color, text=text)
+    return figure
+
+
+def _make_confidence_figure(data, b_id):
+    # TODO: Implement confidence figure
+    raise NotImplementedError()
+
+
+def make_matching_figures(data, config, mode="evaluation"):
+    """Make matching figures for a batch.
+
+    Args:
+        data (Dict): a batch updated by PL_LoFTR.
+        config (Dict): matcher config
+    Returns:
+        figures (Dict[str, List[plt.figure]]
+    """
+    assert mode in ["evaluation", "confidence", "gt"]  # 'confidence'
+    figures = {mode: []}
+    for b_id in range(data["image0"].size(0)):
+        if mode == "evaluation":
+            fig = _make_evaluation_figure(
+                data, b_id, alpha=config.TRAINER.PLOT_MATCHES_ALPHA
+            )
+        elif mode == "confidence":
+            fig = _make_confidence_figure(data, b_id)
+        else:
+            raise ValueError(f"Unknown plot mode: {mode}")
+        figures[mode].append(fig)
+    return figures
+
+
+def dynamic_alpha(
+    n_matches, milestones=[0, 300, 1000, 2000], alphas=[1.0, 0.8, 0.4, 0.2]
+):
+    if n_matches == 0:
+        return 1.0
+    ranges = list(zip(alphas, alphas[1:] + [None]))
+    loc = bisect.bisect_right(milestones, n_matches) - 1
+    _range = ranges[loc]
+    if _range[1] is None:
+        return _range[0]
+    return _range[1] + (milestones[loc + 1] - n_matches) / (
+        milestones[loc + 1] - milestones[loc]
+    ) * (_range[0] - _range[1])
+
+
+def error_colormap(err, thr, alpha=1.0):
+    assert alpha <= 1.0 and alpha > 0, f"Invaid alpha value: {alpha}"
+    x = 1 - np.clip(err / (thr * 2), 0, 1)
+    return np.clip(
+        np.stack([2 - x * 2, x * 2, np.zeros_like(x),
+                 np.ones_like(x) * alpha], -1),
+        0,
+        1,
+    )