光流估计中cost volume详解

原创声明：是暮涯啊 https://blog.csdn.net/longshaonihaoa/article/details/124726727

MEMC系列文章：
运动估计运动补偿（Motion estimation and motion compensation，MEMC）入门总结
深度学习MEMC插帧论文列表paper list
光流估计中cost volume详解
插帧中grid_sample函数详解

我认为之所以光流估计中会使用独有的cost volume，其初始形态是传统基于块的光流估计方法中，对每个前一帧图像F1中的某个块B1，计算其在一定范围内对应于后一帧图像2中哪个块B2最接近。假如B1的坐标为（x1，y1）B2的坐标为（x2，y2），对应这个块的光流矢量就是（x2-x1，y2-y1）。

1、FlowNet中的correlation

1.1 介绍

FlowNet作为光流估计的开山鼻祖，对correlation的计算花了大量篇幅。作者说这些分割啊，深度估计啊这些网络可以产生像素级预测，这说明NN用于得到像素级稠密光流估计是有可能的。那怎么涉及网络呢？一种是直接把前后帧图像F1F2堆叠后送到网络中，同时用GT的光流作为监督，让网络自己提取motion信息，也就是FlowNetSimple的结构。
原则上网络够大就行，但实际不确定行不行，所以退一步设计一个不那么通用，但效果好的网络。就先提取F1F2各自的有意义的特征，后来在将它们融合在一起。这大致类似于标准匹配方法，即首先从两个图像的块中提取特征，然后比较这些特征向量。接着就提出了灵魂疑问，怎么让网络计算特征间的相关性呢？
给定特征f1 f2，尺寸为WHC。correlation层计算f1的每个patch和f2的每个patch，现在只考虑f1中中心坐标为x1和f2中中心坐标为x2的两个patch之间的相关性。
patch边长为K:=2k+1，该公式的定义就是一个卷积操作，只不过正常卷积使用filter卷data，这个使用其他data卷data。计算一个c(x1,x2)需要$c{K}^2$次乘操作（从这个计算量来看是per-channel操作），这还是一个点，如果计算全局就要$W^2{H}^2$次这样的操作。因此作者引入搜索范围和步长。对f2的搜索范围界定在D:=2d+1，x1和x2的步长分别为s1，s2。（这里其实就是传统快搜索里将全局搜索转为局部搜索），最后生成$[WH{D}^2]$ 的相关性矩阵。$D^2$意味着将patch间2D的相关性堆叠到一个维度。

1.2 代码实现

from spatial_correlation_sampler import SpatialCorrelationSampler
函数详细介绍可以看这里
更深一点想看c++的原始代码可以点这里 ，还是蛮标准的循环取点相乘再取平均。

2 PWCNet 里的cost Volume

2.1 介绍

PWCNet其实和SpyNet很像，都是金字塔，warp。但是PWCnet不同的地方有两个1）warp是对特征，2）就是我们这次介绍的主角cost volume。摘要对cost volume的介绍是：用warped 的特征和第一张图的特征计算构造一个cost volume，它被一个CNN处理用来估计光流。
相关工作中说 cost volume用于存储 像素与其相关联的下一帧的对应像素 的 数据匹配成本。在介绍网络结构时，作者使用特征计算匹配代价（matching cost），定义它为第一张图的特征和第二张图warp后的特征间的相关性（correlation）
T为转置操作，N为列向量$c_1^l(x_1)$的长度，上标l表示金字塔的第l层。作者认为对L层的金字塔，只需要计算很小的D个像素内的cost volume就好了，因为top层的一个像素的motion对应到最大尺度层就是$2^{L-1}$个像素了，因此可以把D设的很小。输出的结果是$[D^2{H}^l{W}^l]$，后两个表示l层feature map的宽高。
其实这里看计算和FlowNet的计算基本是相同的，都是对f1中的每个坐标（x,y）对应的特征向量与f2中坐标（x，y）周围范围d的特征向量计算匹配代价。区别在于FlowNet是逐通道计算两个patch间距离。PWCNet是逐WH计算特征向量间距离。这个向量间转置点乘的距离就是余弦相似度。如果FlowNet的步长=D，那理论上两者是相同的。
知乎有个对PWCnet中cost volume更仔细的介绍

2.2 代码实现

torch代码，代码来源

    def corr(self, refimg_fea, targetimg_fea):maxdisp=4b,c,h,w = refimg_fea.shape# 通过F.unfold取出f2的窗口范围d=2*maxdisp+1内的特征向量，# 这里由于F.unfold的步长默认为1，所以每个WH都能提取出周围d*d的特征向量，# 因此可以view到（b,c,2*maxdisp+1, 2*maxdisp+1,h,w）# ps：个人认为第二个 2*maxdisp+1**2 的 次方是 个错误，应也是2*maxdisp+1，但不影响targetimg_fea = F.unfold(targetimg_fea, (2*maxdisp+1,2*maxdisp+1), padding=maxdisp).view(b,c,2*maxdisp+1, 2*maxdisp+1**2,h,w)# 对f1扩维， ps： targetimg_fea后的view应该无效cost = refimg_fea.view(b,c,h,w)[:,:,np.newaxis, np.newaxis]*targetimg_fea.view(b,c,2*maxdisp+1, 2*maxdisp+1**2,h,w)cost = cost.sum(1)b, ph, pw, h, w = cost.size()cost = cost.view(b, ph * pw, h, w)/refimg_fea.size(1)return cost

下面这个代码虽然和上面一样，但我觉得更艺术，就贴下来了

def compute_cost_volume(feat1, feat2, param_dict):"""only implemented for:kernel_size = 1stride1 = 1stride2 = 1"""max_disp = param_dict["max_disp"]_, _, height, width = feat1.size()num_shifts = 2 * max_disp + 1feat2_padded = tf.pad(feat2, (max_disp, max_disp, max_disp, max_disp), "constant", 0)cost_list = []for i in range(num_shifts):for j in range(num_shifts):corr = torch.mean(feat1 * feat2_padded[:, :, i:(height + i), j:(width + j)], axis=1, keepdims=True)cost_list.append(corr)cost_volume = torch.cat(cost_list, axis=1)return cost_volume

3 LiteFlowNet中的correction

文章对这部分的介绍很少，只是说I1 I2间的点相关性通过高层特征向量间相关性计算得到。
F1 F2 标傲世金字塔特征，d限制搜索范围，N表示特征向量长度。最后的代价矩阵c也会被集成到一个3D grid中。（其实和PWCNet的计算应该是相同的）

作者通过三个操作减少计算负担。1）每个特征level计算短距离匹配， 也就是限制d的大小，和PWC相同。2）通过f-warp拉进F1F2间特征空间距离。3）在高空间分辨率的level，只在采样点计算得到稀疏cost volume。然后对稀疏的cost volume进行空间方向的插值

4 MaskFlownet 中的cost volume

本文是对PWCnet的改进，核心就是一句话：做匹配代价容量计算时，应当排除那些在一帧中被遮挡的像素。
实现在于通过卷积预测mask后，将sigmoid（mask）与warp后的光流相乘。同时会加一个平衡项。该操作图示如下：

该操作对金字塔每层都实现一次，下面以第5层为例，代码实现如下：

warp5 = F.broadcast_mul(warp5, F.sigmoid(mask5)) + self.conv5f(tradeoff5)
warp5 = self.leakyRELU(warp5)
corr5 = self.corr(F, c15, warp5) 
corr5 = self.leakyRELU(corr5)

其中的相关性计算self.corr通过 F.Correlation 函数实现。但这个F是MXnet框架下的，详细介绍可以查看文档
文档中的相关性计算公式和FlowNet的完全相同。

5 RAFT 中 multi-scale 4D correlation volume

作者希望使用4D 金字塔 相关性 volume来计算视觉相似度（Visual Similarity）。其实现也更加简单，对图1的特征f1=[B C HW] 和 图2的特征f2=[B C H W] 直接点乘。也可以说就是将前面PWCnet的d设置为0，对应D=1，也就是只计算当前位置。 对应的f1中每个c与f2的HW个c相乘。
代码实现如下：

def corr(fmap1, fmap2):batch, dim, ht, wd = fmap1.shapefmap1 = fmap1.view(batch, dim, ht*wd)fmap2 = fmap2.view(batch, dim, ht*wd) # fimap转换维度后为（b，ht*wd，dim）*（b,dim,ht*wd）= (b,ht*wd, ht*wd)# 论文说的4D correlation volume 就是[h w h w]   corr = torch.matmul(fmap1.transpose(1,2), fmap2)corr = corr.view(batch, ht, wd, 1, ht, wd)return corr  / torch.sqrt(torch.tensor(dim).float())

6 总结

总的来说谈谈为什么计算光流的网络会有这个组件，我觉得一个是从传统方法延续过来的，可以采用coarse2fine的过程中逐渐指导下一层光流计算。另一个是说假如得到的d*d某个通道数值都很大，说明需要再往那个通道对应的运动向量移动。相当于嵌入了位置编码信息。但cost volume并非是计算光流所必需的，比如很早的SpyNet就没有使用该模块，毕竟计算cost volume的计算开销还是很大的。

以上如有错误请多多交流指正，创作不意，有用的话可以多多点赞收藏。