🔥 新闻

• 2024.11.30: 发布更便捷的推理代码用于处理您自己的图片。
• 2024.10.25: 发布分割与检测代码、预训练模型。
• 2024.10.25: 发布 RealLQ250 基准测试，其中包含 250 张真实世界的低质量（LQ）图像。
• 2024.10.25: 发布 DreamClear 的训练与推理代码、预训练模型。
• 2024.10.24: 创建本仓库。

📸 真实世界图像恢复（IR）结果

🔧 依赖项与安装

1. 克隆此仓库并进入 DreamClear 文件夹

   git clone https://github.com/shallowdream204/DreamClear.git
   cd DreamClear
   

2. 创建 Conda 环境并安装包

   conda create -n dreamclear python=3.9 -y
   conda activate dreamclear
   pip3 install -r requirements.txt
   

3. 下载预训练模型（除 llava 外，所有模型均可在 Huggingface 方便地下载。）

   基础模型:

   • PixArt-α-1024: PixArt-XL-2-1024-MS.pth
   • VAE (变分自编码器): sd-vae-ft-ema
   • T5 文本编码器: t5-v1_1-xxl
   • LLaVA: llava-v1.6-vicuna-13b
   • SwinIR: general_swinir_v1.ckpt

   我们提供的模型:

   • DreamClear: DreamClear-1024.pth
   • RMT for Segmentation: rmt_uper_s_2x.pth
   • RMT for Detection: rmt_maskrcnn_s_1x.pth

🎰 训练

一、准备训练数据

类似于 SeeSR，我们预先准备好用于 IR（图像恢复）模型训练的高清 - 低清（HQ-LQ）图像对。运行以下命令以生成配对数据用于训练：

python3 tools/make_paired_data.py \
--gt_path gt_path1 gt_path2 ... \ 
--save_dir /path/to/save/folder/ \
--epoch 1 # number of epochs to generate paired data

生成配对数据后，您可以使用 MLLM（多模态大语言模型，例如，LLaVA）为高清图像生成详细的文本提示。然后您需要使用 T5 提取文本特征以节省训练时间。运行：

python3 tools/extract_t5_features.py \
--t5_ckpt /path/to/t5-v1_1-xxl \
--caption_folder /path/to/caption/folder \
--save_npz_folder /path/to/save/npz/folder

最后，训练数据集的目录结构应如下所示

training_datasets_folder/
    └── gt
        └── 0000001.png # GT , (1024, 1024, 3)
        └── ...
    └── sr_bicubic
        └── 0000001.png # LQ + bicubic upsample, (1024, 1024, 3)
        └── ...
    └── caption
        └── 0000001.txt # Caption files (not used in training)
        └── ...
    └── npz
        └── 0000001.npz # T5 features
        └── ...

二、DreamClear 训练

运行以下命令以默认设置训练 DreamClear：

python3 -m torch.distributed.launch --nproc_per_node=8 --nnodes=... --node_rank=... --master_addr=... --master_port=... \
    train_dreamclear.py configs/DreamClear/DreamClear_Train.py \
    --load_from /path/to/PixArt-XL-2-1024-MS.pth \
    --vae_pretrained /path/to/sd-vae-ft-ema \
    --swinir_pretrained /path/to/general_swinir_v1.ckpt \
    --val_image /path/to/RealLQ250/lq/val_image.png \
    --val_npz /path/to/RealLQ250/npz/val_image.npz \
    --work_dir experiments/train_dreamclear

请在 configs/DreamClear/DreamClear_Train.py 中修改训练数据集的路径。您也可以在根据您自己的 GPU 机器在此文件中修改训练超参数（例如，lr, train_batch_size, gradient_accumulation_steps）。

⚡ 推理

我们提供了 RealLQ250 基准测试集，可从 Google Drive 下载。

测试 DreamClear 用于图像恢复

运行以下命令以恢复 LQ (低质量) 图像（代码默认使用 2 个 GPU (图形处理器) 进行推理）：

python3 -m torch.distributed.launch --nproc_per_node 1 --master_port 1234 \
    test.py configs/DreamClear/DreamClear_Test.py \
    --dreamclear_ckpt /path/to/DreamClear-1024.pth \
    --swinir_ckpt /path/to/general_swinir_v1.ckpt \
    --vae_ckpt /path/to/sd-vae-ft-ema \
    --t5_ckpt /path/to/t5-v1_1-xxl \
    --llava_ckpt /path/to/llava-v1.6-vicuna-13b \
    --lre --cfg_scale 4.5 --color_align wavelet \
    --image_path /path/to/input/images \
    --save_dir validation \
    --mixed_precision fp16 \
    --upscale 4

在高级基准测试上的评估

分割 和 检测 的测试说明可在各自的文件夹中找到。

🪪 许可证

提供的代码和预训练权重均遵循 Apache 2.0 许可证。

🤗 致谢

本代码基于 PixArt-α、BasicSR 和 RMT。部分代码源自 SeeSR、StableSR、DiffBIR 和 LLaVA。感谢各位作者的出色工作。

📧 联系方式

如果您有任何问题，欢迎随时通过邮箱 shallowdream555@gmail.com 与我联系。

📖 引用

如果您的研究受益于我们的工作，请考虑引用我们的论文：

@article{ai2024dreamclear,
    title={DreamClear: High-Capacity Real-World Image Restoration with Privacy-Safe Dataset Curation},
    author={Ai, Yuang and Zhou, Xiaoqiang and Huang, Huaibo and Han, Xiaotian and Chen, Zhengyu and You, Quanzeng and Yang, Hongxia},
    journal={Advances in Neural Information Processing Systems},
    volume={37},
    pages={55443--55469},
    year={2024}
}