LLFormer - 高效处理超高清低光照图像的Transformer模型

超高清低光图像增强：基准和基于Transformer的方法（AAAI 2023 -- 口头报告）

Tao Wang，Kaihao Zhang，Tianrun Shen，Wenhan Luo，Bjorn Stenger，Tong Lu

摘要： 随着光学传感器质量的提高，处理大规模图像的需求日益增长。特别是设备捕捉超高清（UHD）图像和视频的能力对图像处理流程提出了新的要求。在本文中，我们考虑了低光图像增强（LLIE）任务，并引入了一个由4K和8K分辨率图像组成的大规模数据库。我们进行了系统的基准测试研究，并对当前的LLIE 算法进行了比较。作为第二个贡献，我们引入了LLFormer，一种基于transformer的低光增强方法。LLFormer的核心组件是基于轴的多头自注意力和跨层注意力融合块，显著将复杂度降低到线性水平。在新数据集以及现有公开数据集上的大量实验表明，LLFormer的性能优于最先进的方法。我们还表明，使用在我们的基准上训练的现有LLIE方法作为预处理步骤显著提高了下游任务的性能，例如在低光条件下的人脸检测。源代码和预训练模型可在https://github.com/TaoWangzj/LLFormer获取。

新闻

2023年1月11日： 我们的论文被选为AAAI 2023的口头报告
2022年12月24日： 预训练模型已发布！
2022年12月23日： 代码已发布！
2022年12月23日： 主页已发布！

本仓库包含我们论文的数据集、代码和预训练模型。请参考我们的项目页面以快速了解项目概况。

用于图像增强的UHDLOL基准

我们创建了一个新的大规模UHD-LLIE数据集，称为UHDLOL，用于对现有LLIE方法的性能进行基准测试并探索UHD-LLIE问题。它由两个子集组成：UHD-LOL4K和UHD-LOL8K。UHD-LOL4K子集包含8,099对4K低光/正常光图像。其中，5,999对图像用于训练，2,100对用于测试。UHD-LOL8K子集包括2,966对8K低光/正常光图像，分为2,029对用于训练，937对用于测试。更多详情请参考我们的项目页面。

网络架构

我们提出了一种新的基于transformer的UHD-LLIE方法。LLFormer的核心设计包括一个基于轴的transformer块和一个跨层注意力融合块。在前者中,基于轴的多头自注意力机制在通道维度上沿高度和宽度轴依次执行自注意力,以降低计算复杂度,而双门控前馈网络采用门控机制更关注有用的特征。跨层注意力融合块在融合不同层特征时学习注意力权重。

定量结果

UHD-LOL数据集上的结果

LOL和MIT-Adobe FiveK数据集上的结果

开始使用

依赖和安装

创建Conda环境

conda create -n LLFormer python=3.7
conda activate LLFormer
conda install pytorch=1.8 torchvision=0.3 cudatoolkit=10.1 -c pytorch
pip install matplotlib scikit-image opencv-python yacs joblib natsort h5py tqdm

克隆仓库

git clone https://github.com/TaoWangzj/LLFormer.git

安装预热调度器

cd LLFormer
cd pytorch-gradual-warmup-lr; python setup.py install; cd ..

数据集

您可以使用以下链接下载数据集

UHD-LOL4K [OneDrive | 百度网盘]
UHD-LOL8K [OneDrive | 百度网盘]
LOL [链接]
MIT-Adobe FiveK [Google云端硬盘 | 百度网盘]

预训练模型

我们提供了在不同数据集上预训练的模型：

在UHD-LOL4K上训练的LLFormer [Google云端硬盘 | 百度网盘]，使用训练配置文件 ./configs/UHD-LOL4K/train/training_UHD_4K.yaml。
在UHD-LOL8K上训练的LLFormer [Google云端硬盘 | 百度网盘]，使用训练配置文件 ./configs/UHD-LOL8K/train/training_UHD_8K.yaml。
在LOL数据集上训练的LLFormer [Google drive | 百度网盘]，使用训练配置文件 ./configs/LOL/train/training_LOL.yaml
在MIT-Adobe FiveK数据集上训练的LLFormer [Google drive | 百度网盘]，使用训练配置文件 ./configs/MIT-Adobe-FiveK/train/training_MIT_5K.yaml。

视觉对比结果

我们提供了表2中所有现有方法在LOL和MIT-Adobe FiveK数据集上的视觉结果

LOL数据集的视觉结果 [OneDrive | 百度网盘]。
MIT-Adobe FiveK数据集的视觉结果 [OneDrive | 百度网盘]。

测试

您可以按如下方式直接测试预训练模型

修改数据集和预训练模型的路径。

# 测试参数
input_dir # 数据路径
result_dir # 结果保存路径
weights # 预训练模型的权重路径

测试LOL和MIT-Adobe FiveK数据集的模型

您需要指定数据路径 input_dir、result_dir 和模型路径 weight_path。然后运行

python test.py --input_dir 您的数据路径 --result_dir 您的保存路径 --weights 权重路径

测试UHD-LOL数据集的模型

您需要指定数据路径 input_dir、result_dir 和模型路径 weight_path。然后运行

python test_UHD.py --input_dir 您的数据路径 --result_dir 您的保存路径 --weights 权重路径

（由于GPU内存限制，我们建议用户对UHD图像使用基于patch的模式进行测试）

训练

下载UHD-LOL训练和测试数据
从UHD-LOL数据集的全分辨率训练图像中生成图像patch

python scripts/extract_subimages_UHD.py

要训练LLFormer，请运行

python train.py -yml_path 您的配置路径

您需要修改配置以适应您自己的训练环境

引用

如果UHDLOL基准和LLFormer对您的研究或工作有帮助，请考虑引用：

@inproceedings{wang2023ultra,
  title={Ultra-high-definition low-light image enhancement: A benchmark and transformer-based method},
  author={Wang, Tao and Zhang, Kaihao and Shen, Tianrun and Luo, Wenhan and Stenger, Bjorn and Lu, Tong},
  booktitle={Proceedings of the AAAI Conference on Artificial Intelligence},
  volume={37},
  number={3},
  pages={2654--2662},
  year={2023}
}
 @inproceedings{zhang2021benchmarking,
      title={Benchmarking ultra-high-definition image super-resolution},
      author={Zhang, Kaihao and Li, Dongxu and Luo, Wenhan and Ren, Wenqi and Stenger, Bjorn and Liu, Wei and Li, Hongdong and Yang, Ming-Hsuan},
      booktitle={ICCV},
      pages={14769--14778},
      year={2021}
    }