使用RAFT的深度光流

刘二婷ttt

于 2024-08-27 03:12:08 发布

阅读量16

点赞数

文章目录

概要
介绍
RAFT 的基础

特征提取
视觉相似性
相关性金字塔
相关性查找
高效的相关性查找（可选）
迭代更新
更新块
凸上采样
训练

如何使用RAFT
结论
参考

概要

在这篇文章中，我们将了解一种旗舰的光流深度学习方法，该方法获得了 2020 年 ECCV 最佳论文奖，并被引用超过 1000 次。它也是KITTI基准上许多表现最好的模型的基础。此模型称为 RAFT：Recurrent All-Pairs Field Transforms for Optical Flow，在 PyTorch 或 GitHub 上很容易获得。这些实现使其具有高度可访问性，但模型很复杂，理解它可能会令人困惑。在这篇文章中，我们将把 RAFT 分解成它的基本组件，并详细了解它们中的每一个。然后我们将学习如何在 Python 中使用它来估计光流。在第 2 部分中，我们将探索晦涩难懂的细节，并将不同的块可视化，以便我们对它们的工作原理有更深的理解。

介绍

光流是像素在图像序列中的表观运动。为了估计光流，场景中物体的运动必须具有相应的亮度位移。这意味着一张图像中移动的红球在下一张图像中应该具有相同的亮度和颜色，这使我们能够确定它在像素方面移动了多少。图 1 显示了一个光流示例，其中逆时针旋转的吊扇被一系列图像捕获。

使用RAFT的深度光流_人工智能

最右边的彩色图像包含了从第 1 帧到第 2 帧每个像素的明显运动轨迹，并进行了颜色编码，不同的颜色表示像素运动的不同水平和垂直方向。这就是密集光流估计的一个例子。

密集光流估计为每个像素分配一个二维光流向量，描述其在时间间隔内的水平和垂直位移。在稀疏光流中，该向量只分配给与边角等强特征相对应的像素。为了使流向量存在，像素在 t 时刻的强度必须与 t+1 时刻相同，这就是所谓的亮度一致性假设。时间 t 时位置 (x,y) 处的图像强度或亮度由 I(x,y,t) 给出。下面我们以已知像素位移为例（如图 2 所示）来直观地说明这一点，其中 dx 和 dy 是图像的水平和垂直位移，dt 是帧间的时间差。

使用RAFT的深度光流_光流_02

亮度一致性假设意味着（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 ]。

使用RAFT的深度光流_权重_03

水平梯度 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的深度光流_人工智能_04

请注意它如何具有三个主要块：特征编码器块、视觉相似性块和迭代更新块。 RAFT 架构有两种规模，一种是具有 480 万个参数的大型架构，一种是具有 100 万个参数的小型架构。在这篇文章中，我们将重点关注大型架构，但是一旦我们了解了大型架构，了解小型架构就没有什么意义了。

特征提取

RAFT 使用由六个残差块组成的卷积神经网络 (CNN) 对两个输入图像执行特征提取，并使用 D 特征图将每个图像下采样至 1/8 分辨率。

使用RAFT的深度光流_权重_05

特征编码器网络 g 对具有共享权重的两个图像进行操作，而上下文编码器网络 f 仅对 I₁ 进行操作并提取用作流估计的主要参考的特征。除了细微的差别外，两个网络的整体架构几乎相同。上下文网络使用批量归一化，而特征网络使用实例归一化，上下文网络提取 C = c + h 特征图，其中 c 是上下文特征图的数量，h 是将初始化隐藏特征图的隐藏特征图的数量。迭代更新块的状态。

使用RAFT的深度光流_权重_06

注意：原始论文经常使用简写符号：HxW 来引用特征图尺寸 H/8xW/8。这可能会令人困惑，因此我们将遵循 H’ = H/8 的约定，使得特征图大小为 H’xW’。我们还将从 I₁ 中提取的特征图张量称为 s g 1 ，对于 I2 也是如此。

视觉相似性

视觉相似度是 4D H’xW’xH’xW’ 全对相关体积 C，通过取特征图的点积计算得出。

使用RAFT的深度光流_光流_07

在相关性体积中，特征图 g¹ 中的每个像素都与特征图 g² 中的每个像素都有计算出的相关性，我们将这些相关性中的每一个称为 2D 响应图（见图 5）。在 4D 中思考可能具有挑战性，因此请想象将体积的前两个维度展平：（H’xW’）xH’xW’ ，我们现在有一个 3D 体积，其中 g¹ 的每个像素都有自己的 2D 响应图，显示它与 g² 的每个像素位置的相关性。由于特征是从图像派生的，因此响应映射实际上指示了给定的 I₁ 像素与 I₂ 的每个像素的相关性。

视觉相似度是一个全对相关卷，通过计算每个像素位置上每个特征图的相关性，将 I₁ 的像素与 I₂ 的每个像素联系起来。

高效的相关性查找（可选）

相关性查找以 O（N²）进行缩放，其中 N 是像素数，这可能是大型图像的瓶颈，但有一个等效的操作与 O（NM）缩放，其中 M 是金字塔级别的数。此操作将相关性金字塔与查找相结合，而操作则利用了内部积的线性和平均池化。2mx2m网格上的平均相关响应Cm（金字塔水平m）如下所示。

使用RAFT的深度光流_人工智能_10

对于给定的金字塔级别 m，我们不需要对特征图 g1 求和，这意味着可以通过特征图 g1 与平均池化特征图 g² 的内积来计算相关性，其复杂度为 O (N) .由于这仅对单个金字塔级别 m 有效，因此我们必须计算每个级别的内积，使其按 O(M) 缩放，总复杂度为 O(NM) 。我们没有预先计算金字塔的相关性，而是只预先计算池化特征图，并在查找发生时按需计算相关性值。

迭代更新

更新算子估算一系列流量： {f₀ , f ₁ ,…, f ₙ } 从一个初始起点 f₀ 开始，这个起点可以是全部 0，也可以是前推的上一次流量估计值（热启动）。在每次迭代 k 中，它都会产生一个流量更新方向 Δf，并将其添加到当前估计值中：fₖ₊₁ = fₖ + Δfₖ。更新算子模仿优化算法，经过训练可提供更新，使估计的流量序列收敛到一个固定点：fₖ → f* 。

更新块

更新块将以下输入作为输入：相关特征、电流估计、上下文特征和隐藏特征。其架构与突出显示的子块如下所示。

使用RAFT的深度光流_权重_11

更新块内的子块是：

特征提取块 — 从相关性、流和 I₁（上下文网络）中提取运动特征。
循环更新块 - 循环计算流更新
Flow Head - 最终卷积层，将流量估计调整为 H/8 x W/8 x 2

如图 6 所示，循环更新块的输入是流、相关性和上下文特征的串联。潜在隐藏状态使用上下文网络中的隐藏特征进行初始化。（上下文网络提取一堆 2D 特征图，然后将其分为上下文和隐藏特征图）。循环更新块由 2 个可分离的 ConvGRU 组成，可在不显着增加网络规模的情况下增加感受野。在每次更新时，来自循环更新块的隐藏状态被传递到流头以获得大小为 H/8 x W/8 x 2 的流估计。然后使用凸上采样对该估计进行上采样。

凸上采样

RAFT 的作者对双线性和凸上采样进行了实验，发现凸上采样可以显着提高性能。

使用RAFT的深度光流_迭代_12

凸上采样将每个精细像素估计为其相邻 3x3 粗像素网格的凸组合

让我们来分解一下凸上采样的工作原理，下面的图 8 提供了一个很好的视觉效果。

使用RAFT的深度光流_人工智能_13

首先，我们假设精细分辨率像素是其最近的粗邻居的 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的深度光流_人工智能_14

如何使用RAFT

我们可以使用 RAFT 来估计我们自己图像上的密集光流。首先，我们需要克隆 GitHub 存储库并下载模型。本教程的代码位于 GitHub 上。

!git clone https://github.com/princeton-vl/RAFT.git

%cd RAFT
!./download_models.sh
%cd ..

接下来，我们将 RAFT 的核心添加到路径中

现在，我们需要一些辅助函数来连接 RAFT 类。注意：这些辅助函数仅针对 CUDA 编写，但您可以使用 Colab .NET 轻松访问 GPU。

import torch
from raft import RAFT
from utils import flow_viz
from utils.utils import InputPadder


def process_img(img, device='cuda'):
    return torch.from_numpy(img).permute(2, 0, 1).float()[None].to(device)


def load_model(weights_path, args):
    """ Loads model to CUDA only """
    model = RAFT(args)
    pretrained_weights = torch.load(weights_path, map_location=torch.device("cpu"))
    model = torch.nn.DataParallel(model)
    model.load_state_dict(pretrained_weights)
    model.to("cuda")
    return model


def inference(model, frame1, frame2, device='cuda', pad_mode='sintel',
              iters=12, flow_init=None, upsample=True, test_mode=True):

    model.eval（)
    with torch.no_grad():
        # preprocess
        frame1 = process_img(frame1, device)
        frame2 = process_img(frame2, device)

        padder = InputPadder(frame1.shape, mode=pad_mode)
        frame1, frame2 = padder.pad(frame1, frame2)

        # predict flow
        if test_mode:
          flow_low, flow_up = model(frame1,
                                    frame2,
                                    iters=iters,
                                    flow_init=flow_init,
                                    upsample=upsample,
                                    test_mode=test_mode)
          return flow_low, flow_up

        else:
            flow_iters = model(frame1,
                               frame2,
                               iters=iters,
                               flow_init=flow_init,
                               upsample=upsample,
                               test_mode=test_mode)

            return flow_iters


def get_viz(flo):
    flo = flo[0].permute(1,2,0).cpu().numpy()
    return flow_viz.flow_to_image(flo)

注意 inference() 中的输入填充，我们需要确保所有图像都能被 8 整除。 raft.py 代码可以轻松地从命令行访问，但如果我们想与之交互，我们将需要重写其中一些，或者我们可以创建一个特殊的类来向它传递参数。

# class to interface with RAFT
class Args():
  def __init__(self, model='', path='', small=False, 
               mixed_precision=True, alternate_corr=False):
    self.model = model
    self.path = path
    self.small = small
    self.mixed_precision = mixed_precision
    self.alternate_corr = alternate_corr

  """ Sketchy hack to pretend to iterate through the class objects """
  def __iter__(self):
    return self

  def __next__(self):
    raise StopIteration

Args 类的默认初始化将直接与任何大型 RAFT 模型交互。为了演示 RAFT，我们将使用缓慢旋转的吊扇视频中的帧。现在我们可以加载模型并估计光流。

model = load_model("RAFT/models/raft-sintel.pth", args=Args())
flow_low, flow_up = inference(model, frame1, frame2, device='cuda', test_mode=True)

测试模式将返回 1/8 分辨率流和凸上采样流。

使用RAFT的深度光流_光流_15

结论

在这篇文章中，我们了解了 RAFT，这是一种能够估计准确流场的高级模型。 RAFT 能够通过从提取的特征图中计算所有对相关量来捕获所有像素之间的关系。构建相关金字塔是为了捕获大像素位移和小像素位移。查找算子根据当前流量估计从相关金字塔中提取新的相关特征。更新块使用相关特征和当前流量估计来提供迭代更新，迭代更新收敛到最终流量估计，该最终流量估计通过凸上采样进行上采样。在第 2 部分中，我们将解压网络并了解一些关键块的工作原理。

参考

https://towardsdatascience.com/optical-flow-with-raft-part-1-f984b4a33993

原创作者: u_15699099 转载于: https://blog.51cto.com/u_15699099/11843632

刘二婷ttt

关注

0
点赞
踩
0

收藏

觉得还不错? 一键收藏
0
评论
使用RAFT的深度光流

文章目录概要介绍RAFT 的基础特征提取视觉相似性相关性金字塔相关性查找高效的相关性查找（可选）迭代更新更新块凸上采样训练如何使用RAFT结论参考概要在这篇文章中，我们将了解一种旗舰的光流深度学习方法，该方法获得了 2020 年 ECCV 最佳论文奖，并被引用超过 1000 次。它也是KITTI基准上许多表现最好的模型...
复制链接

扫一扫