用于视频修复的STTN

论文 | 项目 | 幻灯片 |BibTex

学习联合空时变换以进行视频修复

曾艳红, 傅建龙, 和 晁洪洋。
发表于ECCV 2020。

引用

如果我们的论文和仓库中的任何部分对您的工作有所帮助，请慷慨地使用以下方式引用：

@inproceedings{yan2020sttn,
  author = {Zeng, Yanhong and Fu, Jianlong and Chao, Hongyang,
  title = {Learning Joint Spatial-Temporal Transformations for Video Inpainting},
  booktitle = {The Proceedings of the European Conference on Computer Vision (ECCV)},
  year = {2020}
}

简介

高质量的视频修复技术能够填补视频帧中的缺失区域，这是一项充满前景但极具挑战性的任务。

在本文中，我们提出了一种用于视频修复的联合空时Transformer网络（STTN）。具体而言，我们通过提出的多尺度补丁注意力模块，同时填充所有输入帧中的缺失区域。STTN由空时对抗损失进行优化。

为了展示所提模型的优越性，我们使用标准的静态掩码以及更真实的动态物体掩码进行了定量和定性评估。

安装

克隆此仓库。

git clone git@github.com:researchmm/STTN.git
cd STTN/

我们的项目基于PyTorch和Python构建。对于所需的完整Python包集合，我们建议从提供的YAML文件创建一个Conda环境，例如：

conda env create -f environment.yml 
conda activate sttn

使用预训练模型完成视频

可以使用预训练模型生成结果视频。作为参考，我们提供了一个在Youtube-VOS数据集上预训练的模型（Google Drive文件夹）。

1. 从Google Drive文件夹下载预训练模型，并将其保存到checkpoints/目录下。

2. 使用预训练模型完成视频。例如：

python test.py --video examples/schoolgirls_orig.mp4 --mask examples/schoolgirls  --ckpt checkpoints/sttn.pth 

输出视频将保存在examples/目录中。

数据集准备

我们在datasets/目录中提供了数据集划分。

准备Youtube-VOS（2018）数据集。 该数据集可以从这里下载。特别地，我们遵循标准的训练/验证/测试划分（3,471/474/508）。数据集应按照如下目录结构组织：

datasets
    ｜- youtube-vos
        |- JPEGImages
           |- <video_id>.zip
           |- <video_id>.zip
        |- test.json 
        |- train.json 

准备DAVIS（2018）数据集。 该数据集可以从这里下载。特别地，有90个带有密集标注对象掩码的视频，以及60个未标注的视频。数据集应按照如下目录结构组织：

datasets
    ｜- davis
        |- JPEGImages
          |- cows.zip
          |- goat.zip
        |- Annoatations
          |- cows.zip
          |- goat.zip
        |- test.json 
        |- train.json 

训练新模型

一旦数据集准备就绪，就可以使用以下命令训练新模型。例如：

python train.py --config configs/youtube-vos.json --model sttn 

测试

测试过程与使用预训练模型完成视频类似。

python test.py --video examples/schoolgirls_orig.mp4 --mask examples/schoolgirls  --ckpt checkpoints/sttn.pth 

输出视频将保存在examples/目录中。

可视化

我们在visualization.ipynb中提供了一个可视化注意力图的例子。

训练监控

我们可以通过运行以下命令来监控训练过程中的损失：

tensorboard --logdir release_mode                                                    

联系方式

如果您对本论文有任何疑问或建议，请随时联系我（zengyh7@mail2.sysu.edu.cn）。

STTN 视频修复快速上手指南

STTN (Spatial-Temporal Transformer Network) 是一个用于高质量视频修复（Video Inpainting）的开源项目，曾发表于 ECCV 2020。它通过联合时空变换网络，能够同时处理多帧中的缺失区域，适用于去除视频中不需要的物体或修复损坏片段。

环境准备

在开始之前，请确保您的开发环境满足以下要求：

• 操作系统: Linux 或 macOS (Windows 需使用 WSL 或类似环境)
• Python: 推荐 Python 3.6+
• 深度学习框架: PyTorch
• 包管理工具: Conda (强烈推荐，可自动解决依赖冲突)
• 硬件: 建议使用带有 CUDA 支持的 NVIDIA GPU 进行推理和训练

安装步骤

1. 克隆代码仓库

首先将项目代码下载到本地：

git clone git@github.com:researchmm/STTN.git
cd STTN/

2. 创建并激活 Conda 环境

项目提供了 environment.yml 文件以简化依赖安装。执行以下命令创建名为 sttn 的虚拟环境并激活它：

conda env create -f environment.yml 
conda activate sttn

提示：如果下载 environment.yml 中的包速度较慢，可尝试配置国内镜像源（如清华源或中科大源）： conda config --add channels https://mirrors.tuna.tsinghua.edu.cn/anaconda/pkgs/main/

3. 下载预训练模型

为了直接体验效果，需要下载在 Youtube-VOS 数据集上预训练的模型。

• 下载地址: Google Drive
• 存放路径: 下载完成后，将文件重命名为 sttn.pth 并放入项目根目录下的 checkpoints/ 文件夹中。
  • 若没有该文件夹，请手动创建：mkdir checkpoints

基本使用

安装完成并准备好模型后，即可运行最简单的视频修复示例。以下命令将读取示例视频和对应的掩码（Mask），生成修复后的视频。

运行推理命令：

python test.py --video examples/schoolgirls_orig.mp4 --mask examples/schoolgirls  --ckpt checkpoints/sttn.pth 

参数说明：

• --video: 输入的原视频路径。
• --mask: 对应的掩码文件路径（指示需要修复的区域）。
• --ckpt: 预训练模型权重文件的路径。

结果输出： 修复完成的视频将自动保存在 examples/ 目录下。您可以使用任何视频播放器查看修复效果。