RAFT：深度学习光流估计新范式

scott198512

已于 2024-05-28 00:21:07 修改

阅读量337

点赞数 3

分类专栏：深度学习文章标签：深度学习人工智能

于 2024-05-14 14:28:46 首次发布

原文链接：https://learnopencv.com/optical-flow-using-deep-learning-raft/

版权

深度学习专栏收录该内容

15 篇文章 3 订阅

订阅专栏

0.概述

在这篇文章中，我们将讨论两种基于深度学习的光流运动估计方法。FlowNet是第一个用于计算光流的CNN方法，RAFT是当前最先进的估计光流的方法。在本文中，我们还将看到如何使用作者提供的经过训练的模型来使用PyTorch对新数据进行推理。

1. Optical Flow task

在上一篇文章中，我们讨论了估计光流的基本组件和一些方法。需要提醒的是，光流任务包括估计两个连续帧之间的逐个像素运动。我们的目标是找到稀疏特征集或所有图像像素的位移，以计算它们的运动矢量。在上一篇文章中，我们回顾了OpenCV的一个重要理论和实际解决方案，现在我们将用深度学习方法继续这个主题来估计光流。

2. FlowNet

FlowNet architecture 于2015年推出，是第一个预测光流的CNN方法。作者受到CNN架构在分类、深度估计和语义分割任务中的成功结果的启发。随着深度学习方法和CNN成为解决许多计算机视觉任务的有利策略，作者又引入了两个用于光流估计的神经网络。

2.1 FlowNet Architecture

FlowNetS和FlowNetCorr架构都包含类似于U-Net架构的编码器和解码器部分。编码器从两个连续的图像中提取特征，而解码器升级编码器特征图并获得最终的光流预测。下面的内容将更深入地剖析FlowNetS和FlowNetCorr网络。

2.1.1 FlowNetS encoder

流网络（也称为FlowNetSimple）中的输入数据是两个连续帧的串联。这两幅图像被放置到6通道张量中，其中前三个通道属于第一幅图像，其余三个通道属于第二幅图像。编码器部分由几个卷积层组成，后面是激活函数。这种架构允许网络自己决定如何处理两个堆叠的图像，并为以下结果细化提供特征图。

2.1.2 FlowNetCorr encoder

FlowNetCorr 体系结构一次只接受一帧作为输入，因此图像不会在此处堆叠。在这个网络中，作者使用具有共享权重的CNN的第一阶段分别从两个连续的图像中提取特征。下一步是分别合并第一帧和第二帧中的两个计算特征图 $\mathbf f_1, \mathbf f_2$ .

为了提供这一点，作者引入了一种称为关联层的新技术，这是 FlowNetS 和 FlowNetCorr 之间的主要区别。相关性的计算等同于卷积运算，但这里没有可训练的核——一个特征图与另一个特征图卷积。因此，该层在两个特征图 $\mathbf f_1, \mathbf f_2$

之间执行乘法补丁比较，并且没有可训练的权重。

以第一和第二特征图为中心的

$x_1，x_2$

两个正方形斑块的相关公式分别定义为：

$corr（x_1， x_2） = \！\！\！\！\！\！\sum \limits_{o \in [-k， k]\times [-k， k]} \！\！\！\！\！\left \langle \mathbf{f}_1（x_1 + o）， \mathbf{f}_1（x_2 + o） \right \rangle$

方形块大小定义为 $\boldsymbol K =2k+1$

值得一提的是，作者不会在两个特征图之间做完全的关联，而是按补丁的方式进行。

在完成两个特征图的匹配过程后，关联结果将前向传递到其余的卷积层以提取高级特征。下图显示了 FlowNetCorr 的体系结构：

2.1.3 FlowNetS 和 FlowNetCorr decoder

编码器的输出特征图比输入图像分辨率小 64 倍，因此我们需要放大结果。对于这两种架构，此处的策略是相等的。可训练的上卷积层用于解码器阶段，以放大编码器光流输出。每个解码器阶段将上一阶段的放大结果与编码器相应层的特征图连接起来。使用来自编码器的数据有助于预测精细细节，就像在 U-Net 中所做的那样。作者只使用了四个升级阶段，因为使用更多的阶段可以略微提高质量。要将输出放大到初始图像分辨率，最好使用计算成本较低的双线性上采样。

网络输出是一个具有两个通道的张量，第一个通道包括 x 轴每个像素的运动位移，第二个通道包括 y 轴的运动位移。因此，我们为图像中的每个像素都有一个运动矢量。

2.2 FlowNet 损失函数

对于 FlowNet 来说，一个好的惩罚策略是多尺度训练损失。最后的预测是放大的小尺寸光流，因此小的预测对下一个大尺寸的预测有很大的影响。使用影响递减参数计算每个解码器级的损耗，以获得每个解码器级的精细预测。因此，小规模预测的计算损失对总损失的贡献更大，而不是对较大的损失的贡献。

每个解码器阶段的损失公式基于预测

$V_{pred} = （\Delta x_{pred}， \Delta y_{pred}）$

和地面实况数据

$V_{gt}=（\Delta x_{gt}， \Delta y_{gt}）$

之间的端点误差：

$L = ||V_{gt}-V_{calc}||_{1} = \sqrt{（\Delta x_{gt}-\Delta x_{calc}）^2 + （\Delta y_{gt}-\Delta y_{calc}）^2}$

3. RAFT

根据 SINTEL 基准测试，目前最先进的方法是 2020 年推出的 CNN 和 RNN 架构
的组合。这种新方法称为 RAFT – Recurrent All-Pairs Field Transforms。与之前的架构一样，它也有两种不同的类型——RAFT 和 RAFT-S。这两种架构有一个共同点，但 RAFT-S 是 RAFT 的轻量级版本。在这里，我们将看一下 RAFT 架构的基本组件和思想，然后将其与 RAFT-S 进行比较。

3.1 RAFT Architecture

根据论文，RAFT可以分为三个阶段：

特征提取器。网络输入由两个连续帧组成。为了从这两张图像中提取特征，作者使用了两个具有共享权重的CNN。这种方法类似于 FlowNetCorr 架构，在 FlowNetCorr 体系结构中，我们分别从两个图像中提取特征。CNN 的架构由 6 个残差层组成，就像 ResNet 的层一样，随着频道数量的增加，每第二层的分辨率就会降低一半。在这里我们可以看到 RAFT 编码器结构：

此外，Context 网络使用相同的 CNN 架构，它仅从第一张图像生成特征。归一化方法只有一个区别——特征提取器使用实例归一化，而上下文网络使用批量归一化。上下文网络的特征图稍后将在递归块中使用。
视觉相似性。视觉相似度计算为所有特征图对的内积。因此，我们将有一个名为 Correlation volumes 的 4D 张量，它提供有关小像素位移和大像素位移的关键信息。此方法不应与 FlowNetCorr 中的关联层混淆。在 FlowNetCorr 中，我们使用补丁相关性，而在 RAFT 中，我们计算两个特征图的全对相关性，没有任何固定大小的窗口。为了清楚起见，两个特征图
$\mathbf{f_1， f_2}$
之间的相关性计算如下：
$C_{ijkl} = \sum \limits_{d} \mathbf{f_1}_{ijd} \cdot \mathbf{f_2}_{kld}$

之后，通过将该 4D 张量的最后两个维度与大小为 1、2、4、8 的内核合并来构建 4 层相关金字塔。前三层的二维切片如下图所示：

相关金字塔用于创建多尺度图像相似性特征，使突然的运动更加明显。因此，金字塔提供了有关小位移和大位移的信息。
迭代更新。迭代更新是一系列门控循环单元（GRU）单元，它们结合了我们之前计算的所有数据。GRU 单元模拟迭代优化算法，但有一个改进——那里有可训练的卷积层，其中有共享权重。一次更新迭代会生成新的光流更新
$\Delta \mathbf{f}$
，以使每个新步骤的预测更加准确：
$\mathbf{f}_{k+1} = \Delta \mathbf{f} + \mathbf{f}_{k+1}$
。经典的 GRU 单元格方案如下所示：

这些 GRU 单元格输入是先前的隐藏状态
$h_{t-1}$
，也就是
$x_t$
流、相关性和上下文特征的串联。具有卷积层而不是全连接层的 GRU 单元定义为：

$\newline z_t = \sigma（Conv_{3\times 3}（[h_{t-1}， x_t]， W_z）） \newline r_t = \sigma（Conv_{3\times 3}（[h_{t-1}， x_t]， W_r）） \newline \widetilde{h_t} = \tanh （Conv_{3\times 3}（[r_t\odot h_{t-1}， x_t]， W_h）） \newline h_t = （1 - z_t）\odot h_{t-1} + z_t \odot \widetilde{h_t}$
作者也只使用卷积层而不是 GRU 块，但与 GRU 相比，结果更差。这可以通过以下事实来解释：使用门控激活时，流量预测序列的收敛性更好。GRU 细胞在意识形态上与循环网络中的 LSTM 细胞接近，您可以使用此链接阅读有关它们的更多信息。

3.2 RAFT上采样模块

GRU单元以初始图像分辨率的1/8输出光流，因此作者提出了两种不同的上采样方法，以匹配地面实况分辨率。第一个是光流的双线性插值。这是一种简单而快速的方法，但这种方法的质量不如称为凸上采样的学习上采样模块：

凸上采样方法指出，全分辨率光流是GRU单元预测的

$3 \!\!\次\！\!3$

加权网格的凸组合。8 倍图像上采样意味着必须将一个像素扩展为 64 （

$8 \!\!\次\！\!8$

）个像素。凸上采样模块由两个卷积层和末端的softmax激活组成，用于预测上采样光流预测中每个新像素的

$H/8 \times W/8 \times （8 \times 8 \times 9）$

掩模。现在，上采样图像上的每个像素都是先前粗分辨率像素的凸组合，这些像素由带有

${w_1， w_2， ...， w_9}$

系数的预测掩码加权：

根据作者的说法，这种方法预测了更准确的光流输出，特别是在运动边界附近。

3.3 RAFT损失函数

损失函数定义为地面实况和预测之间的 L1 距离，等于 FlowNet 中的损失。所有上采样的循环单元输出都会创建一系列光流预测 $\mathbf f_1, ...,\mathbf f_N$

总loss是地面实况和上采样预测之间每个循环块输出的损耗之和：

$L = \sum_{i=1}^N \gamma^{i-N} ||gt - \mathbf{f}_i||_1， \： \： \： \gamma = 0.8$

3.4 RAFT 与 RAFT-S 对比

正如我们之前提到的，在原始论文中介绍了两种类型的网络架构——RAFT 和 RAFT-S。让我们看一下它们的网络结构对比情况：

首先，Feature 和 Context Extractor 具有不同数量的通道，在 RAFT-S 中，残差单元被瓶颈残差单元替换。其次，RAFT 架构在一个块中有两个卷积 GRU 单元，具有 $1 \times 5$ 和 $5 \times 1$ 卷积核大小，而 RAFT-S 只有一个 $3\times 3$ 卷积GRU 单元。因此，RAFT-S 是 RAFT 的轻量级版本，参数数量较少，同时在质量方面不太准确。

可以在Optical-flow-as-deep-learning任务文件夹中找到与这篇文章相关的README.md文件。

4.带有 PyTorch 的 RAFT

作者还附加了带有预训练权重的 RAFT 架构的开源实现。这很棒，因此我们可以检查此架构并使用作者的开发创建推理脚本。要开始演示，您需要下载我们的代码并安装所需的库。您可以在 README.md 文件中找到有关安装的所有信息。

现在，让我们看一下推理脚本的代码部分。首先，我们需要定义模型并加载预训练的权重：

def inference(args):
    # get the RAFT model
    model = RAFT(args)
    # load pretrained weights
    pretrained_weights = torch.load(args.model)

要对 CPU 和 GPU 设备运行推理，我们需要检查 GPU 可用性。默认情况下，该脚本使用您的 CPU 设备：

if torch.cuda.is_available():
    device = "cuda"
    # parallel between available GPUs
    model = torch.nn.DataParallel(model)
    # load the pretrained weights into model
    model.load_state_dict(pretrained_weights)
    model.to(device)
else:
    device = "cpu"
    # change key names for CPU runtime
    pretrained_weights = get_cpu_model(pretrained_weights)
    # load the pretrained weights into model
    model.load_state_dict(pretrained_weights)

预训练模型在两个 GPU 上使用进行训练，因此预训练模型中的字典键具有额外的前缀名称。要在 CPU 上运行推理，我们需要删除此前缀并获取适合 CPU 的新模型：torch.nn.DataParallelmodule


def get_cpu_model(model):
    new_model = OrderedDict()
    # get all layer's names from model
    for name in model:
        # create new name and update new model
        new_name = name[7:]
        new_model[new_name] = model[name]
    return new_model

现在，我们可以继续我们的演示任务。这部分代码读取两个连续帧，然后输出光流预测。之后，我们对流量预测进行编码并将其转换为 BGR 图像.

with torch.no_grad():
    while True:
        # read the next frame
        ret, frame_2 = cap.read()
        if not ret:
            break
 
        # preprocessing
        frame_2 = frame_preprocess(frame_2, device)
 
        # predict the flow
        flow_low, flow_up = model(frame_1, frame_2, iters=20, test_mode=True)
 
        # vusialize the results
        ret = vizualize_flow(frame_1, flow_up, save, counter)
        if not ret:
            break
        frame_1 = frame_2
        counter += 1

为了可视化流程，我们使用标准 RAFT 的编码器。可视化策略几乎类似于我们在上一篇文章中使用的 HSV。在获得 BGR 格式的光流结果后，我们将图像和光流连接起来进行可视化：

def vizualize_flow(img, flo, save, counter):
    # convert CWH to WHC format and change device if necessary  
    img = img[0].permute(1, 2, 0).cpu().numpy()
    flo = flo[0].permute(1, 2, 0).cpu().numpy()
 
    # map flow to rgb image
    flo = flow_viz.flow_to_image(flo)
    flo = cv2.cvtColor(flo, cv2.COLOR_RGB2BGR)
 
    # concatenate, save and show images
    img_flo = np.concatenate([img, flo], axis=0)
    if save:
        cv2.imwrite(f"demo_frames/frame_{str(counter)}.jpg", img_flo)
    cv2.imshow("Optical Flow", img_flo / 255.0)
 
    k = cv2.waitKey(25) & 0xFF
    if k == 27:
        return False
    return True

要运行演示，可以使用以下命令：

python3 inference.py --model=./models/raft-sintel.pth --video ./videos/crowd.mp4

代码运行时将为您提供带有视频和流预测的推理窗口：

5.模型评估

5.1 指标：终点误差

光流任务的标准误差度量称为端点误差，定义为图像中每个像素的地面实况

$V_{gt} = （\Delta x_{gt}， \Delta y_{gt}）$

值和计算

$V_{calc} = （\Delta x_{calc}， \Delta y_{calc}）$

出的光流值之间的欧几里得距离（对于密集光流估计的情况）。如您所见，FlowNet 和 RAFT 的度量和损失函数相等：

$EPE = ||V_{gt}-V_{calc}||_1 = \sqrt{（\Delta x_{gt}-\Delta x_{calc}）^2 + （\Delta y_{gt}-\Delta y_{calc}）^2}$

5.2 评估基准

有一些开源的光流数据集和基准测试描述了深度学习解决方案在质量和推理时间方面优于经典解决方案的优势。最著名的数据集之一是 SINTEL 数据集，它有自己的基准。FlowNet 架构有相当不错的效果，但它仍然比 DeepFlow 和 EpicFlow 等一些算法方法差。关于RAFT架构，它在SINTEL基准测试中排名第一。您可以比较我们在上一篇文章中回顾的 RLOF 方法对结果的可视化，以及最先进的 RAFT 架构。

此外，还有一个 KITTY 基准测试，其中 RAFT 排名第三，仍然优于任何其他算法方法。

6. 结论

Optical Flow 发现它适用于许多视频编辑任务，如稳定、压缩、慢动作等。此外，一些跟踪和动作识别系统也使用光流数据。在这篇文章中，我们回顾了一些深度学习方法，这些方法以 FlowNet 和 RAFT 架构为例提高了质量。如今，RAFT架构在SINTEL数据集上表现最好，在KITTY数据集上表现最佳。