1.前言
论文：https://arxiv.org/pdf/2109.07547.pdf
代码库：https://github.com/princeton-vl/RAFT-Stereo
1.1 什么是双目立体匹配
在标准的立体匹配设置中，输入包括两个帧：左帧和右帧。任务是估计输入图像之间的像素级位移图。在矫正过的立体视觉中，每个像素的位移被限制在水平线上。这个位移图，称为视差，可以与相机参数一起用来恢复深度。下面的第一节可看可不看，直接看方法部分也没问题。
1.2 常规操作
双目深度估计工作集中在问题的两个关键部分：(1) 特征匹配和 (2) 规则化。给定两个图像，特征匹配旨在计算一对图像块之间的匹配成本，常用的方法包括互信息、归一化交叉相关。规则化旨在恢复一个一致的深度图，受到如平滑性和平面性等先验的约束。
在立体视觉中，主要的方法是使用3D卷积神经网络。首先通过枚举整数视差构建一个3D成本体积，然后使用3D卷积网络来建模。然而，使用3D卷积处理成本体积会带来高昂的计算成本。
1.3 光流与双目立体匹配
光流和矫正过的立体视觉问题在本质上是密切相关的。在光流问题中，任务是预测一个像素级位移场，使得对于第一帧中的每个像素，我们都能估计它在第二帧中的对应位置。在矫正过的立体视觉中，任务相同，但我们有额外的约束，即x位移始终为正，且对应点位于水平线上,因此，y位移始终为0。
光流问题通常使用迭代细化方法来处理。RAFT展示了迭代细化可以完全在高分辨率下执行，并提出了一个简单的架构，在标准的光流基准测试中表现良好。RAFT首先从输入图像中提取特征，然后通过计算所有像素对之间的相关性来构建一个4D成本体积。最后，基于GRU的更新操作符使用来自相关性体积的特征迭代更新流场。
1.4 本文特别之处
RAFT-Stereo与以前的立体网络有实质性的不同。现有的工作通常依赖于3D卷积网络来处理立体成本体积。相比之下，RAFT-Stereo沿袭了RAFT，它仅使用2D卷积和使用单个矩阵乘法构建的轻量级成本体积。通过避免3D卷积的高计算和内存成本，RAFT-Stereo可以直接应用于百万像素图像，无需调整大小或分块处理图像。此外，通过使用迭代网络，我们可以轻松地用早期停止来交换准确性和效率。
2.实际方法
2.1 模型结构
特征编码部分用了两个网络，一个feature encoder, 一个context encoder。
feature encoder应用于左图和右图，将每个图像映射到特征图，然后再使用这些特征图构建相关性特征。feature encoder由一系列残差块和下采样层组成，产生的特征图分辨率为输入图像的1/4或1/8，具有256个通道,使用了实例归一化。
context encoder结构和feature encoder一样，只是把实例归一化换成了批归一化，并且只对左图用，之后这个信息是被用来注入到GRU中去的。
2.2 相关性特征构建
$${\bf C}_{i j k}\:=\:\sum_{h}f_{i j h}\cdot{\bf g}_{i k h},\:\:\:{\bf C}\in\mathbb{R}^{H\times W\times{W}}$$
其中$\mathbf{f},\mathbf{g}_{_{}}\in\mathbb{R}^{H\times W\times D}$,也就是左右图提到的特征。想象一下两张图，现在行与行有联系，那就直接把左图中的像素与右图中同一行的所有其他像素特征进行相乘求相似度。
但是这样其实是有问题的，因为相同空间点在左右图中的成像位置，一定是左图的x大于右图的x坐标，也就是视差大于0。所以这里与右图整行比其实有浪费，因为之用往左边比就行了。比如在左图中位置为（3,1），那只用比较右图的（0,1），（1，1），（2,1）就行了。但是这样整体算简单，作者也就这么干了，乘法运算应该耗时也还好，浪费就浪费吧。
作者还对这个相关性特征构建了一个特征金字塔，这和图像金字塔做法一样，就是用平均池化来构建的。
$${\bf C^{k+1}}\in\mathbb{R}^{H\times W\times{W/2^k}}$$
金字塔的优点就是希望在视差层面既看到全部，又看到细节。
2.3 查询操作
前面我们构建了一个相关性特征，现在想一想这么个事情，就是我告诉你视差是多少了，那你把相关性特征相应位置处的内容提出来。
为了简单起见，就看一行吧，比如我现在计算告诉你视差在红色点处，那我每层金字塔都取（2r-1）个点, 也就是黑点的地方,最后这处的特征是：金字塔层数(2*r-1)个点concate在一起。因为网络算出来的视差是小数，所以还得双线性插值，代码如下。
我来解释一下这里的目的是啥吧：其实网络回归一次能够得到红色点，但是红色点不够准确，接着给定红色点周围的信息，相当于缩小范围后让网络重新回归，希望得到更加准确的位置。那初始给定的视差应该为0，不然网络一开始随机取得太离谱，即使左右有个2r可能也覆盖不到正确区域。
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
|
def bilinear_sampler(img, coords, mode='bilinear', mask=False):
""" Wrapper for grid_sample, uses pixel coordinates """
H, W = img.shape[-2:]
xgrid, ygrid = coords.split([1,1], dim=-1)
xgrid = 2*xgrid/(W-1) - 1
if H > 1:
ygrid = 2*ygrid/(H-1) - 1
grid = torch.cat([xgrid, ygrid], dim=-1)
img = F.grid_sample(img, grid, align_corners=True)
if mask:
mask = (xgrid > -1) & (ygrid > -1) & (xgrid < 1) & (ygrid < 1)
return img, mask.float()
return img
out_pyramid = []
for i in range(self.num_levels):
corr = self.corr_pyramid[i]
dx = torch.linspace(-r, r, 2*r+1)
dx = dx.view(2*r+1, 1).to(coords.device)
x0 = dx + coords.reshape(batch*h1*w1, 1, 1, 1) / 2**i
y0 = torch.zeros_like(x0)
coords_lvl = torch.cat([x0,y0], dim=-1)
corr = bilinear_sampler(corr, coords_lvl)
corr = corr.view(batch, h1, w1, -1)
out_pyramid.append(corr)
out = torch.cat(out_pyramid, dim=-1)
|
2.4 多级别更新操作符
假如预测出了一系列视差$d_1,d_2,…,d_N$。每次迭代的时候，使用当前估计出来的视差按照上面说的做法去相关性特征$C^{k}$里面来索引特征。
这些特征结合当前估计的视差会经过两层卷积层，之后再结合context特征，注入到GRU中，GRU来更新隐藏状态，新的隐藏状态又被用来更新视差。
为了提高算法的性能，RAFT-Stereo使用了多个不同分辨率的GRU，这样可以同时处理不同尺度的图像信息，从而更好地处理图像中的大范围无纹理区域。这些GRU之间通过共享隐藏状态来相互连接，以便在不同分辨率之间传递信息。
如果不知道什么是GRU,可以在我的博客里检索一下。我这里简单概括一下，对于迭代输入，GRU能够提取每次输入里的精髓信息，把信息存在自己的隐藏状态里，还是看代码吧。
先补一下他们的GRU实现吧，代码如下：
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
|
class ConvGRU(nn.Module):
def __init__(self, hidden_dim, input_dim, kernel_size=3):
super(ConvGRU, self).__init__()
self.convz = nn.Conv2d(hidden_dim+input_dim, hidden_dim, kernel_size, padding=kernel_size//2)
self.convr = nn.Conv2d(hidden_dim+input_dim, hidden_dim, kernel_size, padding=kernel_size//2)
self.convq = nn.Conv2d(hidden_dim+input_dim, hidden_dim, kernel_size, padding=kernel_size//2)
def forward(self, h, cz, cr, cq, *x_list):
x = torch.cat(x_list, dim=1)
hx = torch.cat([h, x], dim=1)
z = torch.sigmoid(self.convz(hx) + cz)
r = torch.sigmoid(self.convr(hx) + cr)
q = torch.tanh(self.convq(torch.cat([r*h, x], dim=1)) + cq)
h = (1-z) * h + z * q
return h
|
下面代码里inp是需要注入的context信息，net就是一个列表，里面是上一次GRU输出的值，也就是上一轮的隐藏状态。看每个gru，第1个参数就是上一轮的隐藏状态，第2，3，4个参数是注入信息，第4个参数是其他分辨率的隐藏状态，只有在gru08这一层，第4个参数才是视差相关的信息。最后视差由net[0]来解码。
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
|
class BasicMultiUpdateBlock(nn.Module):
def __init__(self, args, hidden_dims=[]):
super().__init__()
self.args = args
self.encoder = BasicMotionEncoder(args)
encoder_output_dim = 128
self.gru08 = ConvGRU(hidden_dims[2], encoder_output_dim + hidden_dims[1] * (args.n_gru_layers > 1))
self.gru16 = ConvGRU(hidden_dims[1], hidden_dims[0] * (args.n_gru_layers == 3) + hidden_dims[2])
self.gru32 = ConvGRU(hidden_dims[0], hidden_dims[1])
self.flow_head = FlowHead(hidden_dims[2], hidden_dim=256, output_dim=2)
factor = 2**self.args.n_downsample
self.mask = nn.Sequential(
nn.Conv2d(hidden_dims[2], 256, 3, padding=1),
nn.ReLU(inplace=True),
nn.Conv2d(256, (factor**2)*9, 1, padding=0))
def forward(self, net, inp, corr=None, flow=None, iter08=True, iter16=True, iter32=True, update=True):
if iter32:
net[2] = self.gru32(net[2], *(inp[2]), pool2x(net[1]))
if iter16:
if self.args.n_gru_layers > 2:
net[1] = self.gru16(net[1], *(inp[1]), pool2x(net[0]), interp(net[2], net[1]))
else:
net[1] = self.gru16(net[1], *(inp[1]), pool2x(net[0]))
if iter08:
motion_features = self.encoder(flow, corr)
if self.args.n_gru_layers > 1:
net[0] = self.gru08(net[0], *(inp[0]), motion_features, interp(net[1], net[0]))
else:
net[0] = self.gru08(net[0], *(inp[0]), motion_features)
if not update:
return net
delta_flow = self.flow_head(net[0])
# scale mask to balence gradients
mask = .25 * self.mask(net[0])
return net, mask, delta_flow
|
最后，为了从低分辨率的视差场生成高分辨率的视差图，算法用的就是简单上采样。
2.5 Slow-FastGRU
看上面的GRU图，更新1/8分辨率的隐藏状态大约需要比更新1/16分辨率的隐藏状态多4倍的FLOPs。为了利用这一事实来加快推理速度，作者训练了RAFT-Stereo的一个版本，在这个版本中，他们多次更新1/16和1/32分辨率的隐藏状态，而每次只更新一次1/8分辨率的隐藏状态。这样做的目的是减少高分辨率更新的频率，从而减少总体的计算量。
在KITTI数据集上，使用32次GRU更新，这种简单的改变将RAFT-Stereo的运行时间从0.132秒减少到0.05秒，降低了52%。模型的推理速度得到了显著提升。
3.总结
这个方法亲测效果不错，而且他们的GRU循环结构能够在TensorRT上大幅度加速，记得是3倍提速。如果你想在实际中使用，记得使用middlebury数据finetune的版本，这个版本对于真实场景效果更好。
