深度白平衡编辑，CVPR 2020（口头报告）

Mahmoud Afifi^1,2 和 Michael S. Brown¹

¹三星AI中心（SAIC）- 多伦多

²约克大学

口头报告视频

论文参考代码：Deep White-Balance Editing。Mahmoud Afifi 和 Michael S. Brown，CVPR 2020。如果您使用此代码或我们的数据集，请引用我们的论文：

@inproceedings{afifi2020deepWB,
  title={Deep White-Balance Editing},
  author={Afifi, Mahmoud and Brown, Michael S},
  booktitle={Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition},
  year={2020}
}

训练数据

1. 下载 Rendered WB 数据集。

2. 将输入图像和真实标签图像分别复制到同一个目录中。每对输入/标签图像应采用以下命名格式：输入图像为 name_WB_picStyle.png，对应的标签图像为 name_G_AS.png。这与 Rendered WB 数据集 中使用的文件名风格一致。请参阅 dataset 目录以获取示例。

代码

我们提供了 Matlab 和 PyTorch 平台的源代码。无法保证训练好的模型会生成完全相同的结果。

1. Matlab（推荐）

前提条件

1. Matlab 2019b 或更高版本
2. 深度学习工具箱

使用说明

运行 install_.m

示例：

1. 运行 demo_single_image.m 或 demo_images.m 分别处理单张图像或图像目录。可用任务包括 AWB、全部和编辑。如果运行 demo_single_image.m，结果将保存在 ../result_images 中，并输出如下图像：

2. 运行 demo_GUI.m 进行 GUI 演示。

训练代码：

运行 training.m 开始训练。在运行代码之前，您需要调整 datasetDir 变量中的训练图像目录。您可以在训练前修改 training.m 中的训练设置。

例如，您可以使用 epochs 和 miniBatch 变量分别更改训练轮数和小批量大小。如果将 fold = 0 和 trainingImgsNum = 0 设置，则训练将使用所有训练数据，不进行交叉验证。如果您希望将训练图像数量限制为 n 张，请将 trainingImgsNum 设置为 n。若要进行 3 折交叉验证，可将 fold 设置为 testing_fold。此时代码将在剩余折数上进行训练，并将选定的一折用于测试。

其他有用选项包括：patchsPerImg 用于选择每张图像的随机补丁数量，以及 patchSize 用于设置训练补丁的大小。要控制学习率下降的速率和因子，请查看位于 utilities 目录中的 get_training_options.m 函数。您可以使用 loadpath 变量从 .mat 格式的训练检查点继续训练。若要从头开始训练，请将 loadpath=[];。

训练开始后，会在 reports_and_checkpoints 目录下生成一个 .csv 文件，可用于可视化训练进度。如果您使用带有图形界面的 Matlab，并希望在训练过程中可视化部分输入/输出补丁，请在 此处 设置断点，并在命令窗口中输入以下代码：

close all; i = 1; figure; subplot(2,3,1);imshow(extractdata(Y(:,:,1:3,i))); subplot(2,3,2);imshow(extractdata(Y(:,:,4:6,i))); subplot(2,3,3);imshow(extractdata(Y(:,:,7:9,i))); subplot(2,3,4); imshow(gather(T(:,:,1:3,i))); subplot(2,3,5); imshow(gather(T(:,:,4:6,i))); subplot(2,3,6); imshow(gather(T(:,:,7:9,i)));

您可以通过更改上述代码中的 i 值来查看当前批次中的不同图像。该图将显示生成的补丁（第一行）及其对应的标签补丁（第二行）。对于无图形界面的情况，您可以编辑自定义代码 此处 来定期保存示例补丁。提示：您可能需要使用 持久变量 来控制这一过程。另一种方法是使用 自定义训练循环。

2. PyTorch

前置条件

1. Python 3.6

2. PyTorch（已测试版本：1.2.0 和 1.5.0）

3. torchvision（已测试版本：0.4.0 和 0.6.0）

4. cudatoolkit

5. tensorboard（可选）

6. numpy

7. Pillow

8. future

9. tqdm

10. matplotlib

11. scipy

12. scikit-learn

代码可能在指定版本之外的库版本上也能运行。

开始使用

示例：

1. 运行 demo_single_image.py 处理单张图片。 应用 AWB + 不同白平衡设置的示例：python demo_single_image.py --input_image ../example_images/00.jpg --output_image ../result_images --show。此示例应在 ../result_images 中保存输出图像，并显示如下图像：

2. 运行 demo_images.py 处理图片目录。示例：python demo_images.py --input_dir ../example_images/ --output_dir ../result_images --task AWB。可用任务包括 AWB、all 和 editing。您也可以在 demo_single_image.py 示例中指定任务。

训练代码：

运行 training.py 开始训练。在运行代码之前，您应调整训练图像目录。

示例：CUDA_VISIBLE_DEVICE=0 python train.py --training_dir ../dataset/ --fold 0 --epochs 500 --learning-rate-drop-period 50 --num_training_images 0。在此示例中，fold = 0 和 num_training_images = 0 表示训练将使用所有训练数据，不进行交叉验证。如果您希望将训练图像数量限制为 n 张，请将 num_training_images 设置为 n。如果您希望进行 3 折交叉验证，请使用 fold = testing_fold。此时，代码将在剩余的折中进行训练，并将选定的折用于测试。

其他有用的选项包括：--patches-per-image 用于选择每张图像的随机补丁数量；--learning-rate-drop-period 和 --learning-rate-drop-factor 分别用于控制学习率下降的周期和因子；以及 --patch-size 用于设置训练补丁的大小。您可以使用 --load 选项从 .pth 训练检查点文件继续训练。

如果您机器上安装了 TensorBoard，在开始训练后运行 tensorboard --logdir ./runs 可以查看训练进度并可视化输入/输出补丁的样本。

结果

本软件仅用于研究目的，不得用于商业用途。

维护者：Mahmoud Afifi（m.3afifi@gmail.com）

相关研究项目

• 当颜色恒常性出错时：首次直接解决白平衡错误问题的工作；内存开销小且速度快（CVPR 2019）。
• 白平衡增强器：一种基于相机白平衡误差的增强技术（ICCV 2019）。
• 交互式白平衡：一种简单的方法，将非线性白平衡校正与用户选择的颜色关联起来，从而实现交互式的白平衡调整（CIC 2020）。
• 曝光校正：一种采用对抗训练的单一粗到细深度学习模型，用于同时校正过度曝光和欠曝光的照片（CVPR 2021）。

Deep_White_Balance 快速上手指南

Deep White-Balance Editing 是 CVPR 2020 Oral 论文开源项目，用于对图像进行自动白平衡（AWB）校正及交互式白平衡编辑。本项目提供 Matlab 和 PyTorch 两个版本的代码，推荐优先使用 Matlab 版本以获得更稳定的复现效果。

环境准备

方案一：Matlab 版本（推荐）

• 操作系统：Windows / Linux / macOS
• 软件要求：Matlab 2019b 或更高版本
• 工具箱依赖：Deep Learning Toolbox

方案二：PyTorch 版本

• 操作系统：Linux / Windows (需配置 CUDA)
• 语言环境：Python 3.6+
• 核心依赖：
  • PyTorch (测试版本：1.2.0, 1.5.0)
  • torchvision (测试版本：0.4.0, 0.6.0)
  • cudatoolkit
• 其他库：numpy, Pillow, future, tqdm, matplotlib, scipy, scikit-learn, tensorboard (可选)

国内加速建议：安装 Python 依赖时，建议使用清华或阿里镜像源加速下载： pip install -r requirements.txt -i https://pypi.tuna.tsinghua.edu.cn/simple

安装步骤

1. 获取代码与数据

克隆仓库后，首先需要下载训练数据集（Rendered WB dataset）。

• 数据集下载地址：Rendered WB dataset
• 数据整理：将输入图像和真值图像放入同一目录。文件名需遵循特定格式：
  • 输入图像：name_WB_picStyle.png
  • 真值图像：name_G_AS.png
  • 参考示例请查看项目中的 dataset 目录结构。

2. 初始化环境

Matlab 用户： 在 Matlab 中运行以下脚本完成路径配置：

install_.m

PyTorch 用户： 确保已安装上述列出的 Python 依赖包。若未安装，可手动执行：

pip install torch torchvision numpy Pillow future tqdm matplotlib scipy scikit-learn

基本使用

以下示例展示如何对单张图片进行白平衡处理。

场景一：使用 Matlab (推荐)

运行演示脚本处理单张图片或整个文件夹。支持的任务模式包括：AWB (自动白平衡), all, editing。

处理单张图片：

demo_single_image.m

执行后结果将保存至 ../result_images 目录，并弹出对比图。

启动图形界面 (GUI) 进行交互式编辑：

demo_GUI.m

场景二：使用 PyTorch

处理单张图片： 以下命令将对输入图片应用 AWB 及多种白平衡设置，并显示结果。

python demo_single_image.py --input_image ../example_images/00.jpg --output_image ../result_images --show

处理整个图片目录：

python demo_images.py --input_dir ../example_images/ --output_dir ../result_images --task AWB

注：--task 参数可选值为 AWB, all, editing。

---

声明：本软件仅供学术研究使用，不可用于商业用途。