Logo
Project Icon

STCN

改进内存覆盖的高效视频对象分割框架

视频目标分割STCN空间时间对应神经网络NeurIPSGithub开源项目

STCN是一个创新的视频对象分割框架,通过改进内存覆盖重新构建时空网络。该方法在多个基准测试中达到了最先进水平,同时保持20+ FPS的高效运行。STCN采用简洁的网络结构,建立图像间亲和力,并使用L2相似度替代点积,显著提升内存利用率。这种方法在准确性和效率间实现了理想平衡,为视频对象分割研究带来新思路。

访问官网访问官网GithubGithub论文论文
介绍相关项目

STCN

重新思考具有改进内存覆盖的时空网络,用于高效视频对象分割

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

NeurIPS 2021

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

查看我们的新工作 Cutie!

bmx pigs

新闻: 在 YouTubeVOS 2021 挑战赛中,STCN 在新类别(未知类别)中获得第一名的准确率,在整体准确率中获得第二名。我们的解决方案也快速且轻量。

我们提出了时空对应网络(STCN)作为一个新的、有效的、高效的框架,用于在视频对象分割背景下建模时空对应关系。 STCN 在多个基准测试中达到了最先进的结果,同时运行速度快达20+ FPS,没有使用任何花哨的技巧。使用混合精度时,其速度甚至更高。 尽管效果显著,但网络本身非常简单,还有很大的改进空间。有关技术细节请参阅论文。

更新 (2021年7月15日)

  1. CBAM 模块:我们尝试了不使用 CBAM 模块,我认为我们确实不太需要它。对于 s03 模型,DAVIS 数据集上降低了 1.2,YouTubeVOS 数据集上提高了 0.1。对于 s012 模型,DAVIS 数据集上提高了 0.1,YouTubeVOS 数据集上提高了 0.1。欢迎您删除这个模块(参见 no_cbam 分支)。总的来说,规模更大的 YouTubeVOS 似乎是一个更好的一致性评估基准。

更新 (2021年8月22日)

  1. 可重复性:我们已更新了下面的包依赖要求。使用该环境,我们在两台不同的机器上进行了多次运行,在 DAVIS 数据集上获得了 [85.1, 85.5] 范围内的 J&F 分数。

更新 (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+)
  • 高性能
  • 仍有很大的改进空间(例如局部性、内存空间压缩)
  • 易于训练:只需要两个11GB的GPU,不需要V100

要求

我们在开发这个项目时使用了以下软件包/版本。

请参考官方 [PyTorch 指南](https://github.com/hkchengrex/STCN/blob/main/https://pytorch.org/ 安装 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&FJFFPSFPS (AMP)
DAVIS 2016验证91.790.493.026.940.8
DAVIS 2017验证85.382.088.620.234.1
DAVIS 2017test-dev79.976.383.514.622.7
数据集分割总体得分J-已见F-已见J-未见F-未见
YouTubeVOS 18验证84.383.287.979.087.2
YouTubeVOS 19验证84.282.687.079.487.7
数据集AUC-J&FJ&F @ 60s
DAVIS 交互式88.488.8

对于DAVIS交互式,我们将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 验证集和测试开发集（由 --split 控制）
  • eval_youtube.py 用于 YouTubeVOS 2018/19 验证集（由 --yv_path 控制）

参数提示应该能让你大致了解如何使用它们。例如，如果你已经使用我们的脚本下载了数据集和预训练模型，你只需要指定输出路径：python eval_davis.py --output [某个路径] 即可进行 DAVIS 2017 验证集评估。对于 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]

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

模型通过不同的阶段进行渐进式训练（0：静态图像；1：BL30K；2：300K 主训练；3：150K 主训练）。每个阶段结束后，我们通过加载最新训练的权重来开始下一个阶段。

（仅在阶段 0 训练的模型不能直接使用。请参阅 model/model.py: load_network 了解我们所做的必要映射。）

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

因此，要训练 s012 模型，我们按顺序启动三个训练步骤，如下所示：

在静态图像上预训练：CUDA_VISIBLE_DEVICES=0,1 OMP_NUM_THREADS=4 python -m torch.distributed.launch --master_port 9842 --nproc_per_node=2 train.py --id retrain_s0 --stage 0

在 BL30K 数据集上预训练：CUDA_VISIBLE_DEVICES=0,1 OMP_NUM_THREADS=4 python -m torch.distributed.launch --master_port 9842 --nproc_per_node=2 train.py --id retrain_s01 --load_network [path_to_trained_s0.pth] --stage 1

主训练：CUDA_VISIBLE_DEVICES=0,1 OMP_NUM_THREADS=4 python -m torch.distributed.launch --master_port 9842 --nproc_per_node=2 train.py --id retrain_s012 --load_network [path_to_trained_s01.pth] --stage 2

要训练 s03 模型，我们按顺序启动两个训练步骤，如下所示： 静态图像预训练：CUDA_VISIBLE_DEVICES=0,1 OMP_NUM_THREADS=4 python -m torch.distributed.launch --master_port 9842 --nproc_per_node=2 train.py --id retrain_s0 --stage 0

主要训练：CUDA_VISIBLE_DEVICES=0,1 OMP_NUM_THREADS=4 python -m torch.distributed.launch --master_port 9842 --nproc_per_node=2 train.py --id retrain_s03 --load_network [path_to_trained_s0.pth] --stage 3

深入了解

  • 添加数据集或进行数据增强：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 and Tai, Yu-Wing and Tang, Chi-Keung},
  booktitle={NeurIPS},
  year={2021}
}

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

如果您想引用数据集： @inproceedings{shapenet2015, 标题 = {{ShapeNet：一个信息丰富的三维模型仓库}}, 作者 = {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}, 会议论文集 = {arXiv:1512.03012}, 年份 = {2015} }


联系方式：<hkchengrex@gmail.com>
相关项目
项目侧边栏1项目侧边栏2
推荐项目
Project Cover

豆包MarsCode

豆包 MarsCode 是一款革命性的编程助手，通过AI技术提供代码补全、单测生成、代码解释和智能问答等功能，支持100+编程语言，与主流编辑器无缝集成，显著提升开发效率和代码质量。

Project Cover

AI写歌

Suno AI是一个革命性的AI音乐创作平台，能在短短30秒内帮助用户创作出一首完整的歌曲。无论是寻找创作灵感还是需要快速制作音乐，Suno AI都是音乐爱好者和专业人士的理想选择。

Project Cover

有言AI

有言平台提供一站式AIGC视频创作解决方案，通过智能技术简化视频制作流程。无论是企业宣传还是个人分享，有言都能帮助用户快速、轻松地制作出专业级别的视频内容。

Project Cover

Kimi

Kimi AI助手提供多语言对话支持，能够阅读和理解用户上传的文件内容，解析网页信息，并结合搜索结果为用户提供详尽的答案。无论是日常咨询还是专业问题，Kimi都能以友好、专业的方式提供帮助。

Project Cover

阿里绘蛙

绘蛙是阿里巴巴集团推出的革命性AI电商营销平台。利用尖端人工智能技术，为商家提供一键生成商品图和营销文案的服务，显著提升内容创作效率和营销效果。适用于淘宝、天猫等电商平台，让商品第一时间被种草。

Project Cover

吐司

探索Tensor.Art平台的独特AI模型，免费访问各种图像生成与AI训练工具，从Stable Diffusion等基础模型开始，轻松实现创新图像生成。体验前沿的AI技术，推动个人和企业的创新发展。

Project Cover

SubCat字幕猫

SubCat字幕猫APP是一款创新的视频播放器，它将改变您观看视频的方式！SubCat结合了先进的人工智能技术，为您提供即时视频字幕翻译，无论是本地视频还是网络流媒体，让您轻松享受各种语言的内容。

Project Cover

美间AI

美间AI创意设计平台，利用前沿AI技术，为设计师和营销人员提供一站式设计解决方案。从智能海报到3D效果图，再到文案生成，美间让创意设计更简单、更高效。

Project Cover

AIWritePaper论文写作

AIWritePaper论文写作是一站式AI论文写作辅助工具，简化了选题、文献检索至论文撰写的整个过程。通过简单设定，平台可快速生成高质量论文大纲和全文，配合图表、参考文献等一应俱全，同时提供开题报告和答辩PPT等增值服务，保障数据安全，有效提升写作效率和论文质量。

使用协议隐私政策广告服务
投诉举报邮箱: service@vectorlightyear.com
@2024 懂AI·鲁ICP备2024100362号-6·鲁公网安备37021002001498号