7、关于LoFTR

7、关于LoFTR

LoFTR论文链接：LoFTR

LoFTR的提出，是将Transformer模型的注意力机制在特征匹配方向的应用，Transformer的提取特征的机制，在自身进行，本文提出可以的两张图像之间进行特征计算，非常适合进行特征匹配。

一、传统匹配模式的局限

1. 首先找到一些关键点(图像梯度较高，角点检测等)，然后计算特征相似度来匹配，这种方法很依赖检测到的特征点，一旦点找不到，那就不用说匹配了

2. 对于位置不同的两个点，如果它们的背景特征相似(与位置无关了)，也无法匹配，两者差别不大的，就无法区分了

二、LoFTR解决方案

1. 不需要先得到特征点，第一个问题就得到了解决
2. 采用end2end（给一个输入，得到结果）方法，使用非常方便

三、模型架构图

模型进行特征匹配主要分为四步：

1. Local Feature CNN，就是一个blockbone，对两张图特征图进行卷积，分别得到原来特征图大小的1/8（用来进行粗粒度的匹配）和1/2（用来进行细粒度的匹配） 大小特征图。
2. Coarse-Level Local Feature Transformer,是对两张1/8的特征图（FA，FB）进行多个串联的self-attention（自己的q和k向量做内积）和cross-attention（FA提供q，FB提供V做内积）来计算自己每个区域内之间的关系和两张图不同区域内的关系。
3. Matching Module，到目前位置都是进行粗粒度匹配，经过了多层的Coarse-Level Local Feature Transformer，FA，FB已经了解了自身的关系，同时也知道了和对方之间的关系，要进行关系的匹配了，采用互近邻(mutual nearest neighbor (MNN))的方式匹配，简单来说，假设FA对某个点信息概率值大， 但是同时也要某个点对A的概率值也大，必须双向的，源码阈值设置为0.2，也就是互相的结果都大于0.2才能匹配成功，筛选得到符合阈值点传入下一层。
4. Coarse-to-Fine Module，在经过粗粒度匹配之后，会得到一些候选点区域，这些候选点区域会传到细粒度，细粒度其实就是在匹配的小区域上再做一次Coarse-Level Local Feature Transformer，同样会互相了解对方的特征，但是这里不是采用MNN机制，假设FA的小区域中的某个点有FB对应小区域中的所有点匹配的概率值，才将这写概率值绘制成类似于热力图，再对整个图计算期望值，得到最终的匹配点坐标。

核心代码

class LoFTR(nn.Module):、# ... 省略初始化代码def forward(self, data):""" 前向传播函数：参数：data (dict): 包含图像和可选掩码的字典'image0': (torch.Tensor): (N, 1, H, W) 第一张图像'image1': (torch.Tensor): (N, 1, H, W) 第二张图像'mask0' (可选): (torch.Tensor): (N, H, W) 第一张图像的掩码，'0' 表示填充位置'mask1' (可选): (torch.Tensor): (N, H, W) 第二张图像的掩码"""# 更新数据字典，添加批次大小和图像尺寸data.update({'bs': data['image0'].size(0),  # 批次大小'hw0_i': data['image0'].shape[2:],  # 第一张图像的高度和宽度'hw1_i': data['image1'].shape[2:]   # 第二张图像的高度和宽度})# 如果两张图像的尺寸相同，进行联合处理if data['hw0_i'] == data['hw1_i']:print("Concatenated Images Shape:", torch.cat([data['image0'], data['image1']], dim=0).shape) # ([2, 1, 480, 640])feats_c, feats_f = self.backbone(torch.cat([data['image0'], data['image1']], dim=0))print("Feature Coarse Shape:", feats_c.shape) # 1/8  ([2, 256, 60, 80])print("Feature Fine Shape:", feats_f.shape) # 1/2  ([2, 128, 240, 320])# 分离两张图像的特征(feat_c0, feat_c1), (feat_f0, feat_f1) = feats_c.split(data['bs']), feats_f.split(data['bs'])print("Feature Coarse Image 0 Shape:", feat_c0.shape) # ([1, 256, 60, 80])print("Feature Coarse Image 1 Shape:", feat_c1.shape) # ([1, 256, 60, 80])print("Feature Fine Image 0 Shape:", feat_f0.shape) # ([1, 128, 240, 320])print("Feature Fine Image 1 Shape:", feat_f1.shape) # ([1, 128, 240, 320])else:# 处理不同尺寸的图像(feat_c0, feat_f0), (feat_c1, feat_f1) = self.backbone(data['image0']), self.backbone(data['image1'])print("Feature Coarse Image 0 Shape:", feat_c0.shape)print("Feature Coarse Image 1 Shape:", feat_c1.shape)print("Feature Fine Image 0 Shape:", feat_f0.shape)print("Feature Fine Image 1 Shape:", feat_f1.shape)# 更新数据字典，添加特征图的尺寸data.update({'hw0_c': feat_c0.shape[2:],  # 第一张图像粗特征的高度和宽度'hw1_c': feat_c1.shape[2:],  # 第二张图像粗特征的高度和宽度'hw0_f': feat_f0.shape[2:],  # 第一张图像细特征的高度和宽度'hw1_f': feat_f1.shape[2:]   # 第二张图像细特征的高度和宽度})# 2. 粗级别局部变换模块# 添加位置编码信息feat_c0 = rearrange(self.pos_encoding(feat_c0), 'n c h w -> n (h w) c')print("Encoded Coarse Feature Image 0 Shape:", feat_c0.shape) # ([1, 4800, 256]) # 总共4800个点，每个点事256维向量feat_c1 = rearrange(self.pos_encoding(feat_c1), 'n c h w -> n (h w) c')print("Encoded Coarse Feature Image 1 Shape:", feat_c1.shape) # ([1, 4800, 256])mask_c0 = mask_c1 = None  # 掩码在训练中有用if 'mask0' in data:mask_c0, mask_c1 = data['mask0'].flatten(-2), data['mask1'].flatten(-2)# 进行粗级别局部变换feat_c0, feat_c1 = self.loftr_coarse(feat_c0, feat_c1, mask_c0, mask_c1) print("Transformed Coarse Feature Image 0 Shape:", feat_c0.shape) # ([1, 4800, 256]) # 计算后输出结果不变print("Transformed Coarse Feature Image 1 Shape:", feat_c1.shape) # ([1, 4800, 256]) # 计算后输出结果不变# 3. 粗级别匹配self.coarse_matching(feat_c0, feat_c1, data, mask_c0=mask_c0, mask_c1=mask_c1) # 4. 细级别细化feat_f0_unfold, feat_f1_unfold = self.fine_preprocess(feat_f0, feat_f1, feat_c0, feat_c1, data)if feat_f0_unfold.size(0) != 0:  # 至少有一个粗级别预测feat_f0_unfold, feat_f1_unfold = self.loftr_fine(feat_f0_unfold, feat_f1_unfold)# 5. 细级别匹配self.fine_matching(feat_f0_unfold, feat_f1_unfold, data)return x

Local Feature CNN

Local Feature CNN，就是一个blockbone，对两张图特征图进行卷积，分别得到原来特征图大小的1/8（用来进行粗粒度的匹配）和1/2（用来进行细粒度的匹配） 大小特征图。

# 来源 feats_c, feats_f = self.backbone(torch.cat([data['image0'], data['image1']], dim=0))
def build_backbone(config):if config['backbone_type'] == 'ResNetFPN':if config['resolution'] == (8, 2):return  ResNetFPN_8_2(config['resnetfpn'])# ...
class ResNetFPN_8_2(nn.Module):"""ResNet+FPN, output resolution are 1/8 and 1/2.Each block has 2 layers."""# 使用ResNet+FPN获得特征图1/8 和 1/2# 关于FPN，FPN（Feature Pyramid Networks）是一种用于目标检测和分割任务的神经网络架构，# 特别适用于处理多尺度的图像特征。它通过建立特征金字塔来捕捉不同尺度的图像信息，增强了模型对不同尺度目标的检测能力。def __init__(self, config):# ... 具体代码太长不展示

Coarse-Level Local Transformer

Coarse-Level Local Feature Transformer,是对两张1/8的特征图（FA，FB）进行多个串联的self-attention（自己的q和k向量做内积）和cross-attention（FA提供q，FB提供V做内积）来计算自己每个区域内之间的关系和两张图不同区域内的关系。

# 来源核心代码中  feat_c0, feat_c1 = self.loftr_coarse(feat_c0, feat_c1, mask_c0, mask_c1)
class LocalFeatureTransformer(nn.Module):def forward(self, feat0, feat1, mask0=None, mask1=None):"""前向传播函数：参数：feat0 (torch.Tensor): 特征图 0，形状为 [N, L, C]feat1 (torch.Tensor): 特征图 1，形状为 [N, S, C]mask0 (torch.Tensor): 特征图 0 的掩码，形状为 [N, L]（可选）mask1 (torch.Tensor): 特征图 1 的掩码，形状为 [N, S]（可选）"""# 确保特征图的通道数与 transformer 的模型维度一致assert self.d_model == feat0.size(2), "the feature number of src and transformer must be equal"# 迭代处理每一层，根据名称决定操作，总共4个self-cross，两个额为一组for layer, name in zip(self.layers, self.layer_names):  # ['self', 'cross', 'self', 'cross', 'self', 'cross', 'self', 'cross']if name == 'self':# 自注意力机制：feat0 本身计算 q 和 kfeat0 = layer(feat0, feat0, mask0, mask0)print("Self Attention Output feat0 Shape:", feat0.shape) # ([1, 4800, 256])# 自注意力机制：feat1 本身计算 q 和 kfeat1 = layer(feat1, feat1, mask1, mask1)print("Self Attention Output feat1 Shape:", feat1.shape) # ([1, 4800, 256])elif name == 'cross':# 交叉注意力机制：feat0 提供 q，feat1 提供 k 和 vfeat0 = layer(feat0, feat1, mask0, mask1)print("Cross Attention Output feat0 Shape:", feat0.shape) # ([1, 4800, 256])# 交叉注意力机制：feat1 提供 q，feat0 提供 k 和 vfeat1 = layer(feat1, feat0, mask1, mask0)print("Cross Attention Output feat1 Shape:", feat1.shape) # ([1, 4800, 256])else:raise KeyError("Unknown layer type")print("Final feat0 Shape:", feat0.shape) # ([1, 4800, 256]) 计算后向量维度是不变的print("Final feat1 Shape:", feat1.shape) # ([1, 4800, 256])return feat0, feat1

Matching Module

Matching Module，到目前位置都是进行粗粒度匹配，经过了多层的Coarse-Level Local Feature Transformer，FA，FB已经了解了自身的关系，同时也知道了和对方之间的关系，要进行关系的匹配了，采用互近邻(mutual nearest neighbor (MNN))的方式匹配，简单来说，假设FA对某个点信息概率值大， 但是同时也要某个点对A的概率值也大，必须双向的，源码阈值设置为0.2，也就是互相的结果都大于0.2才能匹配成功，筛选得到符合阈值点传入下一层。

# 来源核心代码中  self.coarse_matching(feat_c0, feat_c1, data, mask_c0=mask_c0, mask_c1=mask_c1) 
class CoarseMatching(nn.Module):def forward(self, feat_c0, feat_c1, data, mask_c0=None, mask_c1=None):"""前向传播函数：参数：feat_c0 (torch.Tensor): 特征图 0，形状为 [N, L, C]feat_c1 (torch.Tensor): 特征图 1，形状为 [N, S, C]data (dict): 额外数据字典mask_c0 (torch.Tensor): 特征图 0 的掩码，形状为 [N, L]（可选）mask_c1 (torch.Tensor): 特征图 1 的掩码，形状为 [N, S]（可选）"""# 提取批次大小、特征图 0 的长度、特征图 1 的长度和特征图的通道数N, L, S, C = feat_c0.size(0), feat_c0.size(1), feat_c1.size(1), feat_c0.size(2)print(feat_c0.shape) # ([1, 4800, 256])# 对特征图进行归一化feat_c0, feat_c1 = map(lambda feat: feat / feat.shape[-1]**.5, [feat_c0, feat_c1])if self.match_type == 'dual_softmax':# 计算相似度矩阵sim_matrix = torch.einsum("nlc,nsc->nls", feat_c0, feat_c1) / self.temperatureprint("Similarity Matrix Shape:", sim_matrix.shape) # ([1, 4800, 4800]) ,feat_c0中的4800和feat_c1进行计算得到([1, 4800, 4800])if mask_c0 is not None:# 应用掩码，将不需要的位置填充为负无穷大sim_matrix.masked_fill_(~(mask_c0[..., None] * mask_c1[:, None]).bool(),-float('inf'))# 计算置信度矩阵conf_matrix = F.softmax(sim_matrix, 1) * F.softmax(sim_matrix, 2)print("Confidence Matrix Shape:", conf_matrix.shape) # ([1, 4800, 4800])# ... 省略部分没走的代码# 更新数据字典data.update({'conf_matrix': conf_matrix})# 从置信度矩阵中预测粗略匹配data.update(**self.get_coarse_match(conf_matrix, data))

再进行下一步细粒度匹配之前，需要将粗粒度匹配出来结果和1/2特征图大小，进行处理，找出所有待细粒度处理的各个区域

# 来源核心代码中  feat_f0_unfold, feat_f1_unfold = self.fine_preprocess(feat_f0, feat_f1, feat_c0, feat_c1, data)
class FinePreprocess(nn.Module):def forward(self, feat_f0, feat_f1, feat_c0, feat_c1, data):"""前向传播函数：参数：feat_f0 (torch.Tensor): 特征图 0，形状为 [N, C, H, W]feat_f1 (torch.Tensor): 特征图 1，形状为 [N, C, H, W]feat_c0 (torch.Tensor): 粗级特征图 0，形状为 [N, L, C]feat_c1 (torch.Tensor): 粗级特征图 1，形状为 [N, L, C]data (dict): 包含额外信息的数据字典更新：data (dict): 更新数据字典，包括：'W' (int): 窗口大小"""# 获取窗口大小和步幅W = self.Wstride = data['hw0_f'][0] // data['hw0_c'][0]# 更新数据字典data.update({'W': W})# 如果没有有效的批次 ID，返回空张量if data['b_ids'].shape[0] == 0:feat0 = torch.empty(0, self.W**2, self.d_model_f, device=feat_f0.device)feat1 = torch.empty(0, self.W**2, self.d_model_f, device=feat_f0.device)return feat0, feat1# 1. 展开所有局部窗口# 使用 unfold 函数提取特征图中的局部窗口feat_f0_unfold = F.unfold(feat_f0, kernel_size=(W, W), stride=stride, padding=W//2)feat_f0_unfold = rearrange(feat_f0_unfold, 'n (c ww) l -> n l ww c', ww=W**2)feat_f1_unfold = F.unfold(feat_f1, kernel_size=(W, W), stride=stride, padding=W//2)feat_f1_unfold = rearrange(feat_f1_unfold, 'n (c ww) l -> n l ww c', ww=W**2)# 2. 仅选择预测的匹配# 根据数据中的批次 ID 和点 ID 选择特定的匹配特征feat_f0_unfold = feat_f0_unfold[data['b_ids'], data['i_ids']]  # [n, ww, cf]feat_f1_unfold = feat_f1_unfold[data['b_ids'], data['j_ids']]  # [n, ww, cf]# 选项：使用粗级别的 Loftr 特征作为上下文信息：连接和线性变换if self.cat_c_feat:# 从粗级别特征中选择窗口，并进行线性变换feat_c_win = self.down_proj(torch.cat([feat_c0[data['b_ids'], data['i_ids']],feat_c1[data['b_ids'], data['j_ids']]], 0))  # [2n, c]# 合并特征图feat_cf_win = self.merge_feat(torch.cat([torch.cat([feat_f0_unfold, feat_f1_unfold], 0),  # [2n, ww, cf]repeat(feat_c_win, 'n c -> n ww c', ww=W**2),  # [2n, ww, cf]], -1))# 将合并后的特征图分为两部分feat_f0_unfold, feat_f1_unfold = torch.chunk(feat_cf_win, 2, dim=0)return feat_f0_unfold, feat_f1_unfold

Coarse-to-Fine Module

Coarse-to-Fine Module，在经过粗粒度匹配之后，会得到一些候选点区域，这些候选点区域会传到细粒度，细粒度其实就是在匹配的小区域上再做一次Coarse-Level Local Feature Transformer，同样会互相了解对方的特征，但是这里不是采用MNN机制，假设FA的小区域中的某个点有FB对应小区域中的所有点匹配的概率值，才将这写概率值绘制成类似于热力图，再对整个图计算期望值，得到最终的匹配点坐标。

# 来源核心代码中 self.fine_matching(feat_f0_unfold, feat_f1_unfold, data)
class FineMatching(nn.Module):def forward(self, feat_f0, feat_f1, data):"""前向传播函数：参数：feat_f0 (torch.Tensor): 特征图 0，形状为 [M, WW, C]，其中 M 是匹配的数量，WW 是窗口大小的平方，C 是通道数feat_f1 (torch.Tensor): 特征图 1，形状为 [M, WW, C]data (dict): 包含额外信息的数据字典"""M, WW, C = feat_f0.shape  # 提取特征图的形状信息W = int(math.sqrt(WW))  # 计算窗口的边长scale = data['hw0_i'][0] / data['hw0_f'][0]  # 计算缩放因子self.M, self.W, self.WW, self.C, self.scale = M, W, WW, C, scale# 特殊情况处理：如果没有找到粗级别匹配if M == 0:assert not self.training, "在训练阶段，M 应始终大于0，请检查 coarse_matching.py"# logger.warning('在粗级别没有找到匹配。')data.update({'expec_f': torch.empty(0, 3, device=feat_f0.device),  # 返回空的期望位置和标准差'mkpts0_f': data['mkpts0_c'],  # 使用粗级别的关键点作为回退'mkpts1_f': data['mkpts1_c'],})return# 选择特征图中心的特征feat_f0_picked = feat_f0[:, WW//2, :]print(feat_f0_picked.shape) # ([1541, 128]) 全部待计算的特征区域个数# 计算特征之间的相似度矩阵sim_matrix = torch.einsum('mc,mrc->mr', feat_f0_picked, feat_f1)print(sim_matrix.shape) # ([1541, 128]) # 使用 softmax 函数计算热图softmax_temp = 1. / C**.5heatmap = torch.softmax(softmax_temp * sim_matrix, dim=1).view(-1, W, W)print(heatmap.shape) # ([1541, 5, 5]) # 1541 个 5x5的热力图矩阵# 从热图中计算坐标，计算期望值coords_normalized = dsnt.spatial_expectation2d(heatmap[None], True)[0]  # 计算归一化坐标print(coords_normalized.shape) #([1541, 2]) # 1541 个（x,y）坐标,都是小数，表示占图大小比例grid_normalized = create_meshgrid(W, W, True, heatmap.device).reshape(1, -1, 2)  # 创建网格print(grid_normalized.shape) #([1, 25, 2])# 计算标准差var = torch.sum(grid_normalized**2 * heatmap.view(-1, WW, 1), dim=1) - coords_normalized**2  # 计算方差std = torch.sum(torch.sqrt(torch.clamp(var, min=1e-10)), -1)  # 计算标准差，使用 clamp 以保证数值稳定性# 更新数据字典以用于精细级别的监督data.update({'expec_f': torch.cat([coords_normalized, std.unsqueeze(1)], -1)})# 计算绝对的关键点坐标,按图长宽比例还原位置坐标self.get_fine_match(coords_normalized, data)