README.md 11.8 KB
Newer Older
Rayyyyy's avatar
Rayyyyy committed
1
2
# Real-ESRGAN
## 论文
Rayyyyy's avatar
Rayyyyy committed
3
4
`Real-ESRGAN: Training Real-World Blind Super-Resolution with Pure Synthetic Data`
- https://arxiv.org/abs/2107.10833
Rayyyyy's avatar
Rayyyyy committed
5
6
7

## 模型结构

Rayyyyy's avatar
Rayyyyy committed
8
生成网络: 采用ESRGAN的生成网络,对于x4倍的超分辨,网络完全按照ESRGAN的生成器执行;对x2和x1倍的超分辨,网络先进行pixel-unshuffle(pixel-shuffl的反操作,pixel-shuffle可理解为通过压缩图像通道而对图像尺寸进行放大),以降低图像分辨率为前提,对图像通道数进行扩充,然后将处理后的图像输入网络进行超分辨重建。
Rayyyyy's avatar
Rayyyyy committed
9
<div align=center>
Rayyyyy's avatar
Rayyyyy committed
10
    <img src="./doc/ESRGAN.png"/>
Rayyyyy's avatar
Rayyyyy committed
11
12
</div>

Rayyyyy's avatar
Rayyyyy committed
13
对抗网络: 由于使用的复杂的构建数据集的方式,所以需要使用更先进的判别器对生成图像进行判别。使用U-Net判别器可以在像素角度,对单个生成的像素进行真假判断,这能够在保证生成图像整体真实的情况下,注重生成图像细节。
Rayyyyy's avatar
Rayyyyy committed
14
15
16
17
18
<div align=center>
    <img src="./doc/UNet.png"/>
</div>

## 算法原理
Rayyyyy's avatar
Rayyyyy committed
19
通过使用更实用的退化过程合成训练对,扩展强大的ESRGAN以恢复一般的真实世界LR图像。
Rayyyyy's avatar
Rayyyyy committed
20
<div align=center>
Rayyyyy's avatar
Rayyyyy committed
21
    <img src="./doc/pipeline.png"/>
Rayyyyy's avatar
Rayyyyy committed
22
23
24
25
26
27
28
</div>

## 环境配置
-v 路径、docker_name和imageID根据实际情况修改

### Docker(方法一)
```bash
dcuai's avatar
dcuai committed
29
docker pull image.sourcefind.cn:5000/dcu/admin/base/pytorch:2.1.0-ubuntu20.04-dtk24.04.1-py3.8
Rayyyyy's avatar
Rayyyyy committed
30
docker run -it -v /path/your_code_data/:/path/your_code_data/ -v /opt/hyhal/:/opt/hyhal/:ro --shm-size=32G --privileged=true --device=/dev/kfd --device=/dev/dri/ --group-add video --name docker_name imageID bash
Rayyyyy's avatar
Rayyyyy committed
31
32
33
34
35
36
37
38
39
40
41

cd /your_code_path/real-esrgan_pytorch
pip install -r requirements.txt
python setup.py develop
```

### Dockerfile(方法二)
```bash
cd ./docker

docker build --no-cache -t real_esrgan:latest .
Rayyyyy's avatar
Rayyyyy committed
42
docker run -it -v /path/your_code_data/:/path/your_code_data/ -v /opt/hyhal/:/opt/hyhal/:ro --shm-size=32G --privileged=true --device=/dev/kfd --device=/dev/dri/ --group-add video --name docker_name imageID bash
Rayyyyy's avatar
Rayyyyy committed
43
44
45
46
47
48
49
50
51
52
53

cd /your_code_path/real-esrgan_pytorch
pip install -r requirements.txt
python setup.py develop
```

### Anaconda(方法三)

1、关于本项目DCU显卡所需的特殊深度学习库可从光合开发者社区下载安装: https://developer.hpccube.com/tool/

```bash
dcuai's avatar
dcuai committed
54
DTK软件栈:dtk24.04.1
Rayyyyy's avatar
Rayyyyy committed
55
python:python3.8
Rayyyyy's avatar
Rayyyyy committed
56
57
torch:2.1.0
torchvision:0.16.0
Rayyyyy's avatar
Rayyyyy committed
58
59
```

Rayyyyy's avatar
Rayyyyy committed
60
`Tips:以上dtk软件栈、python、torch等DCU相关工具版本需要严格一一对应`
Rayyyyy's avatar
Rayyyyy committed
61
62
63
64
65
66
67
68
69
70

2、其他非特殊库直接按照下面步骤进行安装

```bash
pip install -r requirements.txt
python setup.py develop
```

## 数据集
### 准备数据集
Rayyyyy's avatar
Rayyyyy committed
71
所需数据集为: DF2K(DIV2K和Flickr2K) + OST,仅需要**HR**图片.
Rayyyyy's avatar
Rayyyyy committed
72

Rayyyyy's avatar
Rayyyyy committed
73
[DIV2K](http://113.200.138.88:18080/aidatasets/project-dependency/div2k/-/blob/master/DIV2K_train_HR.zip)
Rayyyyy's avatar
Rayyyyy committed
74

Rayyyyy's avatar
Rayyyyy committed
75
[Flickr2K](http://113.200.138.88:18080/aidatasets/project-dependency/flickr2k/-/blob/master/Flickr2K.tar)
Rayyyyy's avatar
Rayyyyy committed
76
77
78

[OST](https://openmmlab.oss-cn-hangzhou.aliyuncs.com/datasets/OST_dataset.zip)

Rayyyyy's avatar
Rayyyyy committed
79
80
81
82
83
84
85
[ADE20K](https://groups.csail.mit.edu/vision/datasets/ADE20K/)

ADE20K数据需要预处理:

```bash
python gen_20k_val.py --root_path /path/of/ADE20K_2021_17_01 --save_path datasets/ADE20K_val
```
Rayyyyy's avatar
Rayyyyy committed
86

Rayyyyy's avatar
Rayyyyy committed
87
据集的目录结构如下:
Rayyyyy's avatar
Rayyyyy committed
88
89
90
91
92

```bash
├── datasets
│   ├── DF2K
│       ├── DF2K_HR # 将DIV2K和Flickr2K的HR图像放于这里
Rayyyyy's avatar
Rayyyyy committed
93
│       ├── DF2K_HR_sub # 生成的
Rayyyyy's avatar
Rayyyyy committed
94
95
96
97
│       ├── DF2K_multiscale # 生成的
│       ├── DF2K_multiscale_sub # 生成的
│       └── meta_info # 生成的
│   ├── OST
Rayyyyy's avatar
Rayyyyy committed
98
│       ├── train_HR # 将OST的HR图像放于这里
Rayyyyy's avatar
Rayyyyy committed
99
│       ├── train_HR_sub
Rayyyyy's avatar
Rayyyyy committed
100
101
102
103
│   ├── ADE20K_2021_17_01
│       ├── images
│       ├── objects.txt
│   ├── ADE20K_val # 生成的
Rayyyyy's avatar
Rayyyyy committed
104
105
```

Rayyyyy's avatar
Rayyyyy committed
106
Tips: 项目提供了`tiny_datasets`进行快速上手测试,如需使用`tiny_datasets`,下面的数据路径需进行对应修改。
Rayyyyy's avatar
Rayyyyy committed
107

Rayyyyy's avatar
Rayyyyy committed
108
109
### 预处理数据集
#### 1.【可选】生成多尺寸图片
Rayyyyy's avatar
Rayyyyy committed
110
如果只想简单试试而不进行准确训练,那么该过程可选。
Rayyyyy's avatar
Rayyyyy committed
111

Rayyyyy's avatar
Rayyyyy committed
112
针对`DF2K`数据集,我们使用多尺寸缩放策略,对`HR`图像进行下采样,获得多尺寸的标准参考(Ground-Truth)图像。
Rayyyyy's avatar
Rayyyyy committed
113

Rayyyyy's avatar
Rayyyyy committed
114
使用 [scripts/generate_multiscale_DF2K.py](scripts/generate_multiscale_DF2K.py)脚本快速生成多尺寸的图像。
Rayyyyy's avatar
Rayyyyy committed
115

Rayyyyy's avatar
Rayyyyy committed
116
117
118
119
120
121
```bash
# example
python scripts/generate_multiscale_DF2K.py --input datasets/DF2K/DF2K_HR --output datasets/DF2K/DF2K_multiscale
```

#### 2.【可选】裁切为子图像
Rayyyyy's avatar
Rayyyyy committed
122
如果你的IO够好或储存空间有限,那么此步骤是可选的。
Rayyyyy's avatar
Rayyyyy committed
123

Rayyyyy's avatar
Rayyyyy committed
124
使用[scripts/extract_subimages.py](scripts/extract_subimages.py)`DF2K`图像裁切为子图像,以加快IO和处理速度。
Rayyyyy's avatar
Rayyyyy committed
125

Rayyyyy's avatar
Rayyyyy committed
126
127
```bash
# example
Rayyyyy's avatar
Rayyyyy committed
128
python scripts/extract_subimages.py --input datasets/DF2K/DF2K_HR --output datasets/DF2K/DF2K_HR_sub --crop_size 400 --step 200
Rayyyyy's avatar
Rayyyyy committed
129
130
131
```

#### 3. 准备元信息 txt
Rayyyyy's avatar
Rayyyyy committed
132
1. 使用 [scripts/generate_meta_info.py](scripts/generate_meta_info.py) 生成包含图像路径的`txt`文件。
Rayyyyy's avatar
Rayyyyy committed
133

Rayyyyy's avatar
Rayyyyy committed
134
2. 还可以合并多个文件夹的图像路径到一个元信息(meta_info)txt,示例参考如下:
Rayyyyy's avatar
Rayyyyy committed
135
136

```bash
Rayyyyy's avatar
Rayyyyy committed
137
138
139
140
141
# meta_info.txt 内容示例
# DF2K_HR_sub/000001_s001.png
# DF2K_HR_sub/000001_s002.png
# DF2K_HR_sub/000001_s003.png

Rayyyyy's avatar
Rayyyyy committed
142
python scripts/generate_meta_info.py --input datasets/DF2K/DF2K_HR,datasets/DF2K/DF2K_multiscale --root datasets/DF2K,datasets/DF2K --meta_info datasets/DF2K/meta_info/meta_info_DF2Kmultiscale.txt
Rayyyyy's avatar
Rayyyyy committed
143
144
145
146
```

## 训练
### 完整训练
Rayyyyy's avatar
Rayyyyy committed
147
完整训练分为两步,第一步训练得到Real-ESRNet,在Real-ESRNet的基础上训练得到Real-ESRGAN。
Rayyyyy's avatar
Rayyyyy committed
148
149
150
151

1. [训练Real-ESRNet](#训练Real-ESRNet)
2. [训练Real-ESRGAN](#训练Real-ESRGAN)

152
#### 训练Real-ESRNet
Rayyyyy's avatar
Rayyyyy committed
153
1. 下载预训练模型[ESRGAN](http://113.200.138.88:18080/aimodels/findsource-dependency/real-esrgan_pytorch/-/blob/main/ESRGAN_SRx4_DF2KOST_official-ff704c30.pth),将模型放到`experiments/pretrained_models`目录下。
Rayyyyy's avatar
Rayyyyy committed
154

Rayyyyy's avatar
Rayyyyy committed
155
2. 相应地修改`options/train_realesrnet_x4plus.yml`中的内容:
Rayyyyy's avatar
Rayyyyy committed
156
157
158
159
160
161
162
163
164
165
```yml
train:
    name: DF2K+OST
    type: RealESRGANDataset
    dataroot_gt: datasets/DF2K  # 修改为你的数据集文件夹根目录
    meta_info: realesrgan/meta_info/meta_info_DF2Kmultiscale+OST_sub.txt  # 修改为你自己生成的元信息txt
    io_backend:
        type: disk
```

Rayyyyy's avatar
Rayyyyy committed
166
如果需要指定预训练路径到其他文件,请修改`pretrain_network_g`参数,当前默认保存为`experiments/train_RealESRNetx4plus_1000k_B12G4_fromESRGAN/models/net_g_1000000.pth`
Rayyyyy's avatar
Rayyyyy committed
167

Rayyyyy's avatar
Rayyyyy committed
168
3. 如果你想在训练过程中执行验证,就取消注释这些内容并进行相应的修改:
Rayyyyy's avatar
Rayyyyy committed
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
```yml
    # 取消注释这些以进行验证
    # val:
    #   name: validation
    #   type: PairedImageDataset
    #   dataroot_gt: path_to_gt
    #   dataroot_lq: path_to_lq
    #   io_backend:
    #     type: disk

...

    # 取消注释这些以进行验证
    # 验证设置
    # val:
    #   val_freq: !!float 5e3
    #   save_img: True

    #   metrics:
Rayyyyy's avatar
Rayyyyy committed
188
    #     psnr: # 指标名称,可以是任意的
Rayyyyy's avatar
Rayyyyy committed
189
190
191
192
193
    #       type: calculate_psnr
    #       crop_border: 4
    #       test_y_channel: false
```

194
195
#### 训练Real-ESRGAN
1. Real-ESRNet模型训练完成后,得到`experiments/train_RealESRNetx4plus_1000k_B12G4_fromESRGAN/models/net_g_1000000.pth`;
Rayyyyy's avatar
Rayyyyy committed
196

Rayyyyy's avatar
Rayyyyy committed
197
2. 大多数修改与上面列出的类似,相应地修改文件`options/train_realesrgan_x4plus.yml`中的内容. 如果需要指定其他文件的预训练路径,请修改`pretrain_network_g`参数.
198
199

### 微调Real-ESRGAN
Rayyyyy's avatar
Rayyyyy committed
200
你可以用自己的数据集微调`Real-ESRGAN`。一般地,微调程序可以分为两种类型:
Rayyyyy's avatar
Rayyyyy committed
201
202
203
204
205

1. [动态生成降级图像](#动态生成降级图像)
2. [使用**已配对**的数据](#使用已配对的数据)

#### 动态生成降级图像
Rayyyyy's avatar
Rayyyyy committed
206
只需要高分辨率图像,在训练过程中,使用`Real-ESRGAN`描述的降级模型生成低质量图像。
Rayyyyy's avatar
Rayyyyy committed
207

Rayyyyy's avatar
Rayyyyy committed
208
1. 下载预训练模型[RealESRGAN_x4plus.pth](http://113.200.138.88:18080/aimodels/findsource-dependency/real-esrgan_pytorch/-/blob/main/RealESRGAN_x4plus.pth)[RealESRGAN_x4plus_netD.pth](http://113.200.138.88:18080/aimodels/findsource-dependency/real-esrgan_pytorch/-/blob/main/RealESRGAN_x4plus_netD.pth)`experiments/pretrained_models`目录下;
Rayyyyy's avatar
Rayyyyy committed
209

Rayyyyy's avatar
Rayyyyy committed
210
2. 修改选项文件 [options/finetune_realesrgan_x4plus.yml](options/finetune_realesrgan_x4plus.yml),特别是`datasets`部分:
Rayyyyy's avatar
Rayyyyy committed
211
212
213
214
215
216
217
218
219
220
221
222

```yml
train:
    name: DF2K+OST
    type: RealESRGANDataset
    dataroot_gt: datasets/DF2K   # 修改为你的数据集文件夹根目录
    meta_info: realesrgan/meta_info/meta_info_DF2Kmultiscale+OST_sub.txt  # 修改为你自己生成的元信息txt
    io_backend:
        type: disk
```

#### 使用已配对的数据
Rayyyyy's avatar
Rayyyyy committed
223
你还可以用自己已经配对的数据微调`RealESRGAN`,这个过程更类似于微调`ESRGAN`
Rayyyyy's avatar
Rayyyyy committed
224

Rayyyyy's avatar
Rayyyyy committed
225
1. 数据准备:
Rayyyyy's avatar
Rayyyyy committed
226
227
228

假设你已经有两个文件夹(folder):

Rayyyyy's avatar
Rayyyyy committed
229
230
- **gt folder**(标准参考,高分辨率图像):*datasets/DF2K/DIV2K_train_HR_sub*
- **lq folder**(低质量,低分辨率图像):*datasets/DF2K/DIV2K_train_LR_bicubic_X4_sub*
Rayyyyy's avatar
Rayyyyy committed
231

Rayyyyy's avatar
Rayyyyy committed
232
然后,使用脚本 [scripts/generate_meta_info_pairdata.py](scripts/generate_meta_info_pairdata.py) 生成元信息(meta_info)txt 文件。
Rayyyyy's avatar
Rayyyyy committed
233
234
235
236
237

```bash
python scripts/generate_meta_info_pairdata.py --input datasets/DF2K/DIV2K_train_HR_sub datasets/DF2K/DIV2K_train_LR_bicubic_X4_sub --meta_info datasets/DF2K/meta_info/meta_info_DIV2K_sub_pair.txt
```

Rayyyyy's avatar
Rayyyyy committed
238
2. 下载所需预训练模型[RealESRGAN_x4plus.pth](http://113.200.138.88:18080/aimodels/findsource-dependency/real-esrgan_pytorch/-/blob/main/RealESRGAN_x4plus.pth)[RealESRGAN_x4plus_netD.pth](http://113.200.138.88:18080/aimodels/findsource-dependency/real-esrgan_pytorch/-/blob/main/RealESRGAN_x4plus_netD.pth)`experiments/pretrained_models`目录下。
Rayyyyy's avatar
Rayyyyy committed
239
240


241
3. 微调准备
Rayyyyy's avatar
Rayyyyy committed
242

Rayyyyy's avatar
Rayyyyy committed
243
修改 [options/finetune_realesrgan_x4plus_pairdata.yml](options/finetune_realesrgan_x4plus_pairdata.yml)中的配置项,特别是`datasets`部分:
Rayyyyy's avatar
Rayyyyy committed
244
245
246
247
248
249
250
251
252
253
254

```yml
train:
    name: DIV2K
    type: RealESRGANPairedDataset
    dataroot_gt: datasets/DF2K  # 修改为你的 gt folder 文件夹根目录
    dataroot_lq: datasets/DF2K  # 修改为你的 lq folder 文件夹根目录
    meta_info: datasets/DF2K/meta_info/meta_info_DIV2K_sub_pair.txt  # 修改为你自己生成的元信息txt
    io_backend:
        type: disk
```
255

Rayyyyy's avatar
Rayyyyy committed
256
257
### 训练命令
#### 单机多卡
Rayyyyy's avatar
Rayyyyy committed
258
默认auto_resume模式,根据完整训练或者微调训练,修改-opt参数对应的yml文件。当前默认为微调训练
Rayyyyy's avatar
Rayyyyy committed
259
```bash
Rayyyyy's avatar
Rayyyyy committed
260
bash train.sh
Rayyyyy's avatar
Rayyyyy committed
261
262
```

Rayyyyy's avatar
Rayyyyy committed
263
#### 多机多卡
Rayyyyy's avatar
Rayyyyy committed
264
使用多节点的情况下,需要将使用节点写入hostfile文件,多节点每个节点一行,例如: c1xxxxxx slots=4。
Rayyyyy's avatar
Rayyyyy committed
265

Rayyyyy's avatar
Rayyyyy committed
266
默认`auto_resume`模式,根据完整训练或者微调训练,请修改[single_process.sh](./single_process.sh)`CONFIG`参数。
Rayyyyy's avatar
Rayyyyy committed
267

Rayyyyy's avatar
Rayyyyy committed
268
```bash
Rayyyyy's avatar
Rayyyyy committed
269
bash run_train_multi.sh
Rayyyyy's avatar
Rayyyyy committed
270
271
272
```

## 推理
Rayyyyy's avatar
Rayyyyy committed
273
下载预训练模型[RealESRGAN_x4plus.pth](http://113.200.138.88:18080/aimodels/findsource-dependency/real-esrgan_pytorch/-/blob/main/RealESRGAN_x4plus.pth),将其放入`weights`文件夹下,测试结果默认保存在`results`文件夹下。
Rayyyyy's avatar
Rayyyyy committed
274

Rayyyyy's avatar
Rayyyyy committed
275
```bash
Rayyyyy's avatar
Rayyyyy committed
276
277
278
279
280
# 执行推理
python inference_realesrgan.py -n RealESRGAN_x4plus -i inputs --face_enhance
```

## result
Rayyyyy's avatar
Rayyyyy committed
281
Model: RealESRGAN_x4plus
Rayyyyy's avatar
Rayyyyy committed
282
283
284
285
286
287
<div align=center>
    <img src="./doc/00017_gray.jpg"/>
    <img src="./doc/00017_gray_out.jpg"/>
</div>

### 精度
Rayyyyy's avatar
Rayyyyy committed
288
本项目基于ADE20K公开数据集的val数据进行效果验证,需先要对`ADE20K val`数据进行推理,得到推理后的结果(默认在results下)计算 `NIQE` 得分。
Rayyyyy's avatar
Rayyyyy committed
289
290
291
292
293
294

```bash
# 执行推理
python inference_realesrgan.py --model_path weights/RealESRGAN_x4plus.pth -n RealESRGAN_x4plus -i dataset/ADE20K_val --face_enhance
python evalution.py --root_path results
```
Rayyyyy's avatar
Rayyyyy committed
295

296
| NIQE | ADE20K |
Rayyyyy's avatar
Rayyyyy committed
297
| :------: | :------: |
298
299
| V100S | 4.0358 |
| Z100L | 4.0314 |
Rayyyyy's avatar
Rayyyyy committed
300
301
302
303
304
305

## 应用场景
### 算法类别
图像超分

### 热点应用行业
Rayyyyy's avatar
Rayyyyy committed
306
交通,政府,制造
Rayyyyy's avatar
Rayyyyy committed
307
308

## 源码仓库及问题反馈
Rayyyyy's avatar
Rayyyyy committed
309
- https://developer.hpccube.com/codes/modelzoo/real-esrgan_pytorch
Rayyyyy's avatar
Rayyyyy committed
310
311

## 参考资料
Rayyyyy's avatar
Rayyyyy committed
312
- https://github.com/xinntao/Real-ESRGAN