SRGAN-PyTorch
概述
本仓库包含了使用生成对抗网络的照片级真实单图像超分辨率的PyTorch逐步实现。
目录
下载权重
下载所有可用的模型权重。
# 下载 `SRGAN_x4-SRGAN_ImageNet.pth.tar` 权重到 `./results/pretrained_models`
$ bash ./scripts/download_weights.sh SRGAN_x4-SRGAN_ImageNet
# 下载 `SRResNet_x4-SRGAN_ImageNet.pth.tar` 权重到 `./results/pretrained_models`
$ bash ./scripts/download_weights.sh SRResNet_x4-SRGAN_ImageNet
# 下载 `DiscriminatorForVGG_x4-SRGAN_ImageNet.pth.tar` 权重到 `./results/pretrained_models`
$ bash ./scripts/download_weights.sh DiscriminatorForVGG_x4-SRGAN_ImageNet
下载数据集
这些训练图像是从ImageNet2012分类数据集的验证部分随机选择的。
$ bash ./scripts/download_datasets.sh SRGAN_ImageNet
这里也方便下载一些常用的测试数据集。
$ bash ./scripts/download_datasets.sh Set5
如何测试和训练
训练和测试只需修改yaml文件即可。
项目中使用Set5作为测试基准,您可以自行修改。
如果您需要测试模型效果,请下载测试数据集。
$ bash ./scripts/download_datasets.sh Set5
测试srgan_x4
$ python3 test.py --config_path ./configs/test/SRGAN_x4-SRGAN_ImageNet-Set5.yaml
测试srresnet_x4
$ python3 test.py --config_path ./configs/test/SRResNet_x4-SRGAN_ImageNet-Set5.yaml
训练srresnet_x4
首先,将数据集图像分割成几个小图像,以减少IO并保持批处理图像大小统一。
$ python3 ./scripts/split_images.py
然后,运行以下命令训练模型
$ python3 train_net.py --config_path ./configs/train/SRResNet_x4-SRGAN_ImageNet.yaml
恢复训练srresnet_x4
修改 ./configs/train/SRResNet_x4-SRGAN_ImageNet.yaml
文件。
- 第33行:将
RESUMED_G_MODEL
更改为./samples/SRResNet_x4-SRGAN_ImageNet/g_epoch_xxx.pth.tar
。
$ python3 train_net.py --config_path ./configs/train/SRResNet_x4-SRGAN_ImageNet.yaml
训练srgan_x4
$ python3 train_gan.py --config_path ./configs/train/SRGAN_x4-SRGAN_ImageNet.yaml
恢复训练srgan_x4
修改 ./configs/train/SRGAN_x4-SRGAN_ImageNet.yaml
文件。
- 第38行:将
PRETRAINED_G_MODEL
更改为./results/SRResNet_x4-SRGAN_ImageNet/g_last.pth.tar
。 - 第40行:将
RESUMED_G_MODEL
更改为./samples/SRGAN_x4-SRGAN_ImageNet/g_epoch_xxx.pth.tar
。 - 第41行:将
RESUMED_D_MODEL
更改为./samples/SRGAN_x4-SRGAN_ImageNet/d_epoch_xxx.pth.tar
。
$ python3 train_gan.py --config_path ./configs/train/SRGAN_x4-SRGAN_ImageNet.yaml
结果
原始论文结果来源:https://arxiv.org/pdf/1609.04802v5.pdf
在下表中,()
中的psnr值表示项目的结果,-
表示未测试。
Set5 | 比例 | SRResNet | SRGAN |
---|---|---|---|
PSNR | 4 | 32.05(32.16) | 29.40(30.67) |
SSIM | 4 | 0.9019(0.8938) | 0.8472(0.8627) |
Set14 | 尺度 | SRResNet | SRGAN |
:-----: | :-----: | :------------------: | :------------------: |
PSNR | 4 | 28.49(28.57) | 26.02(27.12) |
SSIM | 4 | 0.8184(0.7815) | 0.7397(0.7321) |
BSD100 | 尺度 | SRResNet | SRGAN |
---|---|---|---|
PSNR | 4 | 27.58(27.56) | 25.16(26.22) |
SSIM | 4 | 0.7620(0.7367) | 0.6688(0.6867) |
# 如果你没有自己训练模型,可以下载模型权重并进行测试。
$ bash ./scripts/download_weights.sh SRGAN_x4-SRGAN_ImageNet
$ python3 ./inference.py
输入:
输出:
成功构建 `srresnet_x4` 模型。
成功加载 `srresnet_x4` 模型权重 `SRGAN-PyTorch/results/pretrained_models/SRGAN_x4-SRGAN_ImageNet.pth.tar`。
超分辨率图像已保存至 `./figure/sr_comic.png`
贡献
如果你发现了bug,请创建一个GitHub issue,或者更好的是,提交一个pull request。同样,如果你有任何问题,也可以直接在GitHub issues中提出。
我期待看到社区如何使用这些模型!
致谢
使用生成对抗网络的照片级真实单图像超分辨率
Christian Ledig, Lucas Theis, Ferenc Huszar, Jose Caballero, Andrew Cunningham, Alejandro Acosta, Andrew Aitken, Alykhan Tejani, Johannes Totz, Zehan Wang, Wenzhe Shi
摘要
尽管使用更快、更深的卷积神经网络在单图像超分辨率的准确性和速度上取得了突破,但一个核心问题仍然没有得到解决:当我们进行大规模上采样时,如何恢复更细腻的纹理细节?基于优化的超分辨率方法的行为主要由目标函数的选择驱动。最近的工作主要集中在最小化均方重建误差上。虽然得到的估计结果具有较高的峰值信噪比,但它们往往缺乏高频细节,在感知上不令人满意,因为它们无法匹配在更高分辨率下预期的保真度。在本文中,我们提出了SRGAN,一个用于图像超分辨率的生成对抗网络。据我们所知,这是第一个能够推断4倍上采样因子的照片级真实自然图像的框架。为了实现这一目标,我们提出了一个由对抗损失和内容损失组成的感知损失函数。对抗损失使用经过训练的判别器网络来区分超分辨率图像和原始照片级真实图像,从而将我们的解决方案推向自然图像流形。此外,我们使用基于感知相似性而非像素空间相似性的内容损失。我们的深度残差网络能够从公共基准上的严重下采样图像中恢复照片级真实的纹理。广泛的平均意见得分(MOS)测试显示,使用SRGAN在感知质量上获得了巨大的显著提升。使用SRGAN获得的MOS分数比使用任何最先进方法获得的分数更接近原始高分辨率图像的分数。
@InProceedings{srgan,
author = {Christian Ledig, Lucas Theis, Ferenc Huszar, Jose Caballero, Andrew Cunningham, Alejandro Acosta, Andrew Aitken, Alykhan Tejani, Johannes Totz, Zehan Wang, Wenzhe Shi},
title = {Photo-Realistic Single Image Super-Resolution Using a Generative Adversarial Network},
booktitle = {arXiv},
year = {2016}
}