RANSAC-Flow: generic two-stage image alignment（ECCV 2020）图像对齐论文代码详细分析

1. 前言

代码地址
论文地址

  图像对齐方法大致分为两类：

（1）参数化方法（非像素级对齐）
  通常做法是将一对图像输入网络，最后通过一个全连接层得到一个8维的特征向量，对应图像4个顶点的偏移量，从而得到一个单应性矩阵作为对齐变换，详情参考图像对齐深度学习算法调研。

（2）非参数化方法（像素级对齐）
  通常做法是将源图 $I_s$ 和目标图 $I_t$ 先提取特征得到 $F_s$ 与 $F_t$，然后对两幅特征图计算特征相似度，从而得到特征匹配关系，最终得到像素级的对齐变换。
  先解释一下像素级的对齐变换 $flow$ 是什么样的，$flow$ 的大小为 $[1,h_t,w_t,2]$，分辨率与 $I_t$ 相同，$2$ 代表一个坐标，即取 $I_s$ 中该坐标的像素值，坐标在 $[-1,1]$ 之间时在图像内部取值，否则取黑色（估计是对图像外部做 zeropad）。

  论文解决的问题：
  不难看出，非像素级对齐明显无法将各个细节也对齐的很好，而像素级对齐由于依赖计算特征相似度，计算全局计算量太大，计算局部的话又只适用于差异不太大的图像对。
  RANSAC-Flow 针对这两点把两种方法结合在一起，先弄个单应性变换把图像对齐的差不多，再计算像素级变换。

2. 对齐流程

  RANSAC-Flow 的创新主要在于这个对齐的流程，网络很朴素（没有花里胡哨的操作，也不是根据近期的模型改的，不需要看更多的论文和模型）。

  Stage 1：粗对齐（非像素级对齐）
（1）对 $I_s$ 和 $I_t$ 做 特征提取 + 特征匹配 得到若干对特征匹配点 $P$
（2）利用 RANSAC 得到第一个单应性矩阵 $H_1$
（3）通过 $H_1$ 粗对齐图像，得到 $I_s^{\prime }$ 和 $I_t$

  Stage 2：细对齐（像素级对齐）
（1）$I_s^{\prime }$ 和 $I_t$ 输入给细对齐网络，得到细对齐变换 $flow$ 和匹配度 $mask$
（2）$mask$ 大小为 $[1,1,h_t,w_t]$，用来度量每个像素对齐的效果
（3）通过 $mask$ 过滤得到 $P$ 中剩余还没有对齐好的点，回到 Stage 1 中的第二步得到 $H_2$
（4）反复迭代，将多个对齐变换组合到一起，得到最终完整的像素级对齐变换

3. 网络与流程细节

  粗对齐网络是用的现成的 ResNet-50 就不画图了，只参与测试，不参与训练；细对齐网络是自己训练的。下述流程主要是测试时的迭代流程，在训练的时候不需要这种迭代操作。

3.1 粗对齐网络

具体流程如下案例

Is [3, 853, 1280]
It [3, 1200, 1600]
特征提取网络为预训练好的 ResNet-50(ImageNet 或 moco)下采样16倍
1. 按 minSize = 480 和 7个尺度[2.0, 1.66, 1.33, 1.0, 0.83, 0.66, 0.5] 对 Is 提取特征
[3, 960, 1440] --> [1, 1024, 60, 90]
[3, 800, 1200] --> [1, 1024, 50, 75]
[3, 640,  960] --> [1, 1024, 40, 60]
[3, 480,  720] --> [1, 1024, 30, 45]
[3, 400,  592] --> [1, 1024, 25, 37]
[3, 320,  480] --> [1, 1024, 20, 30]
[3, 240,  352] --> [1, 1024, 15, 22]
将特征拉平并cat到一起得到 featA [1024, 14755]

2. 按 minSize = 480 对 It 提取特征
[3, 480, 640]  --> [1, 1024, 30, 40]
将特征拉平得到 featB [1024, 1200]

3. 计算特征相似度
两幅图像的特征两两做点积
score = torch.mm(featA.transpose(0, 1), featB)
得到 score [14755, 1200]

4. 得到特征匹配
若一对特征的点积在所在行和列中都是最大的, 就保留其索引 index1, index2
即 featA 中与 fb 点积最大的为 fa, 同时 featB 中与 fa 点积最大的为 fb
(index1 为 fa 在 featA 中索引, index2 为 fb 在 featB 中索引)

5. 生成粗对齐变换
对匹配的特征点对做 RANSAC 得到单应性矩阵, 并将其转化为像素级变换 flowCoarse

3.2 细对齐网络

细对齐网络分为4个部分
network = {
    
    'netFeatCoarse' : model.FeatureExtractor(), 
           'netCorr'       : model.CorrNeigh(args.kernelSize),
           'netFlowCoarse' : model.NetFlowCoarse(args.kernelSize), 
           'netMatch'      : model.NetMatchability(args.kernelSize),
           }

（1）netFeatCoarse: 提取特征

（2）netCorr: 计算特征相似度
  ZeroPad2d(3) 对应定义网络时的 kernelSize = 7，意味着 $F_a$ 中每个特征会与 $F_b$ 中对应位置的 $7\ast 7$ 范围内的特征计算相似度。

（3）netFlowCoarse: 计算细对齐变换
  想要得到对齐变换，需要得到与之匹配的特征的坐标。不难想到将其转为分类问题，通过 $softmax$ 分类，得到最相似的特征，但是如何可求导地获得对应特征的坐标？
  这里使用了两个固定的网格 gridX 与 gridY，代表 $7\ast 7$ 的 XY 坐标；用网格与分类得分相乘求和，就能近似得到匹配特征的相对坐标值（与自身坐标的偏移量），通过比例转化再加上原坐标 grid 便可得到细对齐变换。

（4）netMatch: 计算匹配度
  匹配度很自然的用 $sigmoid$ 做回归

3.3 迭代

  按匹配度 $match\ge 1$ 来判断粗对齐得到的匹配特征点集是否对齐好，剩余的生成新的粗对齐变换。
  不太明白为什么阈值取1，$0<sigmoid(x)<1$ ，不过看了一下匹配度矩阵数值，确实有不少等于1，很多是0.999…，当然也有较低的值，后面训练和损失的地方再分析。
  迭代停止判定：（1）迭代次数大于预设次数；（2）$match$ 有效位置的均值小于阈值 0.01，有效位置即去掉已经判定对齐好的剩余位置

3.4 融合对齐变换

  迭代后会得到多个 $flow$ 和 $match$，首先将 $flo{w}_1$ 作为最终的变换， $matc{h}_1<0.95$ 且 $matc{h}_i\ge 0.95$ 则取 $flo{w}_i$ 对应位置的变换取而代之。
  这种做法会使 $match$ 都 $<0.95$ 的位置默认用第一个得到的对齐变换。

4. 数据集

  文章涉及多个数据集和应用，这里主要介绍图像对齐里的 MegaDepth，文件预览：
…/data/MegaDepth
├── MegaDepth_Train/
├── MegaDepth_Train_Org/
├── Val/
 ├── img/
 ├── coarse.pkl
 └── corr.csv
└── Test/
 ├── test1600Pairs
 └── test1600Pairs.csv

  训练用的 MegaDepth_Train 内的图像，3幅图像一组，随机选两个作为训练数据，图像是大致对齐的，因此可以用来直接训练细对齐网络。训练数据预览如下：

  验证用的 Val 内的数据，corr.csv 内存储图像对路径和标注的匹配点坐标，coarse.pkl 内为图像对粗对齐的单应性矩阵参数。
  测试用的 Test 内的数据，与验证集类似，只是没有提供单应性变换了。

5. 训练

  数据预处理：大致上是做了缩放、裁剪和翻转，总之最后输入 $224\ast 224$ 大小的图像。

  粗对齐网络不参与训练，而细对齐网络所需要的功能就是能将大致对齐的两幅图像进一步像素级对齐，以及判断每个像素对齐的好不好，训练集也是初步对齐的图像对，所以上述的迭代对齐方法只是在测试进行，训练的时候并不需要。

5.1 损失函数

$$\mathcal{L}\left(I_s,I_t\right)={\mathcal{L}}_{rec}\left(I_s,I_t\right)+\lambda {\mathcal{L}}_m\left(I_s,I_t\right)+\mu {\mathcal{L}}_c\left(I_s,I_t\right)$$

损失函数分为三个部分：

（1）重建损失 Reconstruction loss：${\mathcal{L}}_{rec}$

$${\mathcal{L}}_{rec}^{SSIM}\left(I_s,I_t\right)=\sum _{(x,y)\in I_t}{M}_t^{cycle}(x,y)\left(1-\mathrm{SSIM}\left(I_s\left(x^{\prime },y^{\prime }\right),I_t(x,y)\right)\right)$$

$$M_t^{\text{cycle}}(x,y)=M_{t\to s}(x,y)M_{s\to t}\left(x^{\prime },y^{\prime }\right)$$

  这个损失是对齐的核心损失，利用最大化 SSIM 使两幅图像对齐，其中 $M_t^{\text{cycle}}(x,y)$ 是 $I_s$ 对齐 $I_t$ 的匹配度乘以 $I_t$ 对齐 $I_s$ 的匹配度，SSIM 的具体细节见 图像质量评估指标：PSNR / SSIM 原理及Python代码。其实就是让觉得对齐好的区域的差异变小。

（2）周期一致性损失 Cycle consistency loss：${\mathcal{L}}_c$

$${\mathcal{L}}_c\left(I_s,I_t\right)=\sum _{(x,y)\in I_t}{M}_t^{cycle}(x,y){\Vert \left(x^{\prime },y^{\prime }\right),{\mathbf{F}}_{t\to s}(x,y)\Vert }_2$$

  这个损失在对齐中经常出现，代码中的实现通常是对 $flo{w}_{s\to t}$ 做 $flo{w}_{t\to s}$ 的变换，所得新的变换 $flo{w}_{s\to t\to s}$ 趋向于不变的变换，也就是两张图的对齐是相互可逆的。

（3）匹配度损失 Matchability loss：${\mathcal{L}}_m$

$${\mathcal{L}}_m\left(I_s,I_t\right)=\sum _{(x,y)\in I_t}\left|M_t^{cycle}(x,y)-1\right|$$

  可以看出前两个损失都匹配度作为权重在计算损失，假如匹配度全为0，那么损失就最小了；为了避免这种情况，这边将匹配度趋向于1，并使用权重 $\lambda$ 平衡损失。

5.2 训练策略

  训练分为三个阶段，结合超参数分析：
（1）
--nEpochs 200 --lr 2e-4 --kernelSize 7 --imgSize 224 --batchSize 16 --margin 88
--lambda-match 0.0 --mu-cycle 0.0 --grad 0.0 --trainMode flow

（2）
--nEpochs 50 --lr 2e-4 --kernelSize 7 --imgSize 224 --batchSize 16 --margin 88
--lambda-match 0.0 --mu-cycle 1.0 --grad 0.0 --trainMode flow

（3）
--nEpochs 50 --lr 2e-4 --kernelSize 7 --imgSize 224 --batchSize 16 --margin 88
--lambda-match 0.01 --mu-cycle 1.0 --grad 0.0 --trainMode flow+match

  首先解释一下通用的（没有改变的）超参数：
--kernelSize 7 特征相似度时在 $7\ast 7$ 的范围计算，也就是在 $7\ast 7$ 的范围内寻找匹配的点
--margin 88 四周 88 范围的区域不计算损失，对应 224 大小的图像，即只计算中心 $48\ast 48$ 范围的损失
--lambda-match 0.0 对应 $\lambda {\mathcal{L}}_m\left(I_s,I_t\right)$ 的系数 $\lambda$
--mu-cycle 0.0 对应 $\mu {\mathcal{L}}_c\left(I_s,I_t\right)$ 的系数 $\mu$
--grad 0.0 是一个梯度损失函数的系数，看起来作用是让每个像素的对齐变换变得平滑（相对统一），论文中也没有用就不详细说明了

  下面解释三个训练阶段的差异：

（1）训练除 netMatch （计算匹配度）以外的网络，且只计算重建损失 ${\mathcal{L}}_{rec}$（算SSIM的）
--lambda-match 0.0 --mu-cycle 0.0 --grad 0.0 --trainMode flow

（2）训练除 netMatch （计算匹配度）以外的网络，加入周期一致性损失 ${\mathcal{L}}_c$（两个对齐一起不变的）
--lambda-match 0.0 --mu-cycle 1.0 --grad 0.0 --trainMode flow

（3）训练整个网络，加入匹配度损失 ${\mathcal{L}}_m$
--lambda-match 0.01 --mu-cycle 1.0 --grad 0.0 --trainMode flow+match

  根据个人实验结果，分析一下参数设定，由于显存不够 batchSize 设的 6，以下损失数值都是未乘系数的，且记录的是第一个 epoch 训练完时的数据到最后一个 epoch 的变化。

stage1：
${\mathcal{L}}_{rec}:0.7346\to 0.5205$，${\mathcal{L}}_c:0.0467\to 0.0537$
${\mathcal{L}}_c$ 没有参与训练，最开始是 0.0712 左右，第一个 epoch 以后便不太下降了

stage2：
${\mathcal{L}}_{rec}:0.5302\to 0.5267$，${\mathcal{L}}_c:0.0181\to 0.0104$
${\mathcal{L}}_c$ 系数为 1.0

stage3：
${\mathcal{L}}_{rec}:0.4282\to 0.4644$，${\mathcal{L}}_c:0.0032\to 0.0026$，${\mathcal{L}}_m:0.4498\to 0.3497$
${\mathcal{L}}_c$ 系数为 1.0，${\mathcal{L}}_m$ 系数为 0.01
可以发现 ${\mathcal{L}}_m$ 的数量级和 ${\mathcal{L}}_{rec}$ 相同，毕竟都是和1做差值的损失

以一个像素的损失做分析：
$\begin{array}{ll}l& =m\ast (1-ssim)+0.01\ast (1-m)\\ & =m-m\ast ssim+0.01-0.01\ast m\\ & =m\ast (0.99-ssim)+0.01\end{array}$

  理论上当 $ssim>0.99$ 时，匹配度会趋向于1，而 $ssim<0.99$ 时，匹配度会趋于0；也许这个系数可以当做 $ssim$ 的分水岭，高于 $1-\lambda$ 的才认为对齐好。
  但是，即使 $ssim<0.99$，网络也会同时升高 $ssim$ 和降低 $m$ 来降低 $loss$，而 $ssim$ 的梯度为 $m$，$m$ 的梯度为 $(0.99-ssim)$；因此，当 $ssim$ 和 $m$ 都较高时，网络会更倾向于降低 $ssim$ 来降低 $loss$；而且 $ssim$ 优先训练了 250个 epoch，能高的地方都差不多了，剩下的基本是实在对不齐的了，这样一来匹配度也不容易训练成全为0的解。

6. 评估

  论文在 KITTI 2015 和 Hpatches 上计算了密集对应的评估，在 RobotCar 和 MegaDepth 上计算了稀疏对应评估。
  密集对应的评价指标 AEE (Average Endpoint Error) 和 Fl-all (Ratio of pixels where flow estimate is wrong by both 3 pixels and ≥ 5%) 不太了解，暂时搁置。论文中表示光流误差主要在遮挡区域和图像边界，因为主要目的不是预测光流所以不关键，而且这两个数据集图像对差异比较小，不能突显优势，密集对应的标注不适合差异大的图像对，所以还是在稀疏对应上评估结果。

  稀疏对应的评价指标是根据标注的坐标对，计算与预测坐标的差距 $\mathrm{pixelDiff}=\sqrt{(x_g-x_p{)}^2+(y_g-y_p{)}^2}$，然后计算 $\mathrm{p}\mathrm{i}\mathrm{x}\mathrm{e}\mathrm{l}\mathrm{D}\mathrm{i}\mathrm{f}\mathrm{f}\le \mathrm{d}$ 的比例。
  表格中 Moco Feature 和 ImageNet Feature 指粗对齐迁移的模型，w/o Multi-H 指不使用迭代的方法（只用第一轮得到的对齐变换作为结果），w/o Fine-tuning 指不在对应数据集上微调细对齐模型。

7. 问题

（1）扭曲现象
  事实上对差异较大的图像对使用像素级对齐的方法会产生严重的扭曲现象，经常人眼都无法辨认图里是啥。

（2）评价指标
  目前没有什么好的评价指标来评价对齐效果，主要也是因为测试数据没法进行像素级标注，PSNR 和 SSIM 因为图像对本身的差异在这里也不太管用。个人感觉即使标注的稀疏点对的还行，但是整体图像存在扭曲难以辨识，是不能算作好的，但是又很难量化这个扭曲程度或者说图像符合人类视觉辨识习惯的程度。