文章目录
- 概要
- 介绍
- RAFT 的基础
- 特征提取
- 视觉相似性
- 相关性金字塔
- 相关性查找
- 高效的相关性查找(可选)
- 迭代更新
- 更新块
- 凸上采样
- 训练
- 如何使用RAFT
- 结论
- 参考
概要
在这篇文章中,我们将了解一种旗舰的光流深度学习方法,该方法获得了 2020 年 ECCV 最佳论文奖,并被引用超过 1000 次。它也是KITTI基准上许多表现最好的模型的基础。此模型称为 RAFT:Recurrent All-Pairs Field Transforms for Optical Flow,在 PyTorch 或 GitHub 上很容易获得。这些实现使其具有高度可访问性,但模型很复杂,理解它可能会令人困惑。在这篇文章中,我们将把 RAFT 分解成它的基本组件,并详细了解它们中的每一个。然后我们将学习如何在 Python 中使用它来估计光流。在第 2 部分中,我们将探索晦涩难懂的细节,并将不同的块可视化,以便我们对它们的工作原理有更深的理解。
介绍
光流是像素在图像序列中的表观运动。为了估计光流,场景中物体的运动必须具有相应的亮度位移。这意味着一张图像中移动的红球在下一张图像中应该具有相同的亮度和颜色,这使我们能够确定它在像素方面移动了多少。图 1 显示了一个光流示例,其中逆时针旋转的吊扇被一系列图像捕获。
最右边的彩色图像包含了从第 1 帧到第 2 帧每个像素的明显运动轨迹,并进行了颜色编码,不同的颜色表示像素运动的不同水平和垂直方向。这就是密集光流估计的一个例子。
密集光流估计为每个像素分配一个二维光流向量,描述其在时间间隔内的水平和垂直位移。在稀疏光流中,该向量只分配给与边角等强特征相对应的像素。 为了使流向量存在,像素在 t 时刻的强度必须与 t+1 时刻相同,这就是所谓的亮度一致性假设。 时间 t 时位置 (x,y) 处的图像强度或亮度由 I(x,y,t) 给出。 下面我们以已知像素位移为例(如图 2 所示)来直观地说明这一点,其中 dx 和 dy 是图像的水平和垂直位移,dt 是帧间的时间差。
亮度一致性假设意味着(x,y,t)处的像素在(x+dx, y+dy, t+dy)处具有相同的亮度。因此 I(x, y, t) = I(x+dx, y+dy, t+dt) .
根据亮度一致性假设,我们可以用关于 (x, y, t) 的 1ˢᵗ阶泰勒逼近法对右边进行扩展,从而得出光流方程 [ 1 ]。
水平梯度 Iₓ 和垂直梯度 Iᵧ 可用索贝尔算子近似,时间梯度 Iₜ 已知,因为我们有 t 和 t+1 时间的图像。流动方程有两个未知数 u 和 v,分别是时间 dt 上的水平位移和垂直位移。单个方程中的两个未知数使其成为一个未决问题,人们曾多次尝试求解 u 和 v。RAFT 是一种估算 u 和 v 的深度学习方法,但它实际上比根据两个框架预测流量更复杂。它是为精确估计光流场而精心设计的,下一节我们将深入探讨它的复杂细节。
RAFT 的基础
RAFT 是一种深度神经网络,能够在给定一对连续图像 I₁ 和 I₂ 的情况下估计密集光流。它能估算出流动位移场 ( f¹ , f² ) ,将 I₁ 中的每个像素 (u, v) 映射到 I₂ 中的相应像素 (u’, v’) ,其中 (u’, v’) = (u + f¹ (u), v + f² (v)) 。它的工作原理是提取特征,找到它们之间的相关性,然后以模仿优化算法的方式迭代更新流量。初始流量要么被初始化为全部为 0 的流量,要么被初始化为前向推算的先前流量估计值,这就是所谓的热启动。整体架构如下所示。
请注意它如何具有三个主要块:特征编码器块、视觉相似性块和迭代更新块。 RAFT 架构有两种规模,一种是具有 480 万个参数的大型架构,一种是具有 100 万个参数的小型架构。在这篇文章中,我们将重点关注大型架构,但是一旦我们了解了大型架构,了解小型架构就没有什么意义了。
特征提取
RAFT 使用由六个残差块组成的卷积神经网络 (CNN) 对两个输入图像执行特征提取,并使用 D 特征图将每个图像下采样至 1/8 分辨率。
特征编码器网络 g 对具有共享权重的两个图像进行操作,而上下文编码器网络 f 仅对 I₁ 进行操作并提取用作流估计的主要参考的特征。除了细微的差别外,两个网络的整体架构几乎相同。上下文网络使用批量归一化,而特征网络使用实例归一化,上下文网络提取 C = c + h 特征图,其中 c 是上下文特征图的数量,h 是将初始化隐藏特征图的隐藏特征图的数量。迭代更新块的状态。
注意:原始论文经常使用简写符号:HxW 来引用特征图尺寸 H/8xW/8。这可能会令人困惑,因此我们将遵循 H’ = H/8 的约定,使得特征图大小为 H’xW’。我们还将从 I₁ 中提取的特征图张量称为 s g 1 ,对于 I2 也是如此。
视觉相似性
视觉相似度是 4D H’xW’xH’xW’ 全对相关体积 C,通过取特征图的点积计算得出。
在相关性体积中,特征图 g¹ 中的每个像素都与特征图 g² 中的每个像素都有计算出的相关性,我们将这些相关性中的每一个称为 2D 响应图(见图 5)。在 4D 中思考可能具有挑战性,因此请想象将体积的前两个维度展平:(H’xW’)xH’xW’ ,我们现在有一个 3D 体积,其中 g¹ 的每个像素都有自己的 2D 响应图,显示它与 g² 的每个像素位置的相关性。由于特征是从图像派生的,因此响应映射实际上指示了给定的 I₁ 像素与 I₂ 的每个像素的相关性。
视觉相似度是一个全对相关卷,通过计算每个像素位置上每个特征图的相关性,将 I₁ 的像素与 I₂ 的每个像素联系起来。
相关性金字塔
相关体积有效地提供了小像素位移的信息,但可能难以捕获较大的位移。为了捕获大像素位移和小像素位移,需要多级相关性。为了解决这个问题,我们构建了一个相关性金字塔,其中包含多个级别的相关体积,其中不同层次的相关体积是通过平均汇集相关体积的最后两个维度来产生的。平均池化操作在体积的最后两个维度中产生粗略的 I₂ 相关特征,这使得 I₁ 的精细特征能够与 I₂ 逐渐粗略的特征相关联。 每个金字塔级别都包含越来越小的 2D 响应图。
图 5 显示了不同级别平均池化的不同 2D 响应图。相应相关体积的维度被堆叠成一个 5D 相关金字塔,其中包含四个具有核大小的级别:1、2、4 和 8。 金字塔提供了关于大位移和小位移的可靠信息,同时保持了相对于 I₁ 的高分辨率。
相关性查找
Correlation Lookup Operator L꜀ 通过对每个级别的相关性金字塔中的特征进行索引来生成新的特征图。给定当前的光流估计值 ( f¹ , f² ) ,I₁ : x = (u, v) 的每个像素都映射到它在 I₂ 中的估计对应关系 : x’ = (u + f¹(u) + v + f²(v)) 。我们在 x’ 周围定义一个局部邻域:
对应关系是像素在 I₂ 中的新位置,基于其流量估计值
所有金字塔级别的恒定半径意味着将在较低级别中纳入更大的上下文。 即半径为 4 对应于原始分辨率下的 256 像素。
在实践中,这个邻域是一个以每个高分辨率像素为中心的方形网格,当 r = 4 时,我们得到一个围绕每个像素的 9x9 网格,其中每个维度的长度为 (2r + 1) 。我们通过对网格定义的位置(边缘位置为零填充)的每个像素周围的相关特征进行双倍重采样来获得新的特征图。由于流量偏移和平均池化,邻域格网值可能是浮点,双线性重采样通过获取附近像素的 2x2 子邻域的加权平均值来轻松处理这个问题。换句话说,重采样将为我们提供亚像素精度。我们在金字塔的每一层的所有像素位置重新采样,这可以使用 PyTorch 的 F.grid_sample() 有效地完成。这些重新采样的特征被称为相关性特征,它们被输入到更新模块中。
高效的相关性查找(可选)
相关性查找以 O(N²) 进行缩放,其中 N 是像素数,这可能是大型图像的瓶颈,但有一个等效的操作与 O(NM) 缩放,其中 M 是金字塔级别的数。此操作将相关性金字塔与查找相结合,而操作则利用了内部积的线性和平均池化。2mx2m网格上的平均相关响应Cm(金字塔水平m)如下所示。
对于给定的金字塔级别 m,我们不需要对特征图 g1 求和,这意味着可以通过特征图 g1 与平均池化特征图 g² 的内积来计算相关性,其复杂度为 O (N) .由于这仅对单个金字塔级别 m 有效,因此我们必须计算每个级别的内积,使其按 O(M) 缩放,总复杂度为 O(NM) 。我们没有预先计算金字塔的相关性,而是只预先计算池化特征图,并在查找发生时按需计算相关性值。
迭代更新
更新算子估算一系列流量: {f₀ , f ₁ ,…, f ₙ } 从一个初始起点 f₀ 开始,这个起点可以是全部 0,也可以是前推的上一次流量估计值(热启动)。在每次迭代 k 中,它都会产生一个流量更新方向 Δf,并将其添加到当前估计值中:fₖ₊₁ = fₖ + Δfₖ。更新算子模仿优化算法,经过训练可提供更新,使估计的流量序列收敛到一个固定点:fₖ → f* 。
更新块
更新块将以下输入作为输入:相关特征、电流估计、上下文特征和隐藏特征。其架构与突出显示的子块如下所示。
更新块内的子块是:
- 特征提取块 — 从相关性、流和 I₁(上下文网络)中提取运动特征。
- 循环更新块 - 循环计算流更新
- Flow Head - 最终卷积层,将流量估计调整为 H/8 x W/8 x 2
如图 6 所示,循环更新块的输入是流、相关性和上下文特征的串联。潜在隐藏状态使用上下文网络中的隐藏特征进行初始化。 (上下文网络提取一堆 2D 特征图,然后将其分为上下文和隐藏特征图)。循环更新块由 2 个可分离的 ConvGRU 组成,可在不显着增加网络规模的情况下增加感受野。在每次更新时,来自循环更新块的隐藏状态被传递到流头以获得大小为 H/8 x W/8 x 2 的流估计。然后使用凸上采样对该估计进行上采样。
凸上采样
RAFT 的作者对双线性和凸上采样进行了实验,发现凸上采样可以显着提高性能。
凸上采样将每个精细像素估计为其相邻 3x3 粗像素网格的凸组合
让我们来分解一下凸上采样的工作原理,下面的图 8 提供了一个很好的视觉效果。
首先,我们假设精细分辨率像素是其最近的粗邻居的 3x3 网格的凸组合。该假设意味着粗像素的加权和必须等于真正的精细分辨率像素,并且限制为权重总和为一且非负。由于我们以八倍进行上采样,因此每个粗像素必须分解为 64 (8x8) 个精细像素(图 8 中的视觉效果未按比例绘制)。我们还注意到,3x3 网格中心的 64 个像素中的每一个都需要自己的权重集,因此所需的权重总数为:(H/8 x W/8 x (8x8x9))。
在实践中,权重通过神经网络进行参数化,凸上采样块使用两个卷积层来预测 (H/8 x W/8 x (8x8x9)) 掩模,然后对九个邻居的权重进行 softmax形状为 (H/8 x W/8 x (8x8)) 的面具。然后,我们使用此掩码获得邻域的加权组合并重塑以获得 HxWx2 流场。
训练
RAFT 的目标函数能够捕获所有迭代流预测。形式上,它是流预测和真值之间的加权 l1 距离之和,权重呈指数增加。
如何使用RAFT
我们可以使用 RAFT 来估计我们自己图像上的密集光流。首先,我们需要克隆 GitHub 存储库并下载模型。本教程的代码位于 GitHub 上。
接下来,我们将 RAFT 的核心添加到路径中
现在,我们需要一些辅助函数来连接 RAFT 类。注意:这些辅助函数仅针对 CUDA 编写,但您可以使用 Colab .NET 轻松访问 GPU。
注意 inference() 中的输入填充,我们需要确保所有图像都能被 8 整除。 raft.py 代码可以轻松地从命令行访问,但如果我们想与之交互,我们将需要重写其中一些,或者我们可以创建一个特殊的类来向它传递参数。
Args 类的默认初始化将直接与任何大型 RAFT 模型交互。为了演示 RAFT,我们将使用缓慢旋转的吊扇视频中的帧。现在我们可以加载模型并估计光流。
测试模式将返回 1/8 分辨率流和凸上采样流。
结论
在这篇文章中,我们了解了 RAFT,这是一种能够估计准确流场的高级模型。 RAFT 能够通过从提取的特征图中计算所有对相关量来捕获所有像素之间的关系。构建相关金字塔是为了捕获大像素位移和小像素位移。查找算子根据当前流量估计从相关金字塔中提取新的相关特征。更新块使用相关特征和当前流量估计来提供迭代更新,迭代更新收敛到最终流量估计,该最终流量估计通过凸上采样进行上采样。在第 2 部分中,我们将解压网络并了解一些关键块的工作原理。
参考
https://towardsdatascience.com/optical-flow-with-raft-part-1-f984b4a33993