[IJCAI 2023 口头报告] 用于低光照图像增强的金字塔扩散模型

论文 | 项目主页 | 补充材料

用于低光照图像增强的金字塔扩散模型
周德伟, 杨宗鑫, 杨毅
发表于IJCAI'2023

整体框架

定量结果

LOL数据集上的评估

LOL数据集上的评估结果如下

依赖和安装

git clone https://github.com/limuloo/PyDIff.git
cd PyDiff
conda create -n PyDiff python=3.7
conda activate PyDiff
conda install pytorch==1.7.0 torchvision torchaudio cudatoolkit=11.0 -c pytorch
cd BasicSR-light
pip install -r requirements.txt
BASICSR_EXT=True sudo $(which python) setup.py develop
cd ../PyDiff
pip install -r requirements.txt
BASICSR_EXT=True sudo $(which python) setup.py develop

数据集

您可以参考以下链接下载LOL数据集，并按以下方式放置：

PyDiff/
    BasicSR-light/
    PyDiff/
    dataset/
        LOLdataset/
            our485/
            eval15/

预训练模型

您可以参考以下链接下载预训练模型，并按以下方式放置：

PyDiff/
    BasicSR-light/
    PyDiff/
    pretrained_models/
        LOLweights.pth

测试

cd PyDiff/
CUDA_VISIBLE_DEVICES=0 python pydiff/train.py -opt options/infer.yaml

注意：在您自己的数据集上测试时，请将'infer.yaml'中的'use_kind_align'设置为false。 详情请参考 https://github.com/limuloo/PyDIff/issues/6。

训练

使用2个GPU进行训练

如论文所述，如果您拥有2个内存容量为24GB或更高的GPU，建议使用以下命令进行训练。

cd PyDiff/
CUDA_VISIBLE_DEVICES=0,1 python -m torch.distributed.launch --nproc_per_node=2 --master_port=22666 pydiff/train.py -opt options/train_v1.yaml --launcher pytorch

使用单个GPU进行训练

否则，您可以使用以下命令进行训练，这需要1个内存大于等于24GB的GPU。

cd PyDiff/
CUDA_VISIBLE_DEVICES=0 python -m torch.distributed.launch --nproc_per_node=1 --master_port=22666 pydiff/train.py -opt options/train_v2.yaml --launcher pytorch

在自定义低级任务数据集上进行训练

请更新PyDiff/options/train_v3.yaml文件中的以下字段：YOUR_TRAIN_DATASET_GT_ROOT、YOUR_TRAIN_DATASET_INPUT_ROOT、YOUR_EVAL_DATASET_GT_ROOT和YOUR_EVAL_DATASET_INPUT_ROOT。如有需要，还请更新PyDiff/pydiff/data/lol_dataset.py。最后，请使用以下命令开始训练过程：

cd PyDiff/
CUDA_VISIBLE_DEVICES=0,1 python -m torch.distributed.launch --nproc_per_node=2 --master_port=22666 pydiff/train.py -opt options/train_v3.yaml --launcher pytorch

请随意自定义其他参数以满足您的特定需求。要使PyDiff在单个GPU上进行训练，可以参考PyDiff/options/train_v2.yaml文件。

引用

如果您认为我们的工作对您的研究有用，请引用我们的论文

@article{zhou2023pyramid,
  title={Pyramid Diffusion Models For Low-light Image Enhancement},
  author={Zhou, Dewei and Yang, Zongxin and Yang, Yi},
  journal={arXiv preprint arXiv:2305.10028},
  year={2023}
}

致谢

我们的代码部分基于BasicSR构建。感谢他们出色工作的贡献者们。