7、关于LoFTR
7、关于LoFTR
LoFTR论文链接：LoFTR
LoFTR的提出，是将Transformer模型的注意力机制在特征匹配方向的应用，Transformer的提取特征的机制，在自身进行，本文提出可以的两张图像之间进行特征计算，非常适合进行特征匹配。
一、传统匹配模式的局限
- 首先找到一些关键点(图像梯度较高，角点检测等)，然后计算特征相似度来匹配，这种方法很依赖检测到的特征点，一旦点找不到，那就不用说匹配了
- 对于位置不同的两个点，如果它们的背景特征相似(与位置无关了)，也无法匹配，两者差别不大的，就无法区分了
二、LoFTR解决方案
- 不需要先得到特征点，第一个问题就得到了解决
- 采用end2end（给一个输入，得到结果）方法，使用非常方便
三、模型架构图
模型进行特征匹配主要分为四步：
- Local Feature CNN，就是一个blockbone，对两张图特征图进行卷积，分别得到原来特征图大小的1/8（用来进行粗粒度的匹配）和1/2（用来进行细粒度的匹配） 大小特征图。
- Coarse-Level Local Feature Transformer,是对两张1/8的特征图（FA，FB）进行多个串联的self-attention（自己的q和k向量做内积）和cross-attention（FA提供q，FB提供V做内积）来计算自己每个区域内之间的关系和两张图不同区域内的关系。
- Matching Module，到目前位置都是进行粗粒度匹配，经过了多层的Coarse-Level Local Feature Transformer，FA，FB已经了解了自身的关系，同时也知道了和对方之间的关系，要进行关系的匹配了，采用互近邻(mutual nearest neighbor (MNN))的方式匹配，简单来说，假设FA对某个点信息概率值大， 但是同时也要某个点对A的概率值也大，必须双向的，源码阈值设置为0.2，也就是互相的结果都大于0.2才能匹配成功，筛选得到符合阈值点传入下一层。
- 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 和 k
feat0 = layer(feat0, feat0, mask0, mask0)
print("Self Attention Output feat0 Shape:", feat0.shape) # ([1, 4800, 256])
# 自注意力机制：feat1 本身计算 q 和 k
feat1 = layer(feat1, feat1, mask1, mask1)
print("Self Attention Output feat1 Shape:", feat1.shape) # ([1, 4800, 256])
elif name == 'cross':
# 交叉注意力机制：feat0 提供 q，feat1 提供 k 和 v
feat0 = layer(feat0, feat1, mask0, mask1)
print("Cross Attention Output feat0 Shape:", feat0.shape) # ([1, 4800, 256])
# 交叉注意力机制：feat1 提供 q，feat0 提供 k 和 v
feat1 = 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.temperature
print("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.W
stride = 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**.5
heatmap = 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)