STCN

通过改进的内存覆盖重新思考时空网络，以实现高效的视频目标分割

Ho Kei Cheng、Yu-Wing Tai、Chi-Keung Tang

NeurIPS 2021

[arXiv] [PDF] [项目页面] [Papers with Code]

欢迎查看我们的新作 Cutie！

新闻: 在 YouTubeVOS 2021 挑战赛 中，STCN 在新颖（未知）类别中取得了准确率第一，在总体准确率中取得了第二名。我们的解决方案还具有速度快、轻量级的特点。

我们提出了时空对应网络（STCN），作为一种新的、高效且有效的框架，用于在视频目标分割的背景下建模时空对应关系。 STCN 在多个基准测试中取得了 SOTA 结果，同时在不使用额外花哨技术的情况下，仍能以 20+ FPS 的速度运行。使用混合精度时，其速度甚至更高。 尽管效果显著，但该网络本身非常简单，仍有很大的改进空间。具体技术细节请参阅论文。

更新（2021年7月15日）

1. CBAM 块：我们尝试过不使用 CBAM 块，我认为我们确实不需要它。对于 s03 模型，我们在 DAVIS 上降低了 1.2 分，在 YouTubeVOS 上提高了 0.1 分。对于 s012 模型，我们在 DAVIS 和 YouTubeVOS 上都提高了 0.1 分。您可以选择移除此块（请参阅 no_cbam 分支）。总的来说，规模更大的 YouTubeVOS 似乎是一个更合适的评估基准，能够更好地衡量一致性。

更新（2021年8月22日）

2. 可重复性：我们已更新了下方的软件包要求。在该环境下，我们在两台不同的机器上多次运行后，得到的 DAVIS J&F 分数范围为 [85.1, 85.5]。

更新（2022年4月27日）

多尺度测试代码（如论文中所述）已添加至 这里。

我们在这里提供了什么？

1. 简要介绍

2. 定量结果和预计算输出

   1. DAVIS 2016
   2. DAVIS 2017 验证集/测试集
   3. YouTubeVOS 2018/2019

3. 在您自己的数据上试用我们的模型（提供交互式 GUI）

4. 复现步骤

   1. 预训练模型
   2. 推理
   3. 训练

5. 如果您想进一步了解

6. 引用

简要介绍

本工作主要有两个贡献：STCN 框架（上图）和 L2 相似度。我们是在图像之间建立亲和力，而不是在（图像，掩码）对之间——这大大提高了速度、节省了内存（因为我们只计算一个亲和力矩阵，而不是多个），并且增强了鲁棒性。此外，我们还用 L2 相似度替代点积，从而极大地提高了记忆库的利用率。

优点

• 简单，运行速度快（使用混合精度时可达 30+ FPS；不使用时也有 20+ FPS）
• 性能优异
• 仍有很大的改进空间（例如局部性、内存空间压缩）
• 易于训练：只需两块 11GB 的 GPU，无需 V100

要求

我们在开发该项目时使用了以下软件包及版本：

• PyTorch 1.8.1
• torchvision 0.9.1
• OpenCV 4.2.0
• Pillow-SIMD 7.0.0.post3
• progressbar2
• thinspline 用于训练（pip install git+https://github.com/cheind/py-thin-plate-spline）
• gitpython 用于训练
• gdown 用于下载预训练模型
• 我环境中的其他软件包，仅供参考。

请参考官方 PyTorch 指南 安装 PyTorch/torchvision，以及 pillow-simd 指南安装 Pillow-SIMD。其余软件包可通过以下命令安装：

pip install progressbar2 opencv-python gitpython gdown git+https://github.com/cheind/py-thin-plate-spline

结果

符号说明

• FPS 是平均值，计算方式为总处理时间除以总帧数，与对象数量无关，即多对象 FPS，并在 RTX 2080 Ti 上测量，排除了 IO 时间。
• 我们还提供了使用自动混合精度（AMP）时的推理速度——性能几乎相同。论文中的速度数据未使用 AMP 测量。
• 所有评估均在 480p 分辨率下进行。test-dev 的 FPS 是在验证集上以相同的内存设置（每三帧保存一次内存）测量的，以确保一致性。

[预计算输出 - Google Drive]

[预计算输出 - OneDrive]

s012 表示带有 BL 预训练的模型，而 s03 表示没有预训练的模型（之前在 MiVOS 中称为 s02）。

数据（s012）

| 数据集 | 总分 | J-Seen | F-Seen | J-Unseen | F-Unseen | | --- | --- | :--:|:--:|:---:|:---:|:---:| | YouTubeVOS 18 | 验证集 | 84.3 | 83.2 | 87.9 | 79.0 | 87.2 | | YouTubeVOS 19 | 验证集 | 84.2 | 82.6 | 87.0 | 79.4 | 87.7 |

对于 DAVIS interactive，我们将 MiVOS 的传播模块从 STM 更改为 STCN。详细信息请参阅 此链接。

在您自己的数据上试用（提供交互式 GUI）

如果您（不知何故）拥有第一帧的分割结果（或者更一般地说，每个对象首次出现时的分割结果），可以使用 eval_generic.py。请查看该文件顶部的说明。

如果您只是想交互式地体验一下，我强烈推荐 我们对 MiVOS 的扩展 :yellow_heart: ——它配备了交互式 GUI，而且非常高效、有效。

复现结果

预训练模型

我们对 YouTubeVOS 和 DAVIS 使用相同的模型。你可以自行下载并将它们放入 ./saves/ 目录，或者使用 download_model.py 脚本。

s012 模型（效果更好）：[Google Drive] [OneDrive]

s03 模型：[Google Drive] [OneDrive]

s0 预训练模型：[GitHub]

s01 预训练模型：[GitHub]

推理

• eval_davis_2016.py 用于 DAVIS 2016 验证集
• eval_davis.py 用于 DAVIS 2017 验证集和 test-dev 集（通过 --split 参数控制）
• eval_youtube.py 用于 YouTubeVOS 2018/19 验证集（通过 --yv_path 参数控制）

命令行参数的提示信息应该能让你大致了解如何使用它们。例如，如果你已经使用我们的脚本下载了数据集和预训练模型，那么对于 DAVIS 2017 验证集的评估，你只需要指定输出路径即可：python eval_davis.py --output [somewhere]。而对于 YouTubeVOS 的评估，则需要将 --yv_path 指向你选择的版本。

多尺度测试代码（如论文中所述）已添加到 这里。

训练

数据准备

我建议要么对现有数据进行软链接（ln -s），要么使用提供的 download_datasets.py 脚本，按照我们的格式组织数据集。download_datasets.py 可能会下载比你需要的更多的内容——只需注释掉你不想要的部分即可。该脚本不会下载 BL30K 数据集，因为它非常庞大（超过 600GB），我们不希望导致你的硬盘崩溃。具体结构如下：

├── STCN
├── BL30K
├── DAVIS
│   ├── 2016
│   │   ├── Annotations
│   │   └── ...
│   └── 2017
│       ├── test-dev
│       │   ├── Annotations
│       │   └── ...
│       └── trainval
│           ├── Annotations
│           └── ...
├── static
│   ├── BIG_small
│   └── ...
├── YouTube
│   ├── all_frames
│   │   └── valid_all_frames
│   ├── train
│   ├── train_480p
│   └── valid
└── YouTube2018
    ├── all_frames
    │   └── valid_all_frames
    └── valid

BL30K

BL30K 是在 MiVOS 中提出的一个合成数据集。

你可以使用自动脚本 download_bl30k.py，也可以手动从 MiVOS 下载。请注意，每个片段大约 115GB 大小——总共约 700GB。运行该脚本时，你需要大约 1TB 的可用磁盘空间（包括解压缓冲区）。 Google 可能会屏蔽 Google Drive 链接。你可以 1) 将文件夹创建快捷方式到你自己的 Google Drive，然后 2) 使用 rclone 从你自己的 Google Drive 复制（这不会计入你的存储限额）。

训练命令

CUDA_VISIBLE_DEVICES=[a,b] OMP_NUM_THREADS=4 python -m torch.distributed.launch --master_port [cccc] --nproc_per_node=2 train.py --id [defg] --stage [h]

我们使用分布式数据并行（DDP）在两块 11GB 显存的 GPU 上进行训练。请将 a, b 替换为 GPU ID，cccc 替换为未使用的端口号，defg 替换为唯一的实验标识符，h 替换为训练阶段（0/1/2/3）。

模型按不同阶段逐步训练（0：静态图像；1：BL30K；2：30 万张主训练数据；3：15 万张主训练数据）。每个阶段完成后，我们会加载最新训练好的权重来开始下一个阶段。

仅在阶段 0 训练得到的模型不能直接使用。所需映射关系请参见 model/model.py: load_network。

以 _checkpoint 为后缀的 .pth 文件用于恢复中断的训练（通过 --load_model 参数），但通常并不需要。一般情况下，你只需要使用 --load_network 并加载最后一个网络权重（文件名中不含 checkpoint）。

更深入地了解

• 若要添加你的数据集，或进行数据增强操作：dataset/static_dataset.py, dataset/vos_dataset.py
• 若想修改相似度函数或记忆读取过程：model/network.py: MemoryReader, inference_memory_bank.py
• 若想调整网络结构：model/network.py, model/modules.py, model/eval_network.py
• 若想改进传播过程：model/model.py, eval_*.py, inference_*.py

引用

如果您觉得本仓库有用，请引用我们的论文（如果您使用 top-k，请引用 MiVOS）！

@inproceedings{cheng2021stcn,
  title={重新思考时空网络：通过改进内存覆盖实现高效的视频目标分割},
  author={Cheng, Ho Kei 和 Tai, Yu-Wing 和 Tang, Chi-Keung},
  booktitle={NeurIPS},
  year={2021}
}

@inproceedings{cheng2021mivos,
  title={模块化交互式视频目标分割：从交互到掩码、传播与差异感知融合},
  author={Cheng, Ho Kei 和 Tai, Yu-Wing 和 Tang, Chi-Keung},
  booktitle={CVPR},
  year={2021}
}

如果您还想引用数据集：

@inproceedings{shi2015hierarchicalECSSD,
  title={基于扩展 CSSD 的层次化图像显著性检测},
  author={Shi, Jianping 和 Yan, Qiong 和 Xu, Li 和 Jia, Jiaya},
  booktitle={TPAMI},
  year={2015},
}

@inproceedings{wang2017DUTS,
  title={利用图像级监督学习检测显著对象},
  author={Wang, Lijun 和 Lu, Huchuan 和 Wang, Yifan 和 Feng, Mengyang 
  和 Wang, Dong，以及 Yin, Baocai 和 Ruan, Xiang}, 
  booktitle={CVPR},
  year={2017}
}

@inproceedings{FSS1000,
  title = {FSS-1000：用于少样本分割的1000类数据集},
  author = {Li, Xiang 和 Wei, Tianhan 和 Chen, Yau Pun 和 Tai, Yu-Wing 和 Tang, Chi-Keung},
  booktitle={CVPR},
  year={2020}
}

@inproceedings{zeng2019towardsHRSOD,
  title = {迈向高分辨率显著目标检测},
  author = {Zeng, Yi 和 Zhang, Pingping 和 Zhang, Jianming 和 Lin, Zhe 和 Lu, Huchuan},
  booktitle = {ICCV},
  year = {2019}
}

@inproceedings{cheng2020cascadepsp,
  title={{CascadePSP}：通过全局与局部细化实现类别无关且超高分辨率的分割},
  author={Cheng, Ho Kei 和 Chung, Jihoon 和 Tai, Yu-Wing 和 Tang, Chi-Keung},
  booktitle={CVPR},
  year={2020}
}

@inproceedings{xu2018youtubeVOS,
  title={Youtube-vos：大规模视频目标分割基准},
  author={Xu, Ning 和 Yang, Linjie 和 Fan, Yuchen 和 Yue, Dingcheng 和 Liang, Yuchen 和 Yang, Jianchao 和 Huang, Thomas},
  booktitle = {ECCV},
  year={2018}
}

@inproceedings{perazzi2016benchmark,
  title={视频目标分割的基准数据集与评估方法},
  author={Perazzi, Federico 和 Pont-Tuset, Jordi 和 McWilliams, Brian 和 Van Gool, Luc 和 Gross, Markus 和 Sorkine-Hornung, Alexander},
  booktitle={CVPR},
  year={2016}
}

@inproceedings{denninger2019blenderproc,
  title={BlenderProc},
  author={Denninger, Maximilian 和 Sundermeyer, Martin 和 Winkelbauer, Dominik 和 Zidan, Youssef 和 Olefir, Dmitry 和 Elbadrawy, Mohamad 和 Lodhi, Ahsan 和 Katam, Harinandan},
  booktitle={arXiv:1911.01911},
  year={2019}
}

@inproceedings{shapenet2015,
  title       = {{ShapeNet：一个信息丰富的3D模型库}},
  author      = {Chang, Angel Xuan 和 Funkhouser, Thomas 和 Guibas, Leonidas 和 Hanrahan, Pat 和 Huang, Qixing 和 Li, Zimo 和 Savarese, Silvio 和 Savva, Manolis 和 Song, Shuran 和 Su, Hao 和 Xiao, Jianxiong 和 Yi, Li 和 Yu, Fisher},
  booktitle   = {arXiv:1512.03012},
  year        = {2015}
}

联系方式：hkchengrex@gmail.com

STCN 快速上手指南

STCN (Space-Time Correspondence Networks) 是一个高效、轻量级的视频目标分割（VOS）框架，在多个基准测试中取得了 SOTA 结果。它无需复杂配置即可在单张显卡上以 20+ FPS 的速度运行，支持混合精度加速后更是可达 30+ FPS。

环境准备

系统要求

• 操作系统: Linux (推荐)
• GPU: 至少需要两张 11GB 显存的显卡进行训练（推理仅需一张）。支持 CUDA。
• Python: 建议 Python 3.6+

前置依赖

本项目基于以下核心版本开发，建议尽量匹配以避免兼容性问题：

• PyTorch 1.8.1
• torchvision 0.9.1
• OpenCV 4.2.0
• Pillow-SIMD 7.0.0.post3 (用于加速图像处理)

注意: 请优先参考 PyTorch 官网 安装与您 CUDA 版本匹配的 PyTorch 和 torchvision。国内用户可使用清华或中科大镜像源加速安装。

安装步骤

1. 克隆仓库

   git clone https://github.com/hkchengrex/STCN.git
   cd STCN
   

2. 安装核心依赖 首先确保已安装 PyTorch 和 torchvision，然后运行以下命令安装其余依赖包：

   pip install progressbar2 opencv-python gitpython gdown
   

3. 安装特殊依赖 (Pillow-SIMD 和 thinspline) Pillow-SIMD 能显著提升数据加载速度，thinspline 用于训练阶段的数据增强。

   安装 Pillow-SIMD:

   pip uninstall pillow
   pip install pillow-simd==7.0.0.post3
   

   (如果上述命令失败，请参考 Pillow-SIMD GitHub 页面的编译安装指南)

   安装 thinspline:

   pip install git+https://github.com/cheind/py-thin-plate-spline
   

基本使用

1. 下载预训练模型

你可以手动下载或使用提供的脚本下载模型。模型将保存在 ./saves/ 目录下。

推荐模型 (s012 - 性能更好):

python download_model.py --model s012

若脚本下载缓慢，可手动从 Google Drive 下载并放入 ./saves/ 目录。

2. 准备数据

你需要准备视频文件和对应的首帧掩码（Mask）。 目录结构示例：

dataset_name/
├── JPEGImages/
│   └── video_id/
│       ├── 00000.jpg
│       ├── 00001.jpg
│       └── ...
└── Annotations/
    └── video_id/
        ├── 00000.png  (首帧掩码，后续帧无需提供)
        └── ...

3. 运行推理 (Inference)

使用 eval_generic.py 对自己的数据进行测试。

python eval_generic.py --dataset [你的数据集路径] --output [输出结果路径] --weights [模型权重路径，如 ./saves/s012.pth]

参数说明：

• --dataset: 包含 JPEGImages 和 Annotations 文件夹的根目录。
• --output: 分割结果保存路径。
• --weights: 预训练模型路径。

4. 交互式体验 (可选)

如果你希望拥有图形界面（GUI）来交互式地标注和分割视频，推荐使用集成了 STCN 的 MiVOS 扩展版：

git clone -b MiVOS-STCN https://github.com/hkchengrex/MiVOS.git
cd MiVOS
# 按照该仓库说明安装并运行 interactive_gui.py

5. 验证基准测试结果 (可选)

如果你已下载 DAVIS 或 YouTubeVOS 数据集，可直接运行官方评估脚本：

DAVIS 2017 验证集:

python eval_davis.py --output [输出路径] --split val

YouTubeVOS 验证集:

python eval_youtube.py --yv_path [YouTubeVOS 数据集路径] --output [输出路径]