看清视频的像素流——使用飞桨框架复现RAFT光流估计模型

2022-11-29 10:40:07 浏览数 (1)

光流的概念

光流的概念是大佬James J. Gibson在1950年首先提出来的,是空间运动物体在成像平面上的像素运动的瞬时速度,是利用图像序列中像素的变化以及相邻帧之间的相关性,来找到上一帧跟当前帧的像素点之间存在的对应关系,从而计算出相邻帧之间像素点的运动信息的一种方法。一般而言,光流是由于场景中前景目标本身的移动、相机的运动,或者两者的共同运动所产生的。

根据是否选取图像稀疏点进行光流估计,可以将光流估计分为稀疏光流和稠密光流,如下图(左)选取了一些特征明显(梯度较大)的点进行光流估计和跟踪,下图(右)为连续帧稠密光流示意图。

稠密光流描述图像每个像素向下一帧运动的光流。为了方便表示,使用不同的颜色和亮度表示光流的大小和方向,如下图的不同颜色。下图展示了一种光流和颜色的映射关系,使用颜色表示光流的方向,亮度表示光流的大小。

最为常用的视觉算法库OpenCV中,提供光流估计算法接口,包括稀疏光流估计算法cv2.calcOpticalFlowPyrLK()和稠密光流估计cv2.calcOpticalFlowFarneback()。其中稀疏光流估计算法为1981年Lucas和Kanade两位科学家提出的Lucas-Kanade算法,该算法最为经典也较容易理解。

近几年出现了基于深度学习的光流估计算法,开山之作是FlowNet[1],于2015年首先使用CNN解决光流估计问题,取得了较好的结果,并且在CVPR2017上发表改进版本FlowNet2.0[2],成为当时State-of-the-art的方法。

截止到现在,FlowNet和FlowNet2.0分别被引用790次和552次,依然是深度学习光流估计算法中引用率最高的论文。随后出现了PWC[3]、RAFT[4]等一系列深度学习模型,并不断刷新EPE(光流估计的评价指标)。需要注意的是,基于深度学习的光流估计算法都是针对稠密光流估计问题。

项目介绍

本项目复现了光流估计的经典模型RAFT,该模型出自ECCV2020年的优秀论文《Recurrent All-Pairs Field Transforms for Optical Flow》。由于AI Studio开源项目和数据集中鲜有涉及光流估计领域,为了让大家更好地理解该领域,本项目从光流估计的基础入手,先对数据集、评价指标进行介绍,然后再介绍RAFT模型。

光流数据集

此处的光流数据集是指深度学习模型所需要的数据集。光流数据集由于标注困难,因此多以合成数据为主。最近提出的数据集autoflow[5]将光流合成的参数变为可学习的参数,使用这类数据集可以达到更快的拟合、更好的效果。而最常用的还是公开的固定数据集FlyingChairs、FlyingThings3D等。

本项目主要使用FlyingChairs公开数据集,它由22872个图像对和相应的光流组成。每个图像显示了在随机背景前移动的3D椅子模型,椅子和背景的运动都是纯平面的。数据集中图片是.ppm格式、光流是.flo格式,所有的文件都在data文件夹下。FlyingChairs数据集提供了一个train-validation split txt,里面的每一行为1或者2,1代表训练集,2代表验证集。

上图对比发现,椅子的相对位置有变化

光流估计评价指标

EPE(Endpoint Error)是光流估计中标准的误差度量,是预测光流向量与真实光流向量的欧氏距离在所有像素上的均值(越低越好)。

RAFT模型搭建

RAFT由3个主要组件组成

  • 从两个输入图像img1和img2中提取每像素特征的特征编码器(Feature Encoder)以及仅从img1中提取特征的上下文编码器(Context Encoder)。
  • 一个相关层,通过取所有特征向量对的内积,构造4D W×H×W×H相关体。4D矩阵的最后2维在多个尺度上汇集,以构建一组多尺度体积。
  • 一种更新操作符,它通过使用当前估计值从相关体积集合中查找值来重复更新光流。

接下来进行详解与代码展示

特征提取

RAFT特征提取分为图片特征提取(Feature Encoder)和上下文特征提取(Context Encoder),这两部分使用了相同的结构,由六个残差块(结构和ResNet相似)组成,但是使用了不同的normalization方法,前者使用了BatchNorm,后者使用了InstanceNorm。残差块的代码如下:

代码语言:javascript复制
import paddle
import paddle.nn as nn
import paddle.nn.functional as F
import numpy as np 
from scipy import interpolate

# 本文件是raft模型中的特征提取和上下文提取模块,模块主要使用了残差卷积

class ResidualBlock(nn.Layer):
    def __init__(self, in_planes, planes, norm_fn='batch', stride=1):
        super(ResidualBlock, self).__init__()
        self.conv1 = nn.Conv2D(in_planes, planes,kernel_size=3, padding=1, stride=stride, weight_attr=nn.initializer.KaimingNormal())
        self.conv2 = nn.Conv2D(planes, planes, kernel_size=3, padding=1, weight_attr=nn.initializer.KaimingNormal())
        self.relu = nn.ReLU()  # in_planes=True

        if norm_fn == 'batch':
            self.norm1 = nn.BatchNorm2D(planes, weight_attr=paddle.ParamAttr(initializer=nn.initializer.Constant(value=1.0)))
            self.norm2 = nn.BatchNorm2D(planes, weight_attr=paddle.ParamAttr(initializer=nn.initializer.Constant(value=1.0)))
            if not stride == 1:
                self.norm3 = nn.BatchNorm2D(planes, weight_attr=paddle.ParamAttr(initializer=nn.initializer.Constant(value=1.0)))
        elif norm_fn == 'instance':
            self.norm1 = nn.InstanceNorm2D(planes, weight_attr=paddle.ParamAttr(initializer=nn.initializer.Constant(value=1.0)))
            self.norm2 = nn.InstanceNorm2D(planes, weight_attr=paddle.ParamAttr(initializer=nn.initializer.Constant(value=1.0)))
            if not stride == 1:
                self.norm3 = nn.InstanceNorm2D(planes, weight_attr=paddle.ParamAttr(initializer=nn.initializer.Constant(value=1.0)))

        if stride == 1:
            self.downsample = None
        else:
            self.downsample = nn.Sequential(
                nn.Conv2D(in_planes, planes, kernel_size=1, stride=stride, weight_attr=nn.initializer.KaimingNormal()), self.norm3
            )

    def forward(self, x):
        y = x
        y = self.relu(self.norm1(self.conv1(y)))
        y = self.relu(self.norm2(self.conv2(y)))
        if self.downsample is not None:
            x = self.downsample(x)
        return self.relu(x y)

成本量计算

这步是文章的一个核心。图片特征提取分别从img1和img2提取出M1:H×W×D、M2:H×W×D,然后对M1和M2进行点乘,得到C如下公式所示:

这样一个相关信息张量是非常大的, 因此在C的最后两个维度上进行汇合来降低维度大小, 产生4个相关信息C1、C2、C3、C4(原文中是生成了四个,实际上还可以更多),每个Ci的维度保持前两个维度不变,后两个维度分别除2的i次方,例如C3的维度:H×W×H/8×W/8。

代码语言:javascript复制
import paddle
import paddle.nn.functional as F 
from utils import *

# 本部分涉及了计算光流的成本量函数

class CorrBlock:
    def __init__(self, fmap1, fmap2, num_levels=4, radius=4):
        self.num_levels = num_levels
        self.radius = radius
        self.corr_pyramid = []

        corr = CorrBlock.corr(fmap1, fmap2)
        batch, h1, w1, dim, h2, w2 = corr.shape
        corr = paddle.reshape(corr, shape=[batch*h1*w1, dim, h2, w2])        
        self.corr_pyramid.append(corr)
        for i in range(self.num_levels-1):
            corr = F.avg_pool2d(corr, 2, stride=2)
            self.corr_pyramid.append(corr)

    def __call__(self, coords):
        r = self.radius
        coords = paddle.transpose(coords, perm=[0, 2, 3, 1])
        batch, h1, w1, _ = coords.shape

        out_pyramid = []
        for i in range(self.num_levels):
            corr = self.corr_pyramid[i]
            # linspace为线性等分 r为搜索半径
            dx = paddle.linspace(-r, r, 2*r 1)
            dy = paddle.linspace(-r, r, 2*r 1)
            delta = paddle.stack(paddle.meshgrid(dy, dx), axis=-1)

            centroid_lvl = paddle.reshape(coords, shape=[batch*h1*w1, 1, 1, 2]) / 2**i
            delta_lvl = paddle.reshape(delta, shape=[1, 2*r 1, 2*r 1, 2])
            coords_lvl = centroid_lvl   delta_lvl

            corr = bilinear_sampler(corr, coords_lvl)
            corr = paddle.reshape(corr, shape=[batch, h1, w1, -1])
            out_pyramid.append(corr)
        out = paddle.concat(out_pyramid, axis=-1)
        return paddle.transpose(out, perm=[0, 3, 1, 2]).astype('float32')
    # transpose、permute 操作虽然没有修改底层一维数组,但是新建了一份Tensor元信息,并在新的元信息中的重新指定stride。

    @staticmethod
    def corr(fmap1, fmap2):
        batch, dim, ht, wd = fmap1.shape
        # 修改tensor的形状
        fmap1 = paddle.reshape(fmap1, shape=[batch, dim, ht*wd])
        fmap2 = paddle.reshape(fmap2, shape=[batch, dim, ht*wd]) 

        corr = paddle.matmul(paddle.transpose(fmap1, perm=[0,2,1]), fmap2)
        corr = paddle.reshape(corr, shape=[batch, ht, wd, 1, ht, wd])

        # 归一化
        return corr / paddle.sqrt(paddle.to_tensor(float(dim)))

为何要进行这样计算?有2个方面原因:

  • 进行这样的计算可以找到前一张图片和后一张图片的像素点之间的联系;
  • 这种相关信息张量可以保证同时捕捉到较大和较小的像素位移。

更新迭代

模型经过多次迭代计算得出最终flow,流程如下:

  • 最开始,将flow初始化为0;
  • 每次迭代,计算出一个△flow;
  • 更新当前flow:flow = flow △flow。

每次迭代计算△flow的流程为:

  • 使用当前的flow,通过前面介绍的Correlation,插值查询到一个相关性矩阵。对于图片img1上的每个像素,该矩阵包括了其在图片img2上所有潜在位置的相关性。
  • 通过一个GRU来计算出本次迭代的△flow以及GRU的hidden status,下次迭代作为GRU的输入。
  • 根据本次迭代计算出来的△flow,更新当前的flow后得出本次迭代的flow 。这里得到的flow的分辨率还是原图的1/8,所以再通过上采样得出跟原图同分辨率的flow。后面会用这个flow来计算loss。

RAFT模型创造性的引入GRU作为迭代的主体,成为了很长一段时间内光流估计流域的通用架构。GRU是LSTM网络的一种效果很好的变体,它较LSTM网络的结构更加简单,而且效果也很好,因此也是当前非常流行的一种网络。GRU既然是LSTM的变体,因此也是可以解决RNN网络中的长依赖问题。GRU的参数较少,因此训练速度更快,GRU能够降低过拟合的风险。但是在RAFT训练过程中迭代了12次,依然是耗费了大量的计算资源。

GRU的结构图如下图所示:

代码语言:javascript复制
class SepConvGRU(nn.Layer):
    def __init__(self, hidden_dim=128, input_dim=192 128):
        super(SepConvGRU, self).__init__()
        self.convz1 = nn.Conv2D(hidden_dim input_dim, hidden_dim, (1,5), padding=(0,2))
        self.convr1 = nn.Conv2D(hidden_dim input_dim, hidden_dim, (1,5), padding=(0,2))
        self.convq1 = nn.Conv2D(hidden_dim input_dim, hidden_dim, (1,5), padding=(0,2))

        self.convz2 = nn.Conv2D(hidden_dim input_dim, hidden_dim, (5,1), padding=(2,0))
        self.convr2 = nn.Conv2D(hidden_dim input_dim, hidden_dim, (5,1), padding=(2,0))
        self.convq2 = nn.Conv2D(hidden_dim input_dim, hidden_dim, (5,1), padding=(2,0))

    def forward(self, h, x):
        # horizontal
        hx = paddle.concat([h, x], axis=1)
        z = F.sigmoid(self.convz1(hx))
        r = F.sigmoid(self.convr1(hx))
        q = F.tanh(self.convq1(paddle.concat([r*h, x], axis=1)))
        h = (1-z) * h   z * q

        # vertical
        hx = paddle.concat([h, x], axis=1)
        z = F.sigmoid(self.convz2(hx))
        r = F.sigmoid(self.convr2(hx)) # 在v1版本中,本行的convr2被写成convz2!

        q = F.tanh(self.convq2(paddle.concat([r*h, x], axis=1)))
        h = (1-z) * h   z* q
        return h

训练策略

官方代码中使用了AdamW优化器 OneCycleLR的学习率策略,AdamW在飞桨比较常见,但是OneCycleLR并没有飞桨版本,这个策略出自2017年《Super-Convergence: Very Fast Training of Neural Networks Using Large Learning Rates》,它的变化曲线如下图所示:

虽然没有飞桨版本,但是飞桨提供了内置lr策略,也可以对lr策略进行组合,所以我使用了线性热启动 线性衰减代替了OneCycleLR,虽然在训练过程中震荡的情况稍微明显一些,但是最后的结果还是可以达到和官方代码一样的水平。

代码语言:javascript复制
 # --------------------------Hyperparameter begin----------------------------------------

    num_steps = 120000                      # 训练次数

    batch_size = 16                         # 一批训练数量

    gamma = 0.85                            # 这个是求损失用到的gamma,官方默认是0.85

    max_flow = 400                          # 像素移动的最大值,超过最大值则忽略不计

    num_workers = 4                         # DataLoader读取数据的进程

    learning_rate = 0.00040                 # 学习率

    log_iter = 100                          # 每训练log_iter次打印一次日志

    VAL_FREQ = 2500                         # 每训练VAL_FREQ次验证一次

    log = Logger('test.txt', batch_size)    # 日志文件的名字,位置是work/log

    warmup_proportion = 0.3

    # 学习率策略:线性热启动 线性衰减

    polynomial_lr = paddle.optimizer.lr.PolynomialDecay(learning_rate=learning_rate, decay_steps=int((1-warmup_proportion)*num_steps)   100, end_lr = learning_rate / 10000)

    scheduler = paddle.optimizer.lr.LinearWarmup(learning_rate=polynomial_lr, warmup_steps=int(warmup_proportion*num_steps), start_lr = 0.00001, end_lr=learning_rate)



    # 优化器:AdamW

    optimizer = paddle.optimizer.AdamW(learning_rate=scheduler, weight_decay=0.00001, epsilon=1e-8, parameters=model.parameters())

    # --------------------------Hyperparameter end---------------------------------------

总结

光流估计是计算机视觉研究中的一个重要方向,其不像其他感知任务会显式的在应用中呈现,但是对于视频理解有很大的帮助。例如视频理解的经典架构双流网络的双流分别是RGB视频流和光流,同时在无人驾驶、人体关键点估计等领域都有应用。

  • 目前基于深度学习的光流估计算法,都使用了大规模的数据集,包括FlyingThings、FlyingChairs等合成数据集,KITTI、Sintel等真实数据集,但是由于AI Studio的存储空间、训练时间有限,所以我在本项目中只使用了FlyingChairs数据集,所以目前的模型在处理真实数据集时还有提升的空间。
  • 推荐使用AI Studio的V100 32G环境训练,如果是16G环境,需要调低batchsize,调低lr,增加num_steps。
  • 我训练了97500次,EPE降至0.88,并使用当前权重进行推理,测试图片如下,左侧真实光流,右侧是推理结果,还可以继续优化的。

参考文献

[1].Dosovitskiy, A., et al. (2015). Flownet: Learning optical flow with convolutional networks. Proceedings of the IEEE international conference on computer vision.

[2].Ilg, E., et al. (2017). Flownet 2.0: Evolution of optical flow estimation with deep networks. Proceedings of the IEEE conference on computer vision and pattern recognition.

[3].Sun, D., et al. (2018). Pwc-net: Cnns for optical flow using pyramid, warping, and cost volume. Proceedings of the IEEE conference on computer vision and pattern recognition.

[4].Teed, Z. and J. Deng (2020). Raft: Recurrent all-pairs field transforms for optical flow. European conference on computer vision, Springer.

[5].Sun, D., et al. (2021). Autoflow: Learning a better training set for optical flow. Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition.

0 人点赞