本篇介绍一篇发表在 ICCV 2019 上的半监督视频目标分割文章:Video Object Segmentation using Space-Time Memory Networks. 将历史帧及 MASK 作为额外的记忆指导当前帧的跟踪与分割。

文章资源

  • 预印本 (Preprint):arxiv
  • 正式发表版本 (Version of Record, VOR):ICCV 2019
  • 代码仓库:github

题目

文章题目是 Video Object Segmentation using Space-Time Memory Networks,指明文章的应用领域是视频目标分割(Video Object Segmentation, VOS),文章使用自己提出的时空记忆网络(STM, Space-Time Memory Networks)来解决视频目标分割(VOS)任务。

作者

作者分别是 Seoung Wug Oh, Joon-Young Lee, Ning Xu, Seon Joo Kim,其中第一作者是 Seoung Wug Oh,他来自于韩国首尔的基督新教私立研究型大学,延世大学(Yonsei University, 1885),该工作是其在 Adobe Research 实习时所做。同时,Seon Joo Kim 也来自该大学。

其他两位作者 Joon-Young Lee, Ning Xu 来自于 Adobe Research.

问题定义

文章把视频分割任务看作是分离视频中所有帧的前景和背景像素问题。文章要解决的是一类半监督视频分割任务,即首帧的目标 MASK 已经给定,预测后续所有帧中目标的 MASKs。视频目标分割任务非常具有挑战性,因为随着时间增长目标的形状很可能会发生剧烈的改变,同时也可能会出现遮挡、漂移(分割到错误的目标上)等情况。

算法关键点示意图

以往的目标跟踪算法是使用前一帧、首帧、首帧+前一帧等进行指导当前帧的跟踪与分割,浪费了丰富的历史帧资源。STM 采用所有历史帧(包含MASK)来指导当前帧的跟踪与分割。

在 STM 框架中,历史帧和 MASK 形成了额外的记忆,当前帧作为一个 query 从记忆的 MASK 中进行分割,具体地,query 和记忆在特征空间中进行稠密的匹配,覆盖所有时空像素位置。这能够很好的处理目标形状改变和遮挡问题。

想法引出

一些视频目标分割算法(作者称为传播的方法)尝试从前一帧的 MASK 来预测当前帧的目标 MASK,这种方法能够很好的解决目标形状的改变问题,但缺乏对遮挡和错误漂移的健壮性。另一些方法(作者称为检查的方法)单纯的使用第一帧的 MASK 来预测后续每一帧的目标 MASK,这种方法和前面的方法真好形成对比,虽然对遮挡和错误漂移具有一定的健壮性,但是不能很好的解决目标形变问题。一些研究者采用首帧和前一帧作为参考帧,提供更多的线索来预测当前帧的目标 MASK,这种方法表现出良好的准确性,是目前最好的方法(作者称为混合方法,混合传播与检测)。

既然使用更多帧(首帧和前一帧)对于视频目标分割有益,那么利用再多一些帧(可能是视频所有帧)是否会更好。文章基于此想法设计一个有效的 DNN 框架利用所有(已预测的)帧的信息解决当前帧目标 MASK 预测问题。提出基于记忆网络的新颖的 DNN 系统,该系统能够对当前帧(query frame)的所有像素计算它与视频多帧的像素的视频注意力(spatio-temporal attention),来预测当前帧每一个像素是否属于前景。该网络称为 STM (Spatio-Temporal Memory) 网络,它利用了丰富的参考帧信息,能够更好的处理目标形状改变、遮挡和漂移等问题。

相关工作

半监督视频目标分割

  • 基于传播的方法。如上图1(a),使用前一帧预测下一帧。
  • 基于检测的方法。如上图1(b),使用首帧预测当前帧。
  • 基于混合的方法,混合检测和传播方法。如图1(c).
  • 在线学习方法。测试时使用初始帧的 MASK 微调网络,虽然能够提高准确性,但是推理时计算昂贵。
  • 离线学习方法。常常设计一个网络能够处理推理时(不可知的)目标分割任务。

文章提出的 STM 属于离线学习方法。通过在额外记忆中维护中间结果(而不是固定帧数),能够自适应的选择必要的信息进行跟踪与分割。

记忆网络

记忆网络是指具有外部存储器的神经网络,其中信息可以按目的写入和读取。可以端到端训练的记忆网络最早是在 NLP 研究中提出的,用于文献问答。通常在这些方法中,可记忆信息分别嵌入到键(输入)和值(输出)特征向量中。键用于寻址相关的记忆,其对应的值被返回。近年来,记忆网络已被应用于个性化图像字幕、视觉跟踪、电影理解和摘要等视觉问题。

虽然文章是基于记忆网络,但扩展了记忆网络的思想,使其适合半监督视频对象分割任务。显然,带有对象掩码的帧被放到记忆中,要分割的帧充当查询。记忆是动态更新的新预测掩模,它无需在线学习。

目标是在给定一组带注释的帧作为记忆的情况下进行逐像素预测。因此,查询帧中的每个像素需要在不同的时空位置访问存储帧中的信息。为此,将记忆转化为包含像素级信息的四维张量,并提出时空记忆读取操作,从四维记忆中定位并读取相关信息。从概念上讲,记忆读取可以被认为是一种时空注意力算法,因为计算每个查询像素的关注时间和地点,以确定像素是否属于前景对象。

Space-Time Memory Networks(STM)

在 STM 框架中,视频帧使用第一帧中给出的 ground truth 注释从第二帧开始顺序处理。在视频处理过程中,将带有对象掩码的过去帧(在第一帧给出或在其他帧估计)作为记忆帧,将没有对象掩码的当前帧作为查询帧。框架概述如下图所示。

当前帧(query)和记忆帧(memory,过去所有帧和 MASK)都通过一个深度 encoder 编码为 key 和 value 特征对。对于当前帧,encoder 的输入仅仅是当前帧图像,而对于记忆帧,encoder 的输入是过去帧的图像和对应的目标 MASK. 但是,输出都是 key 和 value 特征向量对,其中,key 向量用于寻址(addressing)。具体地,query 帧的 key 向量分别于 memory 帧的 key 向量计算相似度,利用该相似度绝对 when-and-where 去检索相关的 memory 帧的 value 特征向量。因此,key 向量被学习为匹配鲁棒的外观变化的视觉语义特征。另一方面,value 向量存储了非常详细的关于产生 MASK 预测的详细信息(如目标对象和目标边界)。query 帧和 memory 帧的 value 向量包含了用于不同目的的信息。具体地,query 帧的 value 向量被学习为存储详细的外观信息方便我们进行准确的解码目标 MASK。而 memory 帧的 value 向量被学习为编码视觉语义和关于每个特征是属于前景还是背景的 MASK 掩码信息。

query 帧和 memory 帧的 key 和 value 对输入到 Space-Time Memory Read 块。query 帧和 memory 帧的 key 向量的每一个“像素” 在视频的真个“时-空”空间上进行稠密的匹配。得到的相对匹配值被用于从 memory 帧中寻址 value 向量,组合这些向量得到一个输出向量。最终,decoder 把 Space-Time Memory Read 块的输出作为输入,重建 query 帧的 MASK.

Key 和 Value 嵌入

Query encoder

query encoder 的输入是 query 帧,即当前帧图像,输出是两个特征图 key 和 value,它们是通过在 backbone 网络后增加两个并行的卷积层(没有非线性层)得到的,这些卷积层当做一个 bottleneck 层,约简 backbone 网络输出的特征图的通道大小,对于 key 特征通道数约简为 1/8,对于 value 特征通道数约简为 1/2。query 帧通过 query encoder 输出得到成对的 key 和 value 对:$k^Q \in \mathbb{R}^{H \times W \times C / 8}, v^Q \in \mathbb{R}^{H \times W \times C / 2}$,这里 $H$ 是高,$W$ 是宽,$C$ 是 backbone 网络输出的特征图的通道数。

Memory encoder

与 query encoder 相比,memory encoder 具有相同的结构,除了输入。memory encoder 的输入是图像帧(RGB 格式,三通道,取值 0~255)和对应的目标 MASK(GRAY 二值图,单通道,取值 0 或者 1) 沿着通道方向进行拼接得到的对象(4通道张量)。

对于包含有多个记忆帧的图像序列,每一个图像帧和 MASK 对分别独立输入 memory encoder 得到 key 和 value 特征对,然后沿着时间维度分别堆叠 key 和 value 向量,得到:$k^M \in \mathbb{R}^{T \times H \times W \times C / 8}, v^M \in \mathbb{R}^{T \times H \times W \times C / 2}$,这里 $T$ 是记忆帧中图像的个数。

文章采用 ResNet50 作为 backbone 网络用于 query encoder 和 memory encoder. 使用 stage-4(res4) 特征图作为基础特征图计算 query 帧和 memory 帧的 key 和 value 特征对。特别地,对于 memory encoder,通过对于第一个卷积层插入额外的单通道卷积核来方便处理 4-通道的张量(拼接的图像帧和对应的 MASK)。网络权重从 ImageNet 预训练模型初始化,除了新增加的单通道卷积核采用随机初始化。

Space-Time Memory Read

在记忆读取(Space-Time Memory Read) 中,首先计算 query 帧和 memory 帧的 key 向量在所有“像素”上的相似度,这种相似度匹配是通过比较 memory 帧 key 特征中的每个时空位置与 query 帧 key 特征中的每个空间位置,以非局部方式执行相似性匹配。然后,通过相似度权重的加权求和来检索 memory 帧的 value,并将其与 query 帧的 value 连接起来。这个记忆读取操作在查询(query)特征图上的每个位置,可以总结为:

$$ \mathtt{y}_i = \lbrack \mathtt{v}_i^Q, {\frac{1}{Z}} \sum_{\forall j}f(\mathtt{k}_i^Q, \mathtt{k}_j^M)\mathtt{v}_j^M \rbrack \tag{1} $$

这里 $i,j$ 分别表示 query 和 memory 的位置索引,$Z=\sum_{\forall j} f(\mathtt{k}_i^Q, \mathtt{k}_j^M)$ 是标准化因子,$\lbrack \cdot, \cdot \rbrack$ 表示拼接。相似度函数定义如下:

$$ f(\mathtt{k}_i^Q, \mathtt{k}_j^M) = \text{exp}(\mathtt{k}_i^Q \circ \mathtt{k}_j^M) \tag{2} $$

这里 $\circ$ 记作点乘(dot-product)。

Decoder

解码器 decoder 接受记忆读取(Space-Time Memory Read)操作的输出并重建当前帧的对象掩码。文章使用 [ ref1](Fast video object segmentation by referenceguided mask propagation)中使用的细化(refinement)模块作为解码器的构建块。首先通过卷积层和残差块将记忆读取(Space-Time Memory Read)的输出压缩为 256 个通道,然后通过多个细化模块逐步将压缩后的特征映射一次提升两倍。每个阶段的细化模块都通过残差连接以相应的比例从查询编码器(query encoder)获取前一阶段的输出和特征映射。最后一个细化块的输出用于通过最后的卷积层重建对象掩码,然后进行 softmax 操作。解码器中的每个卷积层都使用 3×3 滤波器,产生 256 通道输出,但最后一个产生 2 通道输出。解码器以输入图像的 1/4 比例预测掩码。

Multi-object Segmentation

到此,文章针对单目标视频分割已经可以处理,但是对于多目标还不适用。不过,通过 MASK merging 操作能够扩展提出的框架以处理多目标视频分割。为每个对象独立运行上面的模型,并为所有对象计算掩码概率图。然后,使用类似于[ ref1]的软聚合操作合并预测的映射。在[ ref1]中,掩码合并仅在测试期间作为后处理步骤进行。在这项工作中,将操作作为差分网络层,并将其应用于训练和测试中。更多的细节参考文章补充材料。

Two-stage Training

网络首先在静态图像数据生成的模拟数据集上进行预训练。然后,通过主训练对现实世界的视频进行进一步的微调。

Pre-training on images

本文的框架的一个优点是,它不需要很长的训练视频。这是因为该方法在没有任何时间平滑假设的情况下学习了距离像素之间的语义时空匹配。这意味着可以只用几个帧来训练该网络。能够使用图像数据集模拟训练视频。[ ref1] 使用静态图像训练他们的网络,本篇采取了类似的策略。通过对具有不同参数的静态图像应用随机仿射变换(rotation, sheering, zooming, translation, and cropping),生成由 3 帧组成的合成视频剪辑。利用带有目标掩码注释的图像数据集(显著目标检测,语义分割)来预训练网络。通过这样做,期望模型对各种各样的对象外观和类别都是健壮的。

Main training on videos

预训练结束后,使用真实视频数据进行主训练。在这一步中,根据目标评估基准,使用 Youtube-VOS 或DAVIS-2017。为了制作一个训练样本,从一个训练视频中抽取 3 个临时有序的帧。为了了解长时间的外观变化,在采样期间随机跳过帧。在训练过程中,跳过的最大帧数与 curriculum learning 一样,从 0 帧逐渐增加到 25 帧。

Training with dynamic memory

在训练过程中,记忆(memory)会根据网络之前的输出动态更新。当系统逐帧向前移动时,前一步计算的分段输出被添加到 memory 中以用于下一帧。未经阈值处理的原始网络输出成为前景目标的概率图,直接用于 memory 嵌入来建模估计的不确定性。

Training details

文章使用随机裁剪的 384×384 补丁(patches)进行训练。对于所有实验,将 minibatch 大小设置为 4,并禁用所有 batch 规范化层。使用 Adam 优化器最小化交叉熵损失,固定学习率为 1e-5。预训练大约需要 4 天,主训练大约需要 3 天,使用 4 块 NVIDIA GeForce 1080 Ti gpu。

Inference

将所有之前的帧写入 memory 可能会引起实际问题,例如 GPU 记忆溢出和运行速度慢。相反,文章通过一个简单的规则来选择要放到 memory 中的帧。第一帧和前一帧是最重要的参考信息。第一帧总是提供可靠的信息,因为它是 ground-truth MASK。前一帧在外观上与当前帧相似,因此可以期望精确的像素匹配和掩码传播。因此,默认将这两帧放入 memory 中。

对于中间帧,只需每 N 帧保存一个新的 memory 帧。N 是一个控制速度和精度之间权衡的超参数,除非另有说明,否则使用 N = 5。值得注意的是,框架在没有额外训练的情况下实现了在线学习和在线适应的效果。通过不更新模型参数,将之前的帧存储到 memory 中,很容易实现在线模型更新的效果。因此,文章提出的方法比以前的大多数方法运行得快得多,同时达到了最先进的精度。

评估 Evaluation

作者在 Youtube-VOS 和 DAVIS benchmark 上评估了模型。具体地,分别在这两个数据集的训练集上训练得到两个模型。对于 Youtube-VOS,使用官方的数据划分,训练集包含视频 3471 个。对于 DAVIS,使用 DAVIS-2017 的训练集的 60 个视频进行训练,分别在 DAVIS-2016 和 DAVIS-2017 上进行评估。除此之外,又额外使用 Youtube-VOS 进行训练后再在 DAVIS 上进行了评估,给出结果。所有结果都是直接从网络输出后没有进行后处理。

评估标准是 region similarity $\mathcal{J}$ 和 contour accuracy $\mathcal{F}$。

Youtube-VOS

Youtube-VOS 是用于视频对象分割的大规模数据集,由 4453 个带有多个对象注释的视频组成。该数据集比由 120 个视频组成的流行的 DAVIS 基准大 30 倍。它还具有未见对象类别(unseen object categories)的验证数据。因此,它可以很好地衡量不同算法的泛化性能(generalization performance)。验证集由 474 个视频组成,包括 91 个对象类别。它对 65 种可见和 26 种不可见的物体类别有单独的测量(separate measures)。

在 Youtube-VOS 验证集上的多目标视频分割结果,其中 STM 表示文章提出的方法:

Seen Seen Unseen Unseen
Overall $\mathcal{J}$ $\mathcal{F}$ $\mathcal{J}$ $\mathcal{F}$
OSMN 51.2 60.0 60.1 40.6 44.0
MSK 53.1 59.9 59.5 45.0 47.9
RGMP 53.8 59.5 - 45.2 -
OnAVOS 55.2 60.1 62.7 46.6 51.4
RVOS 56.8 63.6 67.2 45.5 51.0
OSVOS 58.8 59.8 60.5 54.2 60.7
S2S 64.4 71.0 70.0 55.5 61.2
A-GAME 66.1 67.8 - 60.8 -
PreMVOS 66.9 71.4 75.9 56.5 63.7
BoLTVOS 71.1 71.6 - 64.3 -
STM 79.4 79.7 84.2 72.8 80.9

DAVIS

单目标(DAVIS-2016)

使用包含 20 个视频的 DAVIS-2016 验证集,每个视频都为单个目标对象标注了高质量的掩码。在下表中,包含了在线学习的使用,并提供了每种方法的大致运行时间。以前大多数表现最好的方法都依赖于在线学习,这严重损害了运行速度。在所有没有在线学习的竞争方法中,文章提出的方法达到了最好的准确率,并且在很短的时间内显示出与基于在线学习的最佳方法的竞争结果。另外,使用来自 Youtube-VOS 的额外数据进行训练,在很大程度上优于所有方法。

下表是在 DAVIS-2016 验证集上的评估结果。OL 表示在线学习。(+YV) 表示额外在 Youtube-VOS 上训练。由于篇幅限制,$\mathcal{J}$ Mean 小于 79 的方法在本文中略去,完整的表格见文章补充资料:

OL $\mathcal{J}$ Mean $\mathcal{F}$ Mean Time
S2S(+YV) $\surd$ 79.1 - 9s
MSK $\surd$ 79.7 75.4 12s
OSVOS $\surd$ 79.8 80.6 9s
MaskRNN $\surd$ 80.7 80.9 -
VideoMatch 81.0 - 0.32s
FEELVOS(+YV) 81.1 82.2 0.45s
RGMP 81.5 82.0 0.13s
A-GAME(+YV) 82.0 82.2 0.07s
FAVOS 82.4 79.5 1.8s
LSE $\surd$ 82.9 80.3 -
CINN $\surd$ 83.4 85.0 >30s
PReMVOS $\surd$ 84.9 88.6 >30s
$\text{OSVOS}^S$ $\surd$ 85.6 86.4 4.5s
OnAVOS $\surd$ 86.1 84.9 13s
DyeNet $\surd$ 86.2 - 2.32s
STM 84.8 88.1 0.16s
STM(+YV) 88.7 89.9 0.16s

多目标(DAVIS-2017)

DAVIS-2017 是 DAVIS-2016 的多目标扩展。验证集包含 30 个视频 59 个目标。在下表中报告了验证集上的多目标视频分割结果。同样,文章提出的方法在没有在线学习的快速方法中表现出最好的性能。使用额外的 Youtube-VOS 数据,文章提出的方法在很大程度上优于之前所有最先进的方法,包括 2018 年 DAVIS 挑战赛的获胜者。在测试开发集上的结果包含在补充材料中。

DAVIS-2017 验证集的定量评价。OL 表示在线学习。(+YV) 表示使用 Youtube-VOS 进行额外。* 表示 $\mathcal{J}$ 均值与 $\mathcal{F}$ 均值的平均值:

OL $\mathcal{J}$ Mean $\mathcal{F}$ Mean
OSMN 52.5 57.1
FAVOS 54.6 61.8
VidMatch 56.5 68.2
OSVOS $\surd$ 56.6 63.9
MaskRNN $\surd$ 60.5 -
OnAVOS $\surd$ 64.5 71.2
$\text{OSVOS}^S$ $\surd$ 64.7 71.3
RGMP 64.8 68.6
CINN $\surd$ 67.2 74.2
A-GAME(+YV) 67.2 72.7
FEELVOS(+YV) 69.1 74.0
DyeNet $\surd$ *74.1
PReMVOS $\surd$ 73.9 81.7
STM 69.2 74.0
STM(+YV) 79.2 84.3

使用额外训练数据带来的巨大性能飞跃表明,由于过度拟合,DAVIS 太小而无法训练可推广的深度网络。它还解释了为什么 DAVIS 基准测试中表现最佳的在线学习方法在大规模 Youtube-VOS 基准测试中没有表现出良好的性能。在线学习方法很难得到大量训练数据的帮助。这些方法通常需要大量的参数搜索(如数据合成方法、优化迭代、学习率和后处理),这对于大规模基准测试来说并不容易做到。)

定性结果

下图显示了文章结果的定性示例。从 Youtube-VOS 和 DAVIS 验证集中选择具有挑战性的视频,并对重要帧进行采样(例如遮挡之前和之后)。如图所示,文章提出的方法对于遮挡和复杂运动具有鲁棒性。更多结果请访问补充材料。

分析

训练数据

通过两个训练阶段训练模型:静态图像的预训练和使用 DAVIS 或 Youtube-VOS 的主训练。在下表中,比较了文章的方法与不同训练数据的性能。此外,提供跨数据集验证来衡量泛化性能。

Youtube-VOS 和 DAVIS-2017 验证集上的训练数据分析。比较通过不同训练阶段训练的模型。此外,报告了交叉验证结果(即使用在 Youtube-VOS 上训练的模型评估 DAVIS,反之亦然)。

Variants Youtube-VOS DAVIS-2017 DAVIS-2017
Overall $\mathcal{J}$ $\mathcal{F}$
Pre-training only 69.1 57.9 62.1
Main-training only 68.2 38.1 47.9
Full training 79.4 69.2 74.0
Cross validation 56.3 78.6 83.5

仅预训练(Pre-training only)。有趣的是,在不使用任何真实视频的情况下,仅预训练模型优于仅主训练模型以及 YouTube-VOS 上的所有其他方法。然而,通过使用这两种训练策略获得了最大的性能。

仅主要训练(Main-training only)。如果没有预训练阶段,模型在 Youtube-VOS 上的总体得分下降了 11.2。这表明尽管 Youtube-VOS 提供了超过 3000 个训练视频,但训练视频数据量仍然不足以发挥文章提出的网络的潜力。此外,DAVIS 上的性能非常低意味着存在严重的过度拟合问题,因为训练损失与完整模型相似(没有应用提前停止)。推测在预训练过程中遇到的不同对象有助于模型的泛化并防止过度拟合。

交叉验证(Cross validation)。评估在 DAVIS 和 Youtube-VOS 上训练的模型,反之亦然。在 DAVIS 上训练的模型在 Youtube-VOS 上表现有限。这是预期的结果,因为 DAVIS 太小,无法学习对其他数据集的泛化能力。另一方面,我们在 Youtube-VOS 上训练的模型在 DAVIS 上表现良好,并且优于所有其他方法。

记忆管理(Memory Management)

为了最小的记忆消耗和最快的运行时间,可以将第一帧和/或前一帧保存在记忆(memory)中。为了获得最大的准确度,最终模型除了第一帧和前一帧之外,每 5 帧保存一个新的中间记忆帧。

在下表中比较了不同的记忆管理规则。将第一个帧和前一个帧都保存到记忆中是最重要的,文章的模型使用两个记忆帧实现了最先进的精度。这是因为模型足够强大,可以处理较大的外观变化,同时通过有效利用记忆对漂移和错误积累具有鲁棒性。最重要的是,拥有中间帧存储器可以通过解决极具挑战性的情况进一步提高性能,如下图所示:

Memory frame(s) Youtube-VOS DAVIS-2016 DAVIS-2017 Time
First 68.9 81.4 67.0 0.06
Previous 69.7 83.2 69.6 0.06
First & Previous 78.4 87.8 77.7 0.07
Every 5 frames 79.4 88.7 79.2 0.16

为了进行更深入的分析,在下图中显示了帧级精度分布。对所有视频帧中所有对象的 Jaccard 分数进行排序,并绘制分数以分析在具有挑战性的场景中的性能。将最终模型(每 5 帧)与 First 和 Previous 进行比较,以检查使用额外记忆帧的效果。虽然两种设置在成功范围(超过 30%)上表现同样出色,但对于困难情况(低于 30%),额外记忆帧的效果变得清晰可见。 10% 和 30% 之间的巨大准确度差距表明,文章提出的网络可以通过额外的记忆帧更好地处理具有挑战性的情况。比较“仅第一个”和“仅前一个”,前一帧看起来对于处理失败情况更有用。

记忆可视化。在下图中,可视化了记忆读取操作,以验证学习到的时空匹配。从可视化中可以看出,记忆读取操作准确匹配查询和记忆帧之间的相应像素。

结论

文章提出了一种用于半监督视频对象分割的新颖时空记忆网络,在准确性和速度方面都具有较好的表现。所提出的时空记忆网络有很大潜力成为许多其他像素级估计问题的突破。框架的预期能够应用在其他领域研究中,包括对象跟踪、交互式图像/视频分割和修复(inpainting)。

扩展

  1. Space-Time Memory Networks for Video Object Segmentation With User Guidance