SeeSR - 基于语义感知的实景图像超分辨率方法

SeeSR：面向语义感知的真实世界图像超分辨率（CVPR2024）

吴荣源^1,2 | 杨涛³ | 孙凌晨^1,2 | 张政强^1,2 | 李帅^1,2 | 张磊^1,2

¹香港理工大学, ²OPPO研究院, ³字节跳动

:star: 如果SeeSR对您的图像或项目有帮助，请给这个仓库点个星。谢谢！:hugs:

🚩已被CVPR2024接收

📢 新闻

2024.06 我们的一步式真实世界图像超分辨率工作OSEDiff，可达到SeeSR级别的质量，但速度快10倍以上。
2024.03.10 支持sd-turbo，SeeSR只需2步就能得到不错的图像⚡️。请参考这里。
2024.01.12 🔥🔥🔥 已集成到试用Replicate在线演示 ❤️ 感谢lucataco的实现。
2024.01.09 🚀 添加Gradio演示，包括turbo模式。
2023.12.25 🎅🎄🎅🎄 圣诞快乐!!!
- 🍺 发布SeeSR-SD2-Base，包括代码和预训练模型。
- 📏 我们还发布了RealLR200。它包含200张真实世界的低分辨率图像。
2023.11.28 创建此仓库。

📌 待办事项

SeeSR-SDXL
SeeSR-SD2-Base-face,text
~~SeeSR加速~~

🔎 总体框架

seesr

📷 真实世界结果

seesr

⚙️ 依赖项和安装

## 克隆此仓库
git clone https://github.com/cswry/SeeSR.git
cd SeeSR

# 创建一个Python版本 >= 3.8的环境
conda create -n seesr python=3.8
conda activate seesr
pip install -r requirements.txt

🚀 快速推理

步骤1：下载预训练模型

从HuggingFace下载预训练的SD-2-base模型。
从GoogleDrive或OneDrive下载SeeSR和DAPE模型。

您可以将模型放入preset/models目录。

步骤2：准备测试数据

您可以将测试图像放入preset/datasets/test_datasets目录。

步骤3：运行测试命令

python test_seesr.py \
--pretrained_model_path preset/models/stable-diffusion-2-base \
--prompt '' \
--seesr_model_path preset/models/seesr \
--ram_ft_path preset/models/DAPE.pth \
--image_path preset/datasets/test_datasets \
--output_dir preset/datasets/output \
--start_point lr \
--num_inference_steps 50 \
--guidance_scale 5.5 \
--process_size 512

更多详情请参阅此处

SD-Turbo步骤

只需从sd-turbo下载权重，并将它们放入preset/models。然后，您就可以运行命令了。更多比较可以在这里找到。请注意，在turbo模式下，guidance_scale固定为1.0。

python test_seesr_turbo.py \
--pretrained_model_path preset/models/sd-turbo \
--prompt '' \
--seesr_model_path preset/models/seesr \
--ram_ft_path preset/models/DAPE.pth \
--image_path preset/datasets/test_datasets \
--output_dir preset/datasets/output \
--start_point lr \
--num_inference_steps 2 \
--guidance_scale 1.0 \
--process_size 512

注意

请仔细阅读test_seesr.py中的参数。我们采用了multidiffusion-upscaler-for-automatic1111提出的平铺VAE方法来节省GPU内存。

Gradio演示

请将所有预训练模型放在preset/models目录下，然后运行以下命令以与gradio网站交互。

python gradio_seesr.py

我们还提供了带有sd-turbo的gradio，祝您使用愉快。🤗

python gradio_seesr_turbo.py

seesr

测试基准

我们在GoogleDrive和OneDrive上发布了我们的RealLR200。您可以从StableSR下载RealSR和DRealSR。我们还在GoogleDrive和OneDrive上提供了其副本。至于合成测试集，您可以通过下面描述的合成方法获得。

🌈 训练

步骤1：下载预训练模型

下载预训练的SD-2-base模型和RAM。您可以将它们放入preset/models目录。

步骤2：准备训练数据

我们为训练过程预先准备训练数据对，这会占用一些内存空间但可以节省训练时间。我们使用COCO数据集训练DAPE，使用常见的低级数据集（如DF2K）训练SeeSR。

要为DAPE训练制作配对数据，可以运行：

python utils_data/make_paired_data_DAPE.py \
--gt_path 路径_1 路径_2 ... \
--save_dir preset/datasets/train_datasets/training_for_dape \
--epoch 1

要为SeeSR训练制作配对数据，可以运行：

python utils_data/make_paired_data.py \
--gt_path 路径_1 路径_2 ... \
--save_dir preset/datasets/train_datasets/training_for_dape \
--epoch 1

--gt_path 高清图像的路径。如果有多个高清图像目录，可以设置为路径1 路径2 路径3 ...
--save_dir 配对图像的保存路径
--epoch 你想要制作的轮次数

make_paired_data_DAPE.py和make_paired_data.py的区别在于，make_paired_data_DAPE.py将整个图像调整为512的分辨率，而make_paired_data.py随机裁剪出512分辨率的子图像。

一旦创建了降质数据对，你可以基于它们通过运行utils_data/make_tags.py来生成标签数据。

数据文件夹应该如下所示：

your_training_datasets/
    └── gt
        └── 0000001.png # 高清图像，(512, 512, 3)
        └── ...
    └── lr
        └── 0000001.png # 低分辨率图像，(512, 512, 3)
        └── ...
    └── tag
        └── 0000001.txt # 标签提示
        └── ...

步骤3：训练DAPE

请在basicsr/options/dape.yaml的第13行指定DAPE训练数据路径，然后运行训练命令：

python basicsr/train.py -opt basicsr/options/dape.yaml

你可以修改dape.yaml中的参数以适应你的具体情况，比如GPU数量、批量大小、优化器选择等。更多细节请参考Basicsr中的设置。

步骤4：训练SeeSR

CUDA_VISIBLE_DEVICES="0,1,2,3,4,5,6,7," accelerate launch train_seesr.py \
--pretrained_model_name_or_path="preset/models/stable-diffusion-2-base" \
--output_dir="./experience/seesr" \
--root_folders 'preset/datasets/training_datasets' \
--ram_ft_path 'preset/models/DAPE.pth' \
--enable_xformers_memory_efficient_attention \
--mixed_precision="fp16" \
--resolution=512 \
--learning_rate=5e-5 \
--train_batch_size=2 \
--gradient_accumulation_steps=2 \
--null_text_ratio=0.5 
--dataloader_num_workers=0 \
--checkpointing_steps=10000

--pretrained_model_name_or_path 步骤1中预训练SD模型的路径
--root_folders 步骤2中训练数据集的路径
--ram_ft_path 步骤3中DAPE模型的路径

总批量大小由CUDA_VISIBLE_DEVICES、--train_batch_size和--gradient_accumulation_steps共同决定。如果你的GPU内存有限，可以考虑减小--train_batch_size并增加--gradient_accumulation_steps。

❤️ 致谢

本项目基于diffusers和BasicSR。部分代码来自PASD和RAM。感谢他们出色的工作。我们也向StableSR的开创性工作致敬。

📧 联系方式

如果你有任何问题，请随时联系：rong-yuan.wu@connect.polyu.hk

🎓引用

如果我们的代码对你的研究或工作有帮助，请考虑引用我们的论文。以下是BibTeX引用：

@inproceedings{wu2024seesr,
  title={Seesr: Towards semantics-aware real-world image super-resolution},
  author={Wu, Rongyuan and Yang, Tao and Sun, Lingchen and Zhang, Zhengqiang and Li, Shuai and Zhang, Lei},
  booktitle={Proceedings of the IEEE/CVF conference on computer vision and pattern recognition},
  pages={25456--25467},
  year={2024}
}

🎫 许可证

本项目在Apache 2.0许可证下发布。

统计数据