基于少量样本块训练的视频交互式风格化

官方实现

基于少量样本块训练的视频交互式风格化
O. Texler, D. Futschik, M. Kučera, O. Jamriška, Š. Sochorová, M. Chai, S. Tulyakov, and D. Sýkora
[项目主页], [论文], [BiBTeX]

运行

下载 testing-data.zip 并解压。请确保 _train 文件夹与 _gen 文件夹位于同级目录。

预训练模型

如果你想快速测试网络，这里提供一些 预训练模型。解压后，按照生成步骤操作。调用 generate.py 时，请确保设置正确的 --checkpoint 路径，例如 _pre-trained-models/Zuzka2/model_00020.pth。

训练

要训练网络，请运行 train.py。参考以下示例命令：

train.py --config "_config/reference_P.yaml" 
		 --data_root "Zuzka2_train" 
		 --log_interval 1000 
		 --log_folder logs_reference_P

每 1000 个周期（log_interval），train.py 会将当前的生成器保存到 logs_reference_P（log_folder）目录中，并在 _gen 数据上验证/运行生成器，结果将保存在 Zuzka2_gen/res__P 目录中。

生成

要生成结果，请运行 generate.py。

generate.py --checkpoint "Zuzka2_train/logs_reference_P/model_00020.pth" 
	    --data_root "Zuzka2_gen"
	    --dir_input "input_filtered"
	    --outdir "Zuzka2_gen/res_00020" 
	    --device "cuda:0"

要在实时网络摄像头画面上生成结果，请运行 generate_webcam.py。要停止生成，请在预览窗口激活时按下 q 键。

generate_webcam.py --checkpoint "Zuzka2_train/logs_reference_P/model_00020.pth" 
	    --device "cuda:0"
	    --resolution 1280 720
	    --show_original 1
	    --resize 256

添加了可选的 resolution 参数，但图像始终会被裁剪为正方形，并调整大小为 resize x resize 以减少延迟。

安装

已在 Windows 10、Python 3.7.8、CUDA 10.2 环境下测试通过。 需要以下 Python 包：

numpy                  1.19.1
opencv-python          4.4.0.40
Pillow                 7.2.0
PyYAML                 5.3.1
scikit-image           0.17.2
scipy                  1.5.2
tensorflow 	       1.15.3 (tensorflow 仅在 logger.py 中使用，我将很快移除这个非必要的依赖)
torch                  1.6.0
torchvision            0.7.0

时间一致性 [可选]

本节为可选内容。它描述了有助于保持结果视频序列时间连贯性的步骤。本节中的所有示例命令和构建脚本均假设在 Windows 环境下运行；然而，在 Linux/MacOS 上构建和运行应该非常简单直接。

由于我们的技术没有显式地强制时间一致性，这带来了许多优势，例如并行处理、快速训练等，但风格化后的结果序列可能包含令人不适的闪烁。虽然时间不一致可能由多种因素引起，但下面我们将讨论如何处理其中最关键的两个因素。

输入序列中的噪声

相机捕捉到的输入视频序列通常包含一定量的时间噪声。虽然这种噪声可能肉眼不可见或看似微不足道，但网络往往会将其放大。为了解决这个问题，我们建议使用时域双边滤波器对输入序列进行滤波。

首先，需要计算光流。使用 _tools/disflow 中的光流工具。请参阅下面的 构建 disflow 部分了解如何构建该工具。 一旦 disflow.exe 构建完成并位于 PATH 环境变量中，查看并修改 _tools/tool_disflow.py 的前几行，然后运行它。它会从 input 文件夹读取 PNG 图像，并将光流数据存储在 flow_fwd 和 flow_bwd 文件夹中。

计算完光流后，使用时域双边滤波器工具 _tools/bilateralAdv 对序列进行滤波。请参阅下面的 构建 bilateralAdv 部分了解如何构建该工具。 一旦 bilateralAdv.exe 构建完成并位于 PATH 环境变量中，查看并修改 _tools/tool_bilateralAdv.py 的前几行，然后运行它。它会从 input 文件夹读取 PNG 图像，并从 flow_fwd 和 flow_bwd 读取光流数据；它将滤波后的序列存储在 input_filtered 文件夹中。 注意，可以随意并行化 _tools/tool_bilateralAdv.py 中的 for 循环，bilateralAdv.exe 使用光流并且可以逐帧独立运行。同时，也欢迎优化 bilateralAdv.exe 以使其使用多 CPU 核心甚至 GPU……我非常期待看到你的拉取请求 :-)

最后，进行训练和推理时，请使用滤波后的 input_filtered 图像，而不是原始的含噪声的 input 图像。希望结果在时间上会更加稳定。

训练数据中的歧义性

由于网络是在较小的（默认为32x32像素）图像块上进行训练的，来自input RGB帧的多个32x32像素块很可能非常相似。例如，如果输入图像的背景中有天空，那么来自天空左右部分的图像块很可能非常相似。问题在于，在风格化的示例中，这些图像块可能会被风格化得略有不同。这就是歧义性所在：在训练过程中，多个相似的输入图像块会被映射到不同的风格化图像块。为了解决这个问题，我们建议使用辅助的RGB输入图像，使所有输入图像块都具有独特性。

首先，需要计算光流。使用_tools/disflow中的光流工具。关于如何构建该工具，请参见下面的构建disflow部分。一旦disflow.exe构建完成并存在于PATH中，请查看并修改_tools/tool_disflow.py的前几行，然后运行它。它会从input文件夹读取PNG图像，并将光流数据存储在flow_fwd和flow_bwd文件夹中。

光流计算完成后，使用_tools/gauss来计算辅助的高斯混合图像。关于如何构建该工具，请参见下面的构建gauss部分。一旦gauss.exe构建完成并存在于PATH中，请查看并修改_tools/tool_gauss.py的前几行，然后运行它。它会从mask文件夹读取掩码图像（这些掩码可以与训练时使用的掩码匹配，但不是必须的，更多信息请参见下面的部分），并从flow_fwd和flow_bwd读取光流数据；然后输出两种不同的高斯混合图像到input_gdisko_gauss_r10_s10（较小的圆圈）和input_gdisko_gauss_r10_s15（较大的圆圈）。选择其中一个，例如input_gdisko_gauss_r10_s10，如果效果不佳，可以尝试另一个。将input_gdisko_gauss_r10_s10文件夹放置在input文件夹旁边，在_gen和_train文件夹中都放置一份。在_train文件夹中，input_gdisko_gauss_r10_s10将仅包含与风格化关键帧对应的帧，例如Maruska640序列的001.png，或Zuzka2序列的000.png、030.png、070.png和103.png。训练时，不要忘记在运行train.py脚本时使用正确的配置文件，例如--config "_config/reference_P_disco1010.yaml"。运行推理脚本generate.py时，使用可选参数--dir_x1 input_gdisko_gauss_r10_s10，这会告诉generate.py从input_gdisko_gauss_r10_s10加载图像。

高斯混合的掩码

运行gauss.exe时，会为每个掩码图像生成高斯混合，并通过光流传播到整个序列。如果提供了多个掩码图像，生成的高斯圆圈会叠加在一起（并覆盖潜在的空白区域）。掩码可以是（并且在大多数情况下应该是）全白的图像。如果不确定选择哪些帧作为掩码，可以选择与关键帧相同的帧，或者序列的第一帧和最后一帧。查看高斯混合的结果，例如input_gdisko_gauss_r10_s10，如果存在较大的黑色区域（大于100x100像素），请在黑色区域最大的帧上添加一个额外的掩码图像。

构建时间一致性工具

构建disflow

在Windows上，尝试使用预构建的disflow.exe。否则，使用_tools/disflow/build_win.bat自行构建disflow.exe（在Linux/MacOS上，可以参考构建脚本，构建过程应该非常简单）。由于它链接到OpenCV-4.2.0，因此需要在PATH中包含opencv_world420.dll。下载OpenCV-4.2.0，他们提供了预构建的Win包。可以自由修改构建脚本以使用不同版本的OpenCV。注意，OpenCV头文件已提供，位于_tools\disflow\opencv-4.2.0\include，Windows的.lib文件已提供，位于_tools\disflow\opencv-4.2.0\lib。

构建bilateralAdv

在Windows上，尝试使用预构建的bilateralAdv.exe。否则，使用_tools/bilateralAdv/build_win.bat自行构建bilateralAdv.exe（在Linux/MacOS上，可以参考构建脚本，构建过程应该非常简单）。

构建gauss

在Windows上，尝试使用预构建的gauss.exe。否则，使用_tools/gauss/build_win.bat自行构建gauss.exe（在Linux/MacOS上，可以参考构建脚本，构建过程应该非常简单）。

其他实现

• 感谢Midas在PyTorch Lightning中重新实现了这个仓库，参见https://github.com/rnwzd/FSPBT-Image-Translation。

致谢

• 这个项目始于Ondrej Texler在Snap Inc.实习期间，由Snap Inc.和捷克技术大学布拉格分校资助。
• 这个仓库的主要工程贡献者是Ondrej Texler、David Futschik和Michal Kučera。
• 时间一致性工具的主要工程贡献者是Ondrej Jamriska和Sarka Sochorova。

许可证

• 基于图像块的训练方法未申请专利，我们也没有计划申请专利。
• 但是，请注意，这个仓库中的部分代码是Ondrej Texler和David Futschik在Snap Inc.任职期间编写的。如果您发现这个项目对您的商业利益有用，请重新实现它。

引用

如果您发现“基于图像块的少样本训练交互式视频风格化”对您的研究或工作有用，请使用以下BibTeX条目。

@Article{Texler20-SIG,
    author    = "Ond\v{r}ej Texler and David Futschik and Michal Ku\v{c}era and Ond\v{r}ej Jamri\v{s}ka and \v{S}\'{a}rka Sochorov\'{a} and Menglei Chai and Sergey Tulyakov and Daniel S\'{y}kora",
    title     = "Interactive Video Stylization Using Few-Shot Patch-Based Training",
    journal   = "ACM Transactions on Graphics",
    volume    = "39",
    number    = "4",
    pages     = "73",
    year      = "2020",
}

Few-Shot-Patch-Based-Training 快速上手指南

Few-Shot-Patch-Based-Training 是一个用于“基于少量样本和图像块训练的交互式视频风格化”项目的官方实现。本指南将帮助您快速配置环境并运行基础功能。

环境准备

系统要求

• 操作系统: 已在 Windows 10 上测试通过。理论上也支持 Linux/MacOS。
• Python: 3.7.8
• CUDA: 10.2 (如需 GPU 加速)

前置依赖

需要安装以下 Python 包及其指定版本：

numpy                  1.19.1
opencv-python          4.4.0.40
Pillow                 7.2.0
PyYAML                 5.3.1
scikit-image           0.17.2
scipy                  1.5.2
tensorflow             1.15.3  # 仅用于 logger.py，未来将移除
torch                  1.6.0
torchvision            0.7.0

安装步骤

1. 克隆或下载项目代码。
2. 创建并激活 Python 虚拟环境（推荐）。
3. 使用 pip 安装依赖包。建议使用国内镜像源（如清华源）加速下载：

   pip install -r requirements.txt -i https://pypi.tuna.tsinghua.edu.cn/simple
   

   如果项目未提供 requirements.txt 文件，请手动安装上述列出的包及对应版本。

基本使用

1. 下载数据与预训练模型

• 测试数据：下载 testing-data.zip 并解压。确保 _train 文件夹与 _gen 文件夹位于同级目录。
• 预训练模型（可选）：如需快速测试，可下载 pre-trained-models.zip 并解压。

2. 训练模型

运行 train.py 脚本。以下是一个示例命令：

train.py --config "_config/reference_P.yaml" 
         --data_root "Zuzka2_train" 
         --log_interval 1000 
         --log_folder logs_reference_P

此命令会每隔1000轮（log_interval）将当前生成器模型保存到 logs_reference_P 文件夹，并在 _gen 数据上进行验证，结果保存在 Zuzka2_gen/res__P。

3. 生成风格化结果

运行 generate.py 脚本。使用上一步训练得到的检查点文件。

generate.py --checkpoint "Zuzka2_train/logs_reference_P/model_00020.pth" 
            --data_root "Zuzka2_gen"
            --dir_input "input_filtered"
            --outdir "Zuzka2_gen/res_00020" 
            --device "cuda:0"

4. 实时摄像头风格化（可选）

运行 generate_webcam.py 脚本，可将风格化效果实时应用于摄像头画面。按 q 键可退出。

generate_webcam.py --checkpoint "Zuzka2_train/logs_reference_P/model_00020.pth" 
                   --device "cuda:0"
                   --resolution 1280 720
                   --show_original 1
                   --resize 256