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

## 模型结构

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

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

## 算法原理
Rayyyyy's avatar
Rayyyyy committed
18
通过使用更实用的退化过程合成训练对, 扩展强大的ESRGAN以恢复一般的真实世界LR图像。
Rayyyyy's avatar
Rayyyyy committed
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
<div align=center>
    <img src="./doc/High-order的pipeline.png"/>
</div>

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

### Docker(方法一)

```bash
docker pull image.sourcefind.cn:5000/dcu/admin/base/pytorch:1.13.1-centos7.6-dtk-23.04.1-py38-latest

docker run -it -v /path/your_code_data/:/path/ your_code_data/ --shm-size=32G --privileged=true --device=/dev/kfd --device=/dev/dri/ --group-add video --name docker_name imageID bash

cd /your_code_path/real-esrgan_pytorch
# 推理和训练都用到了basicsr库
pip install basicsr
# facexlib 和 gfpgan 用于面部增强
pip install facexlib
pip install gfpgan
pip install -r requirements.txt
python setup.py develop
```

### Dockerfile(方法二)

```bash
cd ./docker
cp ../requirements.txt requirements.txt

docker build --no-cache -t real_esrgan:latest .
docker run -it -v /path/your_code_data/:/path/your_code_data/ --shm-size=32G --privileged=true --device=/dev/kfd --device=/dev/dri/ --group-add video --name docker_name imageID bash

cd /your_code_path/real-esrgan_pytorch
# 推理和训练都用到了basicsr库
pip install basicsr
# facexlib 和 gfpgan 用于面部增强
pip install facexlib
pip install gfpgan
pip install -r requirements.txt
python setup.py develop
```

### Anaconda(方法三)

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

```bash
DTK软件栈:dtk23.04.1
python:python3.8
torch:1.13.1
torchvision:0.14.1
```

Tips:以上dtk软件栈、python、torch等DCU相关工具版本需要严格一一对应

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

```bash
# 推理和训练都用到了basicsr库
pip install basicsr
# facexlib 和 gfpgan 用于面部增强
pip install facexlib
pip install gfpgan
pip install -r requirements.txt
python setup.py develop
```

## 数据集
### 准备数据集

所需数据集为: DF2K ( DIV2K 和 Flickr2K ) + OST. 仅需要 HR 图片.

[DIV2K](http://data.vision.ee.ethz.ch/cvl/DIV2K/DIV2K_train_HR.zip)

[Flickr2K](https://cv.snu.ac.kr/research/EDSR/Flickr2K.tar)

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


数据集的目录结构如下:

```bash
├── datasets
│   ├── DF2K
│       ├── DF2K_HR # 将DIV2K和Flickr2K的HR图像放于这里
Rayyyyy's avatar
Rayyyyy committed
105
│       ├── DF2K_HR_sub # 生成的
Rayyyyy's avatar
Rayyyyy committed
106
107
108
109
│       ├── DF2K_multiscale # 生成的
│       ├── DF2K_multiscale_sub # 生成的
│       └── meta_info # 生成的
│   ├── OST
Rayyyyy's avatar
Rayyyyy committed
110
111
│       ├── train_HR # 将OST的HR图像放于这里
│       ├── train_HR_sub # 生成的
Rayyyyy's avatar
Rayyyyy committed
112
113
```

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

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

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

使用 [scripts/generate_multiscale_DF2K.py](scripts/generate_multiscale_DF2K.py) 脚本快速生成多尺寸的图像。

Rayyyyy's avatar
Rayyyyy committed
124
125
126
127
128
129
```bash
# example
python scripts/generate_multiscale_DF2K.py --input datasets/DF2K/DF2K_HR --output datasets/DF2K/DF2K_multiscale
```

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

使用[scripts/extract_subimages.py](scripts/extract_subimages.py) 将 DF2K 图像裁切为子图像, 以加快 IO 和处理速度。

Rayyyyy's avatar
Rayyyyy committed
134
135
```bash
# example
Rayyyyy's avatar
Rayyyyy committed
136
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
137
138
139
```

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

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

```bash
Rayyyyy's avatar
Rayyyyy committed
145
146
147
148
149
150
# meta_info.txt 内容示例
# DF2K_HR_sub/000001_s001.png
# DF2K_HR_sub/000001_s002.png
# DF2K_HR_sub/000001_s003.png

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
151
152
153
154
```

## 训练
### 完整训练
Rayyyyy's avatar
Rayyyyy committed
155
156
157
158
159
完整训练分为两步,第一步训练得到Real-ESRNet, 在Real-ESRNet的基础上训练得到Real-ESRGAN。

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

160
#### 训练Real-ESRNet
Rayyyyy's avatar
Rayyyyy committed
161
1. 下载预训练模型 [ESRGAN](https://github.com/xinntao/Real-ESRGAN/releases/download/v0.1.1/ESRGAN_SRx4_DF2KOST_official-ff704c30.pth), 将模型放到 `experiments/pretrained_models`目录下。
Rayyyyy's avatar
Rayyyyy committed
162
163
164
165
```bash
wget https://github.com/xinntao/Real-ESRGAN/releases/download/v0.1.1/ESRGAN_SRx4_DF2KOST_official-ff704c30.pth -P experiments/pretrained_models
```

166
2. 相应地修改 `options/train_realesrnet_x4plus.yml` 中的内容:
Rayyyyy's avatar
Rayyyyy committed
167
168
169
170
171
172
173
174
175
176
```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
```

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

Rayyyyy's avatar
Rayyyyy committed
179
3. 如果你想在训练过程中执行验证, 就取消注释这些内容并进行相应的修改:
Rayyyyy's avatar
Rayyyyy committed
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
```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
199
    #     psnr: # 指标名称, 可以是任意的
Rayyyyy's avatar
Rayyyyy committed
200
201
202
203
204
    #       type: calculate_psnr
    #       crop_border: 4
    #       test_y_channel: false
```

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

208
209
210
211
2. 大多数修改与上面列出的类似, 相应地修改文件  `options/train_realesrgan_x4plus.yml` 中的内容. 如果需要指定其他文件的预训练路径, 请修改 `pretrain_network_g` 参数.

### 微调Real-ESRGAN
你可以用自己的数据集微调Real-ESRGAN。一般地, 微调(Fine-Tune)程序可以分为两种类型:
Rayyyyy's avatar
Rayyyyy committed
212
213
214
215
216
217

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

#### 动态生成降级图像

Rayyyyy's avatar
Rayyyyy committed
218
只需要高分辨率图像, 在训练过程中, 使用 Real-ESRGAN 描述的降级模型生成低质量图像。
Rayyyyy's avatar
Rayyyyy committed
219
220
221
222
223
224
225
226
227
228
229
230
231

1. 下载预训练模型到 `experiments/pretrained_models` 目录下:

- *RealESRGAN_x4plus.pth*:
    ```bash
    wget https://github.com/xinntao/Real-ESRGAN/releases/download/v0.1.0/RealESRGAN_x4plus.pth -P experiments/pretrained_models
    ```

- *RealESRGAN_x4plus_netD.pth*:
    ```bash
    wget https://github.com/xinntao/Real-ESRGAN/releases/download/v0.2.2.3/RealESRGAN_x4plus_netD.pth -P experiments/pretrained_models
    ```

Rayyyyy's avatar
Rayyyyy committed
232
2. 修改选项文件 [options/finetune_realesrgan_x4plus.yml](options/finetune_realesrgan_x4plus.yml), 特别是 `datasets` 部分:
Rayyyyy's avatar
Rayyyyy committed
233
234
235
236
237
238
239
240
241
242
243
244
245

```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
246
你还可以用自己已经配对的数据微调 RealESRGAN, 这个过程更类似于微调 ESRGAN。
Rayyyyy's avatar
Rayyyyy committed
247
248
249
250
251

1. 数据准备

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

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

Rayyyyy's avatar
Rayyyyy committed
255
然后, 使用脚本 [scripts/generate_meta_info_pairdata.py](scripts/generate_meta_info_pairdata.py) 生成元信息(meta_info)txt 文件。
Rayyyyy's avatar
Rayyyyy committed
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274

```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
```

2. 下载预训练模型

下载所需预训练模型到 `experiments/pretrained_models` 目录下。

- *RealESRGAN_x4plus.pth*:
    ```bash
    wget https://github.com/xinntao/Real-ESRGAN/releases/download/v0.1.0/RealESRGAN_x4plus.pth -P experiments/pretrained_models
    ```

- *RealESRGAN_x4plus_netD.pth*:
    ```bash
    wget https://github.com/xinntao/Real-ESRGAN/releases/download/v0.2.2.3/RealESRGAN_x4plus_netD.pth -P experiments/pretrained_models
    ```

275
3. 微调准备
Rayyyyy's avatar
Rayyyyy committed
276

277
修改 [options/finetune_realesrgan_x4plus_pairdata.yml](options/finetune_realesrgan_x4plus_pairdata.yml) 中的配置项, 特别是 `datasets` 部分:
Rayyyyy's avatar
Rayyyyy committed
278
279
280
281
282
283
284
285
286
287
288

```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
```
289

Rayyyyy's avatar
Rayyyyy committed
290
291
### 训练命令
#### 单机多卡
292
默认auto_resume模式, 根据完整训练或者微调训练, 修改-opt参数对应的yml文件。当前默认为微调训练
Rayyyyy's avatar
Rayyyyy committed
293
```bash
Rayyyyy's avatar
Rayyyyy committed
294
bash train.sh
Rayyyyy's avatar
Rayyyyy committed
295
296
```

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

300
默认auto_resume模式, 根据完整训练或者微调训练, 请修改`$REAL-ESRGAN_ROOT/single_process.sh`中CONFIG参数。
Rayyyyy's avatar
Rayyyyy committed
301

Rayyyyy's avatar
Rayyyyy committed
302
```bash
Rayyyyy's avatar
Rayyyyy committed
303
bash run_train_multi.sh
Rayyyyy's avatar
Rayyyyy committed
304
305
306
```

## 推理
307
下载预训练模型[RealESRGAN_x4plus.pth](https://github.com/xinntao/Real-ESRGAN/releases/download/v0.1.0/RealESRGAN_x4plus.pth), 将其放入 weights 文件夹下, 测试结果默认保存在results文件夹下。
Rayyyyy's avatar
Rayyyyy committed
308

Rayyyyy's avatar
Rayyyyy committed
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
```bash
# 下载预训练模型
wget https://github.com/xinntao/Real-ESRGAN/releases/download/v0.1.0/RealESRGAN_x4plus.pth -P weights
# 执行推理
python inference_realesrgan.py -n RealESRGAN_x4plus -i inputs --face_enhance
```

## result

<div align=center>
    <img src="./doc/00017_gray.jpg"/>
    <img src="./doc/00017_gray_out.jpg"/>
</div>

### 精度

| NIQE | ADE20K val | OST300 |
| :------: | :------: | :------: |
| our | xxx | xxx |
| paper | 3.7886 | 2.8659 |

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

### 热点应用行业
Rayyyyy's avatar
Rayyyyy committed
335
交通,政府,工业
Rayyyyy's avatar
Rayyyyy committed
336
337
338
339
340
341

## 源码仓库及问题反馈
https://developer.hpccube.com/codes/modelzoo/real-esrgan_pytorch

## 参考资料
https://github.com/xinntao/Real-ESRGAN