Spaces:
Sleeping
Sleeping
WhosMeor
commited on
Commit
·
cae212d
1
Parent(s):
b57b9d7
upload files
Browse filesThis view is limited to 50 files because it contains too many changes.
See raw diff
- README.md +282 -4
- VERSION +1 -0
- app.py +153 -0
- deploy.bat +7 -0
- deploy_install.bat +7 -0
- gfpgan/weights/parsing_parsenet.pth +3 -0
- inputs/00003.png +0 -0
- inputs/00017_gray.png +0 -0
- inputs/0014.jpg +0 -0
- inputs/0030.jpg +0 -0
- inputs/ADE_val_00000114.jpg +0 -0
- inputs/OST_009.png +0 -0
- inputs/children-alpha.png +0 -0
- inputs/tree_alpha_16bit.png +0 -0
- inputs/video/onepiece_demo.mp4 +0 -0
- inputs/wolf_gray.jpg +0 -0
- realesrgan.egg-info/PKG-INFO +298 -0
- realesrgan.egg-info/SOURCES.txt +21 -0
- realesrgan.egg-info/dependency_links.txt +1 -0
- realesrgan.egg-info/not-zip-safe +1 -0
- realesrgan.egg-info/requires.txt +9 -0
- realesrgan.egg-info/top_level.txt +1 -0
- realesrgan/__init__.py +6 -0
- realesrgan/__pycache__/__init__.cpython-38.pyc +0 -0
- realesrgan/__pycache__/utils.cpython-38.pyc +0 -0
- realesrgan/__pycache__/version.cpython-38.pyc +0 -0
- realesrgan/archs/__init__.py +10 -0
- realesrgan/archs/__pycache__/__init__.cpython-38.pyc +0 -0
- realesrgan/archs/__pycache__/discriminator_arch.cpython-38.pyc +0 -0
- realesrgan/archs/__pycache__/srvgg_arch.cpython-38.pyc +0 -0
- realesrgan/archs/discriminator_arch.py +67 -0
- realesrgan/archs/srvgg_arch.py +69 -0
- realesrgan/data/__init__.py +10 -0
- realesrgan/data/__pycache__/__init__.cpython-38.pyc +0 -0
- realesrgan/data/__pycache__/realesrgan_dataset.cpython-38.pyc +0 -0
- realesrgan/data/__pycache__/realesrgan_paired_dataset.cpython-38.pyc +0 -0
- realesrgan/data/realesrgan_dataset.py +192 -0
- realesrgan/data/realesrgan_paired_dataset.py +108 -0
- realesrgan/models/__init__.py +10 -0
- realesrgan/models/__pycache__/__init__.cpython-38.pyc +0 -0
- realesrgan/models/__pycache__/realesrgan_model.cpython-38.pyc +0 -0
- realesrgan/models/__pycache__/realesrnet_model.cpython-38.pyc +0 -0
- realesrgan/models/realesrgan_model.py +258 -0
- realesrgan/models/realesrnet_model.py +188 -0
- realesrgan/train.py +11 -0
- realesrgan/utils.py +313 -0
- realesrgan/version.py +5 -0
- requirements.txt +11 -0
- run.bat +7 -0
- setup.py +107 -0
README.md
CHANGED
@@ -1,8 +1,8 @@
|
|
1 |
---
|
2 |
-
title:
|
3 |
-
emoji:
|
4 |
-
colorFrom:
|
5 |
-
colorTo:
|
6 |
sdk: streamlit
|
7 |
sdk_version: 1.35.0
|
8 |
app_file: app.py
|
@@ -11,3 +11,281 @@ license: unknown
|
|
11 |
---
|
12 |
|
13 |
Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
---
|
2 |
+
title: Real ESRGAN Web App Cpu Test
|
3 |
+
emoji: 💻
|
4 |
+
colorFrom: green
|
5 |
+
colorTo: gray
|
6 |
sdk: streamlit
|
7 |
sdk_version: 1.35.0
|
8 |
app_file: app.py
|
|
|
11 |
---
|
12 |
|
13 |
Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
|
14 |
+
|
15 |
+
|
16 |
+
<p align="center">
|
17 |
+
<img src="assets/realesrgan_logo.png" height=120>
|
18 |
+
</p>
|
19 |
+
|
20 |
+
## <div align="center"><b><a href="README.md">English</a> | <a href="README_CN.md">简体中文</a></b></div>
|
21 |
+
|
22 |
+
<div align="center">
|
23 |
+
|
24 |
+
👀[**Demos**](#-demos-videos) **|** 🚩[**Updates**](#-updates) **|** ⚡[**Usage**](#-quick-inference) **|** 🏰[**Model Zoo**](docs/model_zoo.md) **|** 🔧[Install](#-dependencies-and-installation) **|** 💻[Train](docs/Training.md) **|** ❓[FAQ](docs/FAQ.md) **|** 🎨[Contribution](docs/CONTRIBUTING.md)
|
25 |
+
|
26 |
+
[](https://github.com/xinntao/Real-ESRGAN/releases)
|
27 |
+
[](https://pypi.org/project/realesrgan/)
|
28 |
+
[](https://github.com/xinntao/Real-ESRGAN/issues)
|
29 |
+
[](https://github.com/xinntao/Real-ESRGAN/issues)
|
30 |
+
[](https://github.com/xinntao/Real-ESRGAN/blob/master/LICENSE)
|
31 |
+
[](https://github.com/xinntao/Real-ESRGAN/blob/master/.github/workflows/pylint.yml)
|
32 |
+
[](https://github.com/xinntao/Real-ESRGAN/blob/master/.github/workflows/publish-pip.yml)
|
33 |
+
|
34 |
+
</div>
|
35 |
+
|
36 |
+
🔥 **AnimeVideo-v3 model (动漫视频小模型)**. Please see [[*anime video models*](docs/anime_video_model.md)] and [[*comparisons*](docs/anime_comparisons.md)]<br>
|
37 |
+
🔥 **RealESRGAN_x4plus_anime_6B** for anime images **(动漫插图模型)**. Please see [[*anime_model*](docs/anime_model.md)]
|
38 |
+
|
39 |
+
<!-- 1. You can try in our website: [ARC Demo](https://arc.tencent.com/en/ai-demos/imgRestore) (now only support RealESRGAN_x4plus_anime_6B) -->
|
40 |
+
1. :boom: **Update** online Replicate demo: [](https://replicate.com/xinntao/realesrgan)
|
41 |
+
1. Online Colab demo for Real-ESRGAN: [](https://colab.research.google.com/drive/1k2Zod6kSHEvraybHl50Lys0LerhyTMCo?usp=sharing) **|** Online Colab demo for for Real-ESRGAN (**anime videos**): [](https://colab.research.google.com/drive/1yNl9ORUxxlL4N0keJa2SEPB61imPQd1B?usp=sharing)
|
42 |
+
1. Portable [Windows](https://github.com/xinntao/Real-ESRGAN/releases/download/v0.2.5.0/realesrgan-ncnn-vulkan-20220424-windows.zip) / [Linux](https://github.com/xinntao/Real-ESRGAN/releases/download/v0.2.5.0/realesrgan-ncnn-vulkan-20220424-ubuntu.zip) / [MacOS](https://github.com/xinntao/Real-ESRGAN/releases/download/v0.2.5.0/realesrgan-ncnn-vulkan-20220424-macos.zip) **executable files for Intel/AMD/Nvidia GPU**. You can find more information [here](#portable-executable-files-ncnn). The ncnn implementation is in [Real-ESRGAN-ncnn-vulkan](https://github.com/xinntao/Real-ESRGAN-ncnn-vulkan)
|
43 |
+
<!-- 1. You can watch enhanced animations in [Tencent Video](https://v.qq.com/s/topic/v_child/render/fC4iyCAM.html). 欢迎观看[腾讯视频动漫修复](https://v.qq.com/s/topic/v_child/render/fC4iyCAM.html) -->
|
44 |
+
|
45 |
+
Real-ESRGAN aims at developing **Practical Algorithms for General Image/Video Restoration**.<br>
|
46 |
+
We extend the powerful ESRGAN to a practical restoration application (namely, Real-ESRGAN), which is trained with pure synthetic data.
|
47 |
+
|
48 |
+
🌌 Thanks for your valuable feedbacks/suggestions. All the feedbacks are updated in [feedback.md](docs/feedback.md).
|
49 |
+
|
50 |
+
---
|
51 |
+
|
52 |
+
If Real-ESRGAN is helpful, please help to ⭐ this repo or recommend it to your friends 😊 <br>
|
53 |
+
Other recommended projects:<br>
|
54 |
+
▶️ [GFPGAN](https://github.com/TencentARC/GFPGAN): A practical algorithm for real-world face restoration <br>
|
55 |
+
▶️ [BasicSR](https://github.com/xinntao/BasicSR): An open-source image and video restoration toolbox<br>
|
56 |
+
▶️ [facexlib](https://github.com/xinntao/facexlib): A collection that provides useful face-relation functions.<br>
|
57 |
+
▶️ [HandyView](https://github.com/xinntao/HandyView): A PyQt5-based image viewer that is handy for view and comparison <br>
|
58 |
+
▶️ [HandyFigure](https://github.com/xinntao/HandyFigure): Open source of paper figures <br>
|
59 |
+
|
60 |
+
---
|
61 |
+
|
62 |
+
### 📖 Real-ESRGAN: Training Real-World Blind Super-Resolution with Pure Synthetic Data
|
63 |
+
|
64 |
+
> [[Paper](https://arxiv.org/abs/2107.10833)]   [[YouTube Video](https://www.youtube.com/watch?v=fxHWoDSSvSc)]   [[B站讲解](https://www.bilibili.com/video/BV1H34y1m7sS/)]   [[Poster](https://xinntao.github.io/projects/RealESRGAN_src/RealESRGAN_poster.pdf)]   [[PPT slides](https://docs.google.com/presentation/d/1QtW6Iy8rm8rGLsJ0Ldti6kP-7Qyzy6XL/edit?usp=sharing&ouid=109799856763657548160&rtpof=true&sd=true)]<br>
|
65 |
+
> [Xintao Wang](https://xinntao.github.io/), Liangbin Xie, [Chao Dong](https://scholar.google.com.hk/citations?user=OSDCB0UAAAAJ), [Ying Shan](https://scholar.google.com/citations?user=4oXBp9UAAAAJ&hl=en) <br>
|
66 |
+
> [Tencent ARC Lab](https://arc.tencent.com/en/ai-demos/imgRestore); Shenzhen Institutes of Advanced Technology, Chinese Academy of Sciences
|
67 |
+
|
68 |
+
<p align="center">
|
69 |
+
<img src="assets/teaser.jpg">
|
70 |
+
</p>
|
71 |
+
|
72 |
+
---
|
73 |
+
|
74 |
+
<!---------------------------------- Updates --------------------------->
|
75 |
+
## 🚩 Updates
|
76 |
+
|
77 |
+
- ✅ Add the **realesr-general-x4v3** model - a tiny small model for general scenes. It also supports the **-dn** option to balance the noise (avoiding over-smooth results). **-dn** is short for denoising strength.
|
78 |
+
- ✅ Update the **RealESRGAN AnimeVideo-v3** model. Please see [anime video models](docs/anime_video_model.md) and [comparisons](docs/anime_comparisons.md) for more details.
|
79 |
+
- ✅ Add small models for anime videos. More details are in [anime video models](docs/anime_video_model.md).
|
80 |
+
- ✅ Add the ncnn implementation [Real-ESRGAN-ncnn-vulkan](https://github.com/xinntao/Real-ESRGAN-ncnn-vulkan).
|
81 |
+
- ✅ Add [*RealESRGAN_x4plus_anime_6B.pth*](https://github.com/xinntao/Real-ESRGAN/releases/download/v0.2.2.4/RealESRGAN_x4plus_anime_6B.pth), which is optimized for **anime** images with much smaller model size. More details and comparisons with [waifu2x](https://github.com/nihui/waifu2x-ncnn-vulkan) are in [**anime_model.md**](docs/anime_model.md)
|
82 |
+
- ✅ Support finetuning on your own data or paired data (*i.e.*, finetuning ESRGAN). See [here](docs/Training.md#Finetune-Real-ESRGAN-on-your-own-dataset)
|
83 |
+
- ✅ Integrate [GFPGAN](https://github.com/TencentARC/GFPGAN) to support **face enhancement**.
|
84 |
+
- ✅ Integrated to [Huggingface Spaces](https://huggingface.co/spaces) with [Gradio](https://github.com/gradio-app/gradio). See [Gradio Web Demo](https://huggingface.co/spaces/akhaliq/Real-ESRGAN). Thanks [@AK391](https://github.com/AK391)
|
85 |
+
- ✅ Support arbitrary scale with `--outscale` (It actually further resizes outputs with `LANCZOS4`). Add *RealESRGAN_x2plus.pth* model.
|
86 |
+
- ✅ [The inference code](inference_realesrgan.py) supports: 1) **tile** options; 2) images with **alpha channel**; 3) **gray** images; 4) **16-bit** images.
|
87 |
+
- ✅ The training codes have been released. A detailed guide can be found in [Training.md](docs/Training.md).
|
88 |
+
|
89 |
+
---
|
90 |
+
|
91 |
+
<!---------------------------------- Demo videos --------------------------->
|
92 |
+
## 👀 Demos Videos
|
93 |
+
|
94 |
+
#### Bilibili
|
95 |
+
|
96 |
+
- [大闹天宫片段](https://www.bilibili.com/video/BV1ja41117zb)
|
97 |
+
- [Anime dance cut 动漫魔性舞蹈](https://www.bilibili.com/video/BV1wY4y1L7hT/)
|
98 |
+
- [海贼王片段](https://www.bilibili.com/video/BV1i3411L7Gy/)
|
99 |
+
|
100 |
+
#### YouTube
|
101 |
+
|
102 |
+
## 🔧 Dependencies and Installation
|
103 |
+
|
104 |
+
- Python >= 3.7 (Recommend to use [Anaconda](https://www.anaconda.com/download/#linux) or [Miniconda](https://docs.conda.io/en/latest/miniconda.html))
|
105 |
+
- [PyTorch >= 1.7](https://pytorch.org/)
|
106 |
+
|
107 |
+
### Installation
|
108 |
+
|
109 |
+
1. Clone repo
|
110 |
+
|
111 |
+
```bash
|
112 |
+
git clone https://github.com/xinntao/Real-ESRGAN.git
|
113 |
+
cd Real-ESRGAN
|
114 |
+
```
|
115 |
+
|
116 |
+
1. Install dependent packages
|
117 |
+
|
118 |
+
```bash
|
119 |
+
# Install basicsr - https://github.com/xinntao/BasicSR
|
120 |
+
# We use BasicSR for both training and inference
|
121 |
+
pip install basicsr
|
122 |
+
# facexlib and gfpgan are for face enhancement
|
123 |
+
pip install facexlib
|
124 |
+
pip install gfpgan
|
125 |
+
pip install -r requirements.txt
|
126 |
+
python setup.py develop
|
127 |
+
```
|
128 |
+
|
129 |
+
---
|
130 |
+
|
131 |
+
## ⚡ Quick Inference
|
132 |
+
|
133 |
+
There are usually three ways to inference Real-ESRGAN.
|
134 |
+
|
135 |
+
1. [Online inference](#online-inference)
|
136 |
+
1. [Portable executable files (NCNN)](#portable-executable-files-ncnn)
|
137 |
+
1. [Python script](#python-script)
|
138 |
+
|
139 |
+
### Online inference
|
140 |
+
|
141 |
+
1. You can try in our website: [ARC Demo](https://arc.tencent.com/en/ai-demos/imgRestore) (now only support RealESRGAN_x4plus_anime_6B)
|
142 |
+
1. [Colab Demo](https://colab.research.google.com/drive/1k2Zod6kSHEvraybHl50Lys0LerhyTMCo?usp=sharing) for Real-ESRGAN **|** [Colab Demo](https://colab.research.google.com/drive/1yNl9ORUxxlL4N0keJa2SEPB61imPQd1B?usp=sharing) for Real-ESRGAN (**anime videos**).
|
143 |
+
|
144 |
+
### Portable executable files (NCNN)
|
145 |
+
|
146 |
+
You can download [Windows](https://github.com/xinntao/Real-ESRGAN/releases/download/v0.2.5.0/realesrgan-ncnn-vulkan-20220424-windows.zip) / [Linux](https://github.com/xinntao/Real-ESRGAN/releases/download/v0.2.5.0/realesrgan-ncnn-vulkan-20220424-ubuntu.zip) / [MacOS](https://github.com/xinntao/Real-ESRGAN/releases/download/v0.2.5.0/realesrgan-ncnn-vulkan-20220424-macos.zip) **executable files for Intel/AMD/Nvidia GPU**.
|
147 |
+
|
148 |
+
This executable file is **portable** and includes all the binaries and models required. No CUDA or PyTorch environment is needed.<br>
|
149 |
+
|
150 |
+
You can simply run the following command (the Windows example, more information is in the README.md of each executable files):
|
151 |
+
|
152 |
+
```bash
|
153 |
+
./realesrgan-ncnn-vulkan.exe -i input.jpg -o output.png -n model_name
|
154 |
+
```
|
155 |
+
|
156 |
+
We have provided five models:
|
157 |
+
|
158 |
+
1. realesrgan-x4plus (default)
|
159 |
+
2. realesrnet-x4plus
|
160 |
+
3. realesrgan-x4plus-anime (optimized for anime images, small model size)
|
161 |
+
4. realesr-animevideov3 (animation video)
|
162 |
+
|
163 |
+
You can use the `-n` argument for other models, for example, `./realesrgan-ncnn-vulkan.exe -i input.jpg -o output.png -n realesrnet-x4plus`
|
164 |
+
|
165 |
+
#### Usage of portable executable files
|
166 |
+
|
167 |
+
1. Please refer to [Real-ESRGAN-ncnn-vulkan](https://github.com/xinntao/Real-ESRGAN-ncnn-vulkan#computer-usages) for more details.
|
168 |
+
1. Note that it does not support all the functions (such as `outscale`) as the python script `inference_realesrgan.py`.
|
169 |
+
|
170 |
+
```console
|
171 |
+
Usage: realesrgan-ncnn-vulkan.exe -i infile -o outfile [options]...
|
172 |
+
|
173 |
+
-h show this help
|
174 |
+
-i input-path input image path (jpg/png/webp) or directory
|
175 |
+
-o output-path output image path (jpg/png/webp) or directory
|
176 |
+
-s scale upscale ratio (can be 2, 3, 4. default=4)
|
177 |
+
-t tile-size tile size (>=32/0=auto, default=0) can be 0,0,0 for multi-gpu
|
178 |
+
-m model-path folder path to the pre-trained models. default=models
|
179 |
+
-n model-name model name (default=realesr-animevideov3, can be realesr-animevideov3 | realesrgan-x4plus | realesrgan-x4plus-anime | realesrnet-x4plus)
|
180 |
+
-g gpu-id gpu device to use (default=auto) can be 0,1,2 for multi-gpu
|
181 |
+
-j load:proc:save thread count for load/proc/save (default=1:2:2) can be 1:2,2,2:2 for multi-gpu
|
182 |
+
-x enable tta mode"
|
183 |
+
-f format output image format (jpg/png/webp, default=ext/png)
|
184 |
+
-v verbose output
|
185 |
+
```
|
186 |
+
|
187 |
+
Note that it may introduce block inconsistency (and also generate slightly different results from the PyTorch implementation), because this executable file first crops the input image into several tiles, and then processes them separately, finally stitches together.
|
188 |
+
|
189 |
+
### Python script
|
190 |
+
|
191 |
+
#### Usage of python script
|
192 |
+
|
193 |
+
1. You can use X4 model for **arbitrary output size** with the argument `outscale`. The program will further perform cheap resize operation after the Real-ESRGAN output.
|
194 |
+
|
195 |
+
```console
|
196 |
+
Usage: python inference_realesrgan.py -n RealESRGAN_x4plus -i infile -o outfile [options]...
|
197 |
+
|
198 |
+
A common command: python inference_realesrgan.py -n RealESRGAN_x4plus -i infile --outscale 3.5 --face_enhance
|
199 |
+
|
200 |
+
-h show this help
|
201 |
+
-i --input Input image or folder. Default: inputs
|
202 |
+
-o --output Output folder. Default: results
|
203 |
+
-n --model_name Model name. Default: RealESRGAN_x4plus
|
204 |
+
-s, --outscale The final upsampling scale of the image. Default: 4
|
205 |
+
--suffix Suffix of the restored image. Default: out
|
206 |
+
-t, --tile Tile size, 0 for no tile during testing. Default: 0
|
207 |
+
--face_enhance Whether to use GFPGAN to enhance face. Default: False
|
208 |
+
--fp32 Use fp32 precision during inference. Default: fp16 (half precision).
|
209 |
+
--ext Image extension. Options: auto | jpg | png, auto means using the same extension as inputs. Default: auto
|
210 |
+
```
|
211 |
+
|
212 |
+
#### Inference general images
|
213 |
+
|
214 |
+
Download pre-trained models: [RealESRGAN_x4plus.pth](https://github.com/xinntao/Real-ESRGAN/releases/download/v0.1.0/RealESRGAN_x4plus.pth)
|
215 |
+
|
216 |
+
```bash
|
217 |
+
wget https://github.com/xinntao/Real-ESRGAN/releases/download/v0.1.0/RealESRGAN_x4plus.pth -P weights
|
218 |
+
```
|
219 |
+
|
220 |
+
Inference!
|
221 |
+
|
222 |
+
```bash
|
223 |
+
python inference_realesrgan.py -n RealESRGAN_x4plus -i inputs --face_enhance
|
224 |
+
```
|
225 |
+
|
226 |
+
Results are in the `results` folder
|
227 |
+
|
228 |
+
#### Inference anime images
|
229 |
+
|
230 |
+
<p align="center">
|
231 |
+
<img src="https://raw.githubusercontent.com/xinntao/public-figures/master/Real-ESRGAN/cmp_realesrgan_anime_1.png">
|
232 |
+
</p>
|
233 |
+
|
234 |
+
Pre-trained models: [RealESRGAN_x4plus_anime_6B](https://github.com/xinntao/Real-ESRGAN/releases/download/v0.2.2.4/RealESRGAN_x4plus_anime_6B.pth)<br>
|
235 |
+
More details and comparisons with [waifu2x](https://github.com/nihui/waifu2x-ncnn-vulkan) are in [**anime_model.md**](docs/anime_model.md)
|
236 |
+
|
237 |
+
```bash
|
238 |
+
# download model
|
239 |
+
wget https://github.com/xinntao/Real-ESRGAN/releases/download/v0.2.2.4/RealESRGAN_x4plus_anime_6B.pth -P weights
|
240 |
+
# inference
|
241 |
+
python inference_realesrgan.py -n RealESRGAN_x4plus_anime_6B -i inputs
|
242 |
+
```
|
243 |
+
|
244 |
+
Results are in the `results` folder
|
245 |
+
|
246 |
+
---
|
247 |
+
|
248 |
+
## BibTeX
|
249 |
+
|
250 |
+
@InProceedings{wang2021realesrgan,
|
251 |
+
author = {Xintao Wang and Liangbin Xie and Chao Dong and Ying Shan},
|
252 |
+
title = {Real-ESRGAN: Training Real-World Blind Super-Resolution with Pure Synthetic Data},
|
253 |
+
booktitle = {International Conference on Computer Vision Workshops (ICCVW)},
|
254 |
+
date = {2021}
|
255 |
+
}
|
256 |
+
|
257 |
+
## 📧 Contact
|
258 |
+
|
259 |
+
If you have any question, please email `[email protected]` or `[email protected]`.
|
260 |
+
|
261 |
+
<!---------------------------------- Projects that use Real-ESRGAN --------------------------->
|
262 |
+
## 🧩 Projects that use Real-ESRGAN
|
263 |
+
|
264 |
+
If you develop/use Real-ESRGAN in your projects, welcome to let me know.
|
265 |
+
|
266 |
+
- NCNN-Android: [RealSR-NCNN-Android](https://github.com/tumuyan/RealSR-NCNN-Android) by [tumuyan](https://github.com/tumuyan)
|
267 |
+
- VapourSynth: [vs-realesrgan](https://github.com/HolyWu/vs-realesrgan) by [HolyWu](https://github.com/HolyWu)
|
268 |
+
- NCNN: [Real-ESRGAN-ncnn-vulkan](https://github.com/xinntao/Real-ESRGAN-ncnn-vulkan)
|
269 |
+
|
270 |
+
**GUI**
|
271 |
+
|
272 |
+
- [Waifu2x-Extension-GUI](https://github.com/AaronFeng753/Waifu2x-Extension-GUI) by [AaronFeng753](https://github.com/AaronFeng753)
|
273 |
+
- [Squirrel-RIFE](https://github.com/Justin62628/Squirrel-RIFE) by [Justin62628](https://github.com/Justin62628)
|
274 |
+
- [Real-GUI](https://github.com/scifx/Real-GUI) by [scifx](https://github.com/scifx)
|
275 |
+
- [Real-ESRGAN_GUI](https://github.com/net2cn/Real-ESRGAN_GUI) by [net2cn](https://github.com/net2cn)
|
276 |
+
- [Real-ESRGAN-EGUI](https://github.com/WGzeyu/Real-ESRGAN-EGUI) by [WGzeyu](https://github.com/WGzeyu)
|
277 |
+
- [anime_upscaler](https://github.com/shangar21/anime_upscaler) by [shangar21](https://github.com/shangar21)
|
278 |
+
- [Upscayl](https://github.com/upscayl/upscayl) by [Nayam Amarshe](https://github.com/NayamAmarshe) and [TGS963](https://github.com/TGS963)
|
279 |
+
|
280 |
+
## 🤗 Acknowledgement
|
281 |
+
|
282 |
+
Thanks for all the contributors.
|
283 |
+
|
284 |
+
- [AK391](https://github.com/AK391): Integrate RealESRGAN to [Huggingface Spaces](https://huggingface.co/spaces) with [Gradio](https://github.com/gradio-app/gradio). See [Gradio Web Demo](https://huggingface.co/spaces/akhaliq/Real-ESRGAN).
|
285 |
+
- [Asiimoviet](https://github.com/Asiimoviet): Translate the README.md to Chinese (中文).
|
286 |
+
- [2ji3150](https://github.com/2ji3150): Thanks for the [detailed and valuable feedbacks/suggestions](https://github.com/xinntao/Real-ESRGAN/issues/131).
|
287 |
+
- [Jared-02](https://github.com/Jared-02): Translate the Training.md to Chinese (中文).
|
288 |
+
|
289 |
+
|
290 |
+
|
291 |
+
|
VERSION
ADDED
@@ -0,0 +1 @@
|
|
|
|
|
1 |
+
0.3.0
|
app.py
ADDED
@@ -0,0 +1,153 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
import streamlit as st
|
2 |
+
import cv2
|
3 |
+
import os
|
4 |
+
import numpy as np
|
5 |
+
from basicsr.archs.rrdbnet_arch import RRDBNet
|
6 |
+
from basicsr.utils.download_util import load_file_from_url
|
7 |
+
from realesrgan import RealESRGANer
|
8 |
+
from realesrgan.archs.srvgg_arch import SRVGGNetCompact
|
9 |
+
from gfpgan import GFPGANer
|
10 |
+
|
11 |
+
# Function to load the model
|
12 |
+
def load_model(model_name, model_path, denoise_strength, tile, tile_pad, pre_pad, fp32, gpu_id):
|
13 |
+
if model_name == 'RealESRGAN_x4plus':
|
14 |
+
model = RRDBNet(num_in_ch=3, num_out_ch=3, num_feat=64, num_block=23, num_grow_ch=32, scale=4)
|
15 |
+
netscale = 4
|
16 |
+
file_url = ['https://github.com/xinntao/Real-ESRGAN/releases/download/v0.1.0/RealESRGAN_x4plus.pth']
|
17 |
+
elif model_name == 'RealESRGAN_x4plus_anime_6B':
|
18 |
+
model = RRDBNet(num_in_ch=3, num_out_ch=3, num_feat=64, num_block=6, num_grow_ch=32, scale=4)
|
19 |
+
netscale = 4
|
20 |
+
file_url = ['https://github.com/xinntao/Real-ESRGAN/releases/download/v0.2.2.4/RealESRGAN_x4plus_anime_6B.pth']
|
21 |
+
elif model_name == 'RealESRGAN_x2plus':
|
22 |
+
model = RRDBNet(num_in_ch=3, num_out_ch=3, num_feat=64, num_block=23, num_grow_ch=32, scale=2)
|
23 |
+
netscale = 2
|
24 |
+
file_url = ['https://github.com/xinntao/Real-ESRGAN/releases/download/v0.2.1/RealESRGAN_x2plus.pth']
|
25 |
+
|
26 |
+
# Determine model paths
|
27 |
+
if model_path is not None:
|
28 |
+
model_path = model_path
|
29 |
+
else:
|
30 |
+
model_path = os.path.join('weights', model_name + '.pth')
|
31 |
+
if not os.path.isfile(model_path):
|
32 |
+
for url in file_url:
|
33 |
+
# Model_path will be updated
|
34 |
+
model_path = load_file_from_url(
|
35 |
+
url=url, model_dir=os.path.join(os.getcwd(), 'weights'), progress=True, file_name=model_name + '.pth')
|
36 |
+
|
37 |
+
dni_weight = None
|
38 |
+
if model_name == 'realesr-general-x4v3' and denoise_strength != 1:
|
39 |
+
model_path = [model_path, model_path.replace('realesr-general-x4v3', 'realesr-general-wdn-x4v3')]
|
40 |
+
dni_weight = [denoise_strength, 1 - denoise_strength]
|
41 |
+
|
42 |
+
# Use DNI to control the denoise strength
|
43 |
+
dni_weight = None
|
44 |
+
if model_name == 'realesr-general-x4v3' and denoise_strength != 1:
|
45 |
+
wdn_model_path = model_path.replace('realesr-general-x4v3', 'realesr-general-wdn-x4v3')
|
46 |
+
model_path = [model_path, wdn_model_path]
|
47 |
+
dni_weight = [denoise_strength, 1 - denoise_strength]
|
48 |
+
|
49 |
+
# Restorer
|
50 |
+
upsampler = RealESRGANer(
|
51 |
+
scale=netscale,
|
52 |
+
model_path=model_path,
|
53 |
+
dni_weight=dni_weight,
|
54 |
+
model=model,
|
55 |
+
tile=tile,
|
56 |
+
tile_pad=tile_pad,
|
57 |
+
pre_pad=pre_pad,
|
58 |
+
half=not fp32,
|
59 |
+
gpu_id=gpu_id)
|
60 |
+
|
61 |
+
return upsampler
|
62 |
+
|
63 |
+
# Function to download model weights if not present
|
64 |
+
def ensure_model_weights(model_name):
|
65 |
+
weights_dir = 'weights'
|
66 |
+
model_file = f"{model_name}.pth"
|
67 |
+
model_path = os.path.join(weights_dir, model_file)
|
68 |
+
|
69 |
+
if not os.path.exists(weights_dir):
|
70 |
+
os.makedirs(weights_dir)
|
71 |
+
|
72 |
+
if not os.path.isfile(model_path):
|
73 |
+
if model_name == 'RealESRGAN_x4plus':
|
74 |
+
file_url = 'https://github.com/xinntao/Real-ESRGAN/releases/download/v0.1.0/RealESRGAN_x4plus.pth'
|
75 |
+
elif model_name == 'RealESRGAN_x4plus_anime_6B':
|
76 |
+
file_url = 'https://github.com/xinntao/Real-ESRGAN/releases/download/v0.2.2.4/RealESRGAN_x4plus_anime_6B.pth'
|
77 |
+
elif model_name == 'RealESRGAN_x2plus':
|
78 |
+
file_url = 'https://github.com/xinntao/Real-ESRGAN/releases/download/v0.2.1/RealESRGAN_x2plus.pth'
|
79 |
+
|
80 |
+
model_path = load_file_from_url(
|
81 |
+
url=file_url, model_dir=weights_dir, progress=True, file_name=model_file)
|
82 |
+
|
83 |
+
return model_path
|
84 |
+
|
85 |
+
# Streamlit app
|
86 |
+
st.title("Real-ESRGAN Image Enhancement")
|
87 |
+
|
88 |
+
uploaded_file = st.file_uploader("Choose an image...", type=["jpg", "png", "jpeg"])
|
89 |
+
|
90 |
+
# User selects model name, denoise strength, and other parameters
|
91 |
+
model_name = st.selectbox("Model Name", ['RealESRGAN_x4plus', 'RealESRGAN_x4plus_anime_6B', 'RealESRGAN_x2plus'])
|
92 |
+
denoise_strength = st.slider("Denoise Strength", 0.0, 1.0, 0.5)
|
93 |
+
outscale = st.slider("Output Scale", 1, 4, 4)
|
94 |
+
tile = 0
|
95 |
+
tile_pad = 10
|
96 |
+
pre_pad = 0
|
97 |
+
face_enhance = st.checkbox("Face Enhance")
|
98 |
+
fp32 = st.checkbox("Use FP32 Precision")
|
99 |
+
gpu_id = None # or set to 0, 1, etc. if you have multiple GPUs
|
100 |
+
|
101 |
+
if uploaded_file is not None:
|
102 |
+
col1, col2 = st.columns(2)
|
103 |
+
with col1:
|
104 |
+
st.write("### Original Image")
|
105 |
+
st.image(uploaded_file, use_column_width=True)
|
106 |
+
run_button = st.button("Run")
|
107 |
+
|
108 |
+
# Save uploaded image to disk
|
109 |
+
input_image_path = os.path.join("temp", "input_image.png")
|
110 |
+
os.makedirs("temp", exist_ok=True)
|
111 |
+
with open(input_image_path, "wb") as f:
|
112 |
+
f.write(uploaded_file.getbuffer())
|
113 |
+
|
114 |
+
if not run_button:
|
115 |
+
st.warning("Click the 'Run' button to start the enhancement process.")
|
116 |
+
|
117 |
+
if run_button:
|
118 |
+
# Ensure model weights are downloaded
|
119 |
+
model_path = ensure_model_weights(model_name)
|
120 |
+
|
121 |
+
# Load the model
|
122 |
+
upsampler = load_model(model_name, model_path, denoise_strength, tile, tile_pad, pre_pad, fp32, gpu_id)
|
123 |
+
|
124 |
+
# Load the image
|
125 |
+
img = cv2.imdecode(np.frombuffer(uploaded_file.read(), np.uint8), cv2.IMREAD_UNCHANGED)
|
126 |
+
if img is None:
|
127 |
+
st.error("Error loading image. Please try again.")
|
128 |
+
else:
|
129 |
+
img_mode = 'RGBA' if len(img.shape) == 3 and img.shape[2] == 4 else None
|
130 |
+
|
131 |
+
try:
|
132 |
+
if face_enhance:
|
133 |
+
face_enhancer = GFPGANer(
|
134 |
+
model_path='https://github.com/TencentARC/GFPGAN/releases/download/v1.3.0/GFPGANv1.3.pth',
|
135 |
+
upscale=outscale,
|
136 |
+
arch='clean',
|
137 |
+
channel_multiplier=2,
|
138 |
+
bg_upsampler=upsampler)
|
139 |
+
_, _, output = face_enhancer.enhance(img, has_aligned=False, only_center_face=False, paste_back=True)
|
140 |
+
else:
|
141 |
+
output, _ = upsampler.enhance(img, outscale=outscale)
|
142 |
+
except RuntimeError as error:
|
143 |
+
st.error(f"Error: {error}")
|
144 |
+
st.error('If you encounter CUDA out of memory, try to set --tile with a smaller number.')
|
145 |
+
else:
|
146 |
+
# Save and display the output image
|
147 |
+
output_image_path = os.path.join("temp", "output_image.png")
|
148 |
+
cv2.imwrite(output_image_path, output)
|
149 |
+
with col2:
|
150 |
+
st.write("### Enhanced Image")
|
151 |
+
st.image(output_image_path, use_column_width=True)
|
152 |
+
if 'output_image_path' in locals():
|
153 |
+
st.download_button("Download Enhanced Image", data=open(output_image_path, "rb").read(), file_name="output_image.png", mime="image/png")
|
deploy.bat
ADDED
@@ -0,0 +1,7 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
@echo off
|
2 |
+
REM Activate the cuda environment
|
3 |
+
call "%USERPROFILE%\anaconda3\Scripts\activate.bat" cuda
|
4 |
+
REM Change directory to Real-ESRGAN-Web-App
|
5 |
+
cd /d %USERPROFILE%\Real-ESRGAN-Web-App
|
6 |
+
REM Run localtunnel on port 8501
|
7 |
+
lt --port 8501
|
deploy_install.bat
ADDED
@@ -0,0 +1,7 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
@echo off
|
2 |
+
REM Activate the cuda environment
|
3 |
+
call "%USERPROFILE%\anaconda3\Scripts\activate.bat" cuda
|
4 |
+
REM Change directory to Real-ESRGAN-Web-App
|
5 |
+
cd /d %USERPROFILE%\Real-ESRGAN-Web-App
|
6 |
+
REM Install localtunnel
|
7 |
+
npm install -g localtunnel
|
gfpgan/weights/parsing_parsenet.pth
ADDED
@@ -0,0 +1,3 @@
|
|
|
|
|
|
|
|
|
1 |
+
version https://git-lfs.github.com/spec/v1
|
2 |
+
oid sha256:3d558d8d0e42c20224f13cf5a29c79eba2d59913419f945545d8cf7b72920de2
|
3 |
+
size 85331193
|
inputs/00003.png
ADDED
![]() |
inputs/00017_gray.png
ADDED
![]() |
inputs/0014.jpg
ADDED
![]() |
inputs/0030.jpg
ADDED
![]() |
inputs/ADE_val_00000114.jpg
ADDED
![]() |
inputs/OST_009.png
ADDED
![]() |
inputs/children-alpha.png
ADDED
![]() |
inputs/tree_alpha_16bit.png
ADDED
![]() |
inputs/video/onepiece_demo.mp4
ADDED
Binary file (593 kB). View file
|
|
inputs/wolf_gray.jpg
ADDED
![]() |
realesrgan.egg-info/PKG-INFO
ADDED
@@ -0,0 +1,298 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
Metadata-Version: 2.1
|
2 |
+
Name: realesrgan
|
3 |
+
Version: 0.3.0
|
4 |
+
Summary: Real-ESRGAN aims at developing Practical Algorithms for General Image Restoration
|
5 |
+
Home-page: https://github.com/xinntao/Real-ESRGAN
|
6 |
+
Author: Xintao Wang
|
7 |
+
Author-email: [email protected]
|
8 |
+
License: BSD-3-Clause License
|
9 |
+
Keywords: computer vision,pytorch,image restoration,super-resolution,esrgan,real-esrgan
|
10 |
+
Classifier: Development Status :: 4 - Beta
|
11 |
+
Classifier: License :: OSI Approved :: Apache Software License
|
12 |
+
Classifier: Operating System :: OS Independent
|
13 |
+
Classifier: Programming Language :: Python :: 3
|
14 |
+
Classifier: Programming Language :: Python :: 3.7
|
15 |
+
Classifier: Programming Language :: Python :: 3.8
|
16 |
+
Description-Content-Type: text/markdown
|
17 |
+
Requires-Dist: basicsr>=1.4.2
|
18 |
+
Requires-Dist: facexlib>=0.2.5
|
19 |
+
Requires-Dist: gfpgan>=1.3.5
|
20 |
+
Requires-Dist: numpy
|
21 |
+
Requires-Dist: opencv-python
|
22 |
+
Requires-Dist: Pillow
|
23 |
+
Requires-Dist: torch>=1.7
|
24 |
+
Requires-Dist: torchvision
|
25 |
+
Requires-Dist: tqdm
|
26 |
+
|
27 |
+
<p align="center">
|
28 |
+
<img src="assets/realesrgan_logo.png" height=120>
|
29 |
+
</p>
|
30 |
+
|
31 |
+
## <div align="center"><b><a href="README.md">English</a> | <a href="README_CN.md">简体中文</a></b></div>
|
32 |
+
|
33 |
+
<div align="center">
|
34 |
+
|
35 |
+
👀[**Demos**](#-demos-videos) **|** 🚩[**Updates**](#-updates) **|** ⚡[**Usage**](#-quick-inference) **|** 🏰[**Model Zoo**](docs/model_zoo.md) **|** 🔧[Install](#-dependencies-and-installation) **|** 💻[Train](docs/Training.md) **|** ❓[FAQ](docs/FAQ.md) **|** 🎨[Contribution](docs/CONTRIBUTING.md)
|
36 |
+
|
37 |
+
[](https://github.com/xinntao/Real-ESRGAN/releases)
|
38 |
+
[](https://pypi.org/project/realesrgan/)
|
39 |
+
[](https://github.com/xinntao/Real-ESRGAN/issues)
|
40 |
+
[](https://github.com/xinntao/Real-ESRGAN/issues)
|
41 |
+
[](https://github.com/xinntao/Real-ESRGAN/blob/master/LICENSE)
|
42 |
+
[](https://github.com/xinntao/Real-ESRGAN/blob/master/.github/workflows/pylint.yml)
|
43 |
+
[](https://github.com/xinntao/Real-ESRGAN/blob/master/.github/workflows/publish-pip.yml)
|
44 |
+
|
45 |
+
</div>
|
46 |
+
|
47 |
+
🔥 **AnimeVideo-v3 model (动漫视频小模型)**. Please see [[*anime video models*](docs/anime_video_model.md)] and [[*comparisons*](docs/anime_comparisons.md)]<br>
|
48 |
+
🔥 **RealESRGAN_x4plus_anime_6B** for anime images **(动漫插图模型)**. Please see [[*anime_model*](docs/anime_model.md)]
|
49 |
+
|
50 |
+
<!-- 1. You can try in our website: [ARC Demo](https://arc.tencent.com/en/ai-demos/imgRestore) (now only support RealESRGAN_x4plus_anime_6B) -->
|
51 |
+
1. :boom: **Update** online Replicate demo: [](https://replicate.com/xinntao/realesrgan)
|
52 |
+
1. Online Colab demo for Real-ESRGAN: [](https://colab.research.google.com/drive/1k2Zod6kSHEvraybHl50Lys0LerhyTMCo?usp=sharing) **|** Online Colab demo for for Real-ESRGAN (**anime videos**): [](https://colab.research.google.com/drive/1yNl9ORUxxlL4N0keJa2SEPB61imPQd1B?usp=sharing)
|
53 |
+
1. Portable [Windows](https://github.com/xinntao/Real-ESRGAN/releases/download/v0.2.5.0/realesrgan-ncnn-vulkan-20220424-windows.zip) / [Linux](https://github.com/xinntao/Real-ESRGAN/releases/download/v0.2.5.0/realesrgan-ncnn-vulkan-20220424-ubuntu.zip) / [MacOS](https://github.com/xinntao/Real-ESRGAN/releases/download/v0.2.5.0/realesrgan-ncnn-vulkan-20220424-macos.zip) **executable files for Intel/AMD/Nvidia GPU**. You can find more information [here](#portable-executable-files-ncnn). The ncnn implementation is in [Real-ESRGAN-ncnn-vulkan](https://github.com/xinntao/Real-ESRGAN-ncnn-vulkan)
|
54 |
+
<!-- 1. You can watch enhanced animations in [Tencent Video](https://v.qq.com/s/topic/v_child/render/fC4iyCAM.html). 欢迎观看[腾讯视频动漫修复](https://v.qq.com/s/topic/v_child/render/fC4iyCAM.html) -->
|
55 |
+
|
56 |
+
Real-ESRGAN aims at developing **Practical Algorithms for General Image/Video Restoration**.<br>
|
57 |
+
We extend the powerful ESRGAN to a practical restoration application (namely, Real-ESRGAN), which is trained with pure synthetic data.
|
58 |
+
|
59 |
+
🌌 Thanks for your valuable feedbacks/suggestions. All the feedbacks are updated in [feedback.md](docs/feedback.md).
|
60 |
+
|
61 |
+
---
|
62 |
+
|
63 |
+
If Real-ESRGAN is helpful, please help to ⭐ this repo or recommend it to your friends 😊 <br>
|
64 |
+
Other recommended projects:<br>
|
65 |
+
▶️ [GFPGAN](https://github.com/TencentARC/GFPGAN): A practical algorithm for real-world face restoration <br>
|
66 |
+
▶️ [BasicSR](https://github.com/xinntao/BasicSR): An open-source image and video restoration toolbox<br>
|
67 |
+
▶️ [facexlib](https://github.com/xinntao/facexlib): A collection that provides useful face-relation functions.<br>
|
68 |
+
▶️ [HandyView](https://github.com/xinntao/HandyView): A PyQt5-based image viewer that is handy for view and comparison <br>
|
69 |
+
▶️ [HandyFigure](https://github.com/xinntao/HandyFigure): Open source of paper figures <br>
|
70 |
+
|
71 |
+
---
|
72 |
+
|
73 |
+
### 📖 Real-ESRGAN: Training Real-World Blind Super-Resolution with Pure Synthetic Data
|
74 |
+
|
75 |
+
> [[Paper](https://arxiv.org/abs/2107.10833)]   [[YouTube Video](https://www.youtube.com/watch?v=fxHWoDSSvSc)]   [[B站讲解](https://www.bilibili.com/video/BV1H34y1m7sS/)]   [[Poster](https://xinntao.github.io/projects/RealESRGAN_src/RealESRGAN_poster.pdf)]   [[PPT slides](https://docs.google.com/presentation/d/1QtW6Iy8rm8rGLsJ0Ldti6kP-7Qyzy6XL/edit?usp=sharing&ouid=109799856763657548160&rtpof=true&sd=true)]<br>
|
76 |
+
> [Xintao Wang](https://xinntao.github.io/), Liangbin Xie, [Chao Dong](https://scholar.google.com.hk/citations?user=OSDCB0UAAAAJ), [Ying Shan](https://scholar.google.com/citations?user=4oXBp9UAAAAJ&hl=en) <br>
|
77 |
+
> [Tencent ARC Lab](https://arc.tencent.com/en/ai-demos/imgRestore); Shenzhen Institutes of Advanced Technology, Chinese Academy of Sciences
|
78 |
+
|
79 |
+
<p align="center">
|
80 |
+
<img src="assets/teaser.jpg">
|
81 |
+
</p>
|
82 |
+
|
83 |
+
---
|
84 |
+
|
85 |
+
<!---------------------------------- Updates --------------------------->
|
86 |
+
## 🚩 Updates
|
87 |
+
|
88 |
+
- ✅ Add the **realesr-general-x4v3** model - a tiny small model for general scenes. It also supports the **-dn** option to balance the noise (avoiding over-smooth results). **-dn** is short for denoising strength.
|
89 |
+
- ✅ Update the **RealESRGAN AnimeVideo-v3** model. Please see [anime video models](docs/anime_video_model.md) and [comparisons](docs/anime_comparisons.md) for more details.
|
90 |
+
- ✅ Add small models for anime videos. More details are in [anime video models](docs/anime_video_model.md).
|
91 |
+
- ✅ Add the ncnn implementation [Real-ESRGAN-ncnn-vulkan](https://github.com/xinntao/Real-ESRGAN-ncnn-vulkan).
|
92 |
+
- ✅ Add [*RealESRGAN_x4plus_anime_6B.pth*](https://github.com/xinntao/Real-ESRGAN/releases/download/v0.2.2.4/RealESRGAN_x4plus_anime_6B.pth), which is optimized for **anime** images with much smaller model size. More details and comparisons with [waifu2x](https://github.com/nihui/waifu2x-ncnn-vulkan) are in [**anime_model.md**](docs/anime_model.md)
|
93 |
+
- ✅ Support finetuning on your own data or paired data (*i.e.*, finetuning ESRGAN). See [here](docs/Training.md#Finetune-Real-ESRGAN-on-your-own-dataset)
|
94 |
+
- ✅ Integrate [GFPGAN](https://github.com/TencentARC/GFPGAN) to support **face enhancement**.
|
95 |
+
- ✅ Integrated to [Huggingface Spaces](https://huggingface.co/spaces) with [Gradio](https://github.com/gradio-app/gradio). See [Gradio Web Demo](https://huggingface.co/spaces/akhaliq/Real-ESRGAN). Thanks [@AK391](https://github.com/AK391)
|
96 |
+
- ✅ Support arbitrary scale with `--outscale` (It actually further resizes outputs with `LANCZOS4`). Add *RealESRGAN_x2plus.pth* model.
|
97 |
+
- ✅ [The inference code](inference_realesrgan.py) supports: 1) **tile** options; 2) images with **alpha channel**; 3) **gray** images; 4) **16-bit** images.
|
98 |
+
- ✅ The training codes have been released. A detailed guide can be found in [Training.md](docs/Training.md).
|
99 |
+
|
100 |
+
---
|
101 |
+
|
102 |
+
<!---------------------------------- Demo videos --------------------------->
|
103 |
+
## 👀 Demos Videos
|
104 |
+
|
105 |
+
#### Bilibili
|
106 |
+
|
107 |
+
- [大闹天宫片段](https://www.bilibili.com/video/BV1ja41117zb)
|
108 |
+
- [Anime dance cut 动漫魔性舞蹈](https://www.bilibili.com/video/BV1wY4y1L7hT/)
|
109 |
+
- [海贼王片段](https://www.bilibili.com/video/BV1i3411L7Gy/)
|
110 |
+
|
111 |
+
#### YouTube
|
112 |
+
|
113 |
+
## 🔧 Dependencies and Installation
|
114 |
+
|
115 |
+
- Python >= 3.7 (Recommend to use [Anaconda](https://www.anaconda.com/download/#linux) or [Miniconda](https://docs.conda.io/en/latest/miniconda.html))
|
116 |
+
- [PyTorch >= 1.7](https://pytorch.org/)
|
117 |
+
|
118 |
+
### Installation
|
119 |
+
|
120 |
+
1. Clone repo
|
121 |
+
|
122 |
+
```bash
|
123 |
+
git clone https://github.com/xinntao/Real-ESRGAN.git
|
124 |
+
cd Real-ESRGAN
|
125 |
+
```
|
126 |
+
|
127 |
+
1. Install dependent packages
|
128 |
+
|
129 |
+
```bash
|
130 |
+
# Install basicsr - https://github.com/xinntao/BasicSR
|
131 |
+
# We use BasicSR for both training and inference
|
132 |
+
pip install basicsr
|
133 |
+
# facexlib and gfpgan are for face enhancement
|
134 |
+
pip install facexlib
|
135 |
+
pip install gfpgan
|
136 |
+
pip install -r requirements.txt
|
137 |
+
python setup.py develop
|
138 |
+
```
|
139 |
+
|
140 |
+
---
|
141 |
+
|
142 |
+
## ⚡ Quick Inference
|
143 |
+
|
144 |
+
There are usually three ways to inference Real-ESRGAN.
|
145 |
+
|
146 |
+
1. [Online inference](#online-inference)
|
147 |
+
1. [Portable executable files (NCNN)](#portable-executable-files-ncnn)
|
148 |
+
1. [Python script](#python-script)
|
149 |
+
|
150 |
+
### Online inference
|
151 |
+
|
152 |
+
1. You can try in our website: [ARC Demo](https://arc.tencent.com/en/ai-demos/imgRestore) (now only support RealESRGAN_x4plus_anime_6B)
|
153 |
+
1. [Colab Demo](https://colab.research.google.com/drive/1k2Zod6kSHEvraybHl50Lys0LerhyTMCo?usp=sharing) for Real-ESRGAN **|** [Colab Demo](https://colab.research.google.com/drive/1yNl9ORUxxlL4N0keJa2SEPB61imPQd1B?usp=sharing) for Real-ESRGAN (**anime videos**).
|
154 |
+
|
155 |
+
### Portable executable files (NCNN)
|
156 |
+
|
157 |
+
You can download [Windows](https://github.com/xinntao/Real-ESRGAN/releases/download/v0.2.5.0/realesrgan-ncnn-vulkan-20220424-windows.zip) / [Linux](https://github.com/xinntao/Real-ESRGAN/releases/download/v0.2.5.0/realesrgan-ncnn-vulkan-20220424-ubuntu.zip) / [MacOS](https://github.com/xinntao/Real-ESRGAN/releases/download/v0.2.5.0/realesrgan-ncnn-vulkan-20220424-macos.zip) **executable files for Intel/AMD/Nvidia GPU**.
|
158 |
+
|
159 |
+
This executable file is **portable** and includes all the binaries and models required. No CUDA or PyTorch environment is needed.<br>
|
160 |
+
|
161 |
+
You can simply run the following command (the Windows example, more information is in the README.md of each executable files):
|
162 |
+
|
163 |
+
```bash
|
164 |
+
./realesrgan-ncnn-vulkan.exe -i input.jpg -o output.png -n model_name
|
165 |
+
```
|
166 |
+
|
167 |
+
We have provided five models:
|
168 |
+
|
169 |
+
1. realesrgan-x4plus (default)
|
170 |
+
2. realesrnet-x4plus
|
171 |
+
3. realesrgan-x4plus-anime (optimized for anime images, small model size)
|
172 |
+
4. realesr-animevideov3 (animation video)
|
173 |
+
|
174 |
+
You can use the `-n` argument for other models, for example, `./realesrgan-ncnn-vulkan.exe -i input.jpg -o output.png -n realesrnet-x4plus`
|
175 |
+
|
176 |
+
#### Usage of portable executable files
|
177 |
+
|
178 |
+
1. Please refer to [Real-ESRGAN-ncnn-vulkan](https://github.com/xinntao/Real-ESRGAN-ncnn-vulkan#computer-usages) for more details.
|
179 |
+
1. Note that it does not support all the functions (such as `outscale`) as the python script `inference_realesrgan.py`.
|
180 |
+
|
181 |
+
```console
|
182 |
+
Usage: realesrgan-ncnn-vulkan.exe -i infile -o outfile [options]...
|
183 |
+
|
184 |
+
-h show this help
|
185 |
+
-i input-path input image path (jpg/png/webp) or directory
|
186 |
+
-o output-path output image path (jpg/png/webp) or directory
|
187 |
+
-s scale upscale ratio (can be 2, 3, 4. default=4)
|
188 |
+
-t tile-size tile size (>=32/0=auto, default=0) can be 0,0,0 for multi-gpu
|
189 |
+
-m model-path folder path to the pre-trained models. default=models
|
190 |
+
-n model-name model name (default=realesr-animevideov3, can be realesr-animevideov3 | realesrgan-x4plus | realesrgan-x4plus-anime | realesrnet-x4plus)
|
191 |
+
-g gpu-id gpu device to use (default=auto) can be 0,1,2 for multi-gpu
|
192 |
+
-j load:proc:save thread count for load/proc/save (default=1:2:2) can be 1:2,2,2:2 for multi-gpu
|
193 |
+
-x enable tta mode"
|
194 |
+
-f format output image format (jpg/png/webp, default=ext/png)
|
195 |
+
-v verbose output
|
196 |
+
```
|
197 |
+
|
198 |
+
Note that it may introduce block inconsistency (and also generate slightly different results from the PyTorch implementation), because this executable file first crops the input image into several tiles, and then processes them separately, finally stitches together.
|
199 |
+
|
200 |
+
### Python script
|
201 |
+
|
202 |
+
#### Usage of python script
|
203 |
+
|
204 |
+
1. You can use X4 model for **arbitrary output size** with the argument `outscale`. The program will further perform cheap resize operation after the Real-ESRGAN output.
|
205 |
+
|
206 |
+
```console
|
207 |
+
Usage: python inference_realesrgan.py -n RealESRGAN_x4plus -i infile -o outfile [options]...
|
208 |
+
|
209 |
+
A common command: python inference_realesrgan.py -n RealESRGAN_x4plus -i infile --outscale 3.5 --face_enhance
|
210 |
+
|
211 |
+
-h show this help
|
212 |
+
-i --input Input image or folder. Default: inputs
|
213 |
+
-o --output Output folder. Default: results
|
214 |
+
-n --model_name Model name. Default: RealESRGAN_x4plus
|
215 |
+
-s, --outscale The final upsampling scale of the image. Default: 4
|
216 |
+
--suffix Suffix of the restored image. Default: out
|
217 |
+
-t, --tile Tile size, 0 for no tile during testing. Default: 0
|
218 |
+
--face_enhance Whether to use GFPGAN to enhance face. Default: False
|
219 |
+
--fp32 Use fp32 precision during inference. Default: fp16 (half precision).
|
220 |
+
--ext Image extension. Options: auto | jpg | png, auto means using the same extension as inputs. Default: auto
|
221 |
+
```
|
222 |
+
|
223 |
+
#### Inference general images
|
224 |
+
|
225 |
+
Download pre-trained models: [RealESRGAN_x4plus.pth](https://github.com/xinntao/Real-ESRGAN/releases/download/v0.1.0/RealESRGAN_x4plus.pth)
|
226 |
+
|
227 |
+
```bash
|
228 |
+
wget https://github.com/xinntao/Real-ESRGAN/releases/download/v0.1.0/RealESRGAN_x4plus.pth -P weights
|
229 |
+
```
|
230 |
+
|
231 |
+
Inference!
|
232 |
+
|
233 |
+
```bash
|
234 |
+
python inference_realesrgan.py -n RealESRGAN_x4plus -i inputs --face_enhance
|
235 |
+
```
|
236 |
+
|
237 |
+
Results are in the `results` folder
|
238 |
+
|
239 |
+
#### Inference anime images
|
240 |
+
|
241 |
+
<p align="center">
|
242 |
+
<img src="https://raw.githubusercontent.com/xinntao/public-figures/master/Real-ESRGAN/cmp_realesrgan_anime_1.png">
|
243 |
+
</p>
|
244 |
+
|
245 |
+
Pre-trained models: [RealESRGAN_x4plus_anime_6B](https://github.com/xinntao/Real-ESRGAN/releases/download/v0.2.2.4/RealESRGAN_x4plus_anime_6B.pth)<br>
|
246 |
+
More details and comparisons with [waifu2x](https://github.com/nihui/waifu2x-ncnn-vulkan) are in [**anime_model.md**](docs/anime_model.md)
|
247 |
+
|
248 |
+
```bash
|
249 |
+
# download model
|
250 |
+
wget https://github.com/xinntao/Real-ESRGAN/releases/download/v0.2.2.4/RealESRGAN_x4plus_anime_6B.pth -P weights
|
251 |
+
# inference
|
252 |
+
python inference_realesrgan.py -n RealESRGAN_x4plus_anime_6B -i inputs
|
253 |
+
```
|
254 |
+
|
255 |
+
Results are in the `results` folder
|
256 |
+
|
257 |
+
---
|
258 |
+
|
259 |
+
## BibTeX
|
260 |
+
|
261 |
+
@InProceedings{wang2021realesrgan,
|
262 |
+
author = {Xintao Wang and Liangbin Xie and Chao Dong and Ying Shan},
|
263 |
+
title = {Real-ESRGAN: Training Real-World Blind Super-Resolution with Pure Synthetic Data},
|
264 |
+
booktitle = {International Conference on Computer Vision Workshops (ICCVW)},
|
265 |
+
date = {2021}
|
266 |
+
}
|
267 |
+
|
268 |
+
## 📧 Contact
|
269 |
+
|
270 |
+
If you have any question, please email `[email protected]` or `[email protected]`.
|
271 |
+
|
272 |
+
<!---------------------------------- Projects that use Real-ESRGAN --------------------------->
|
273 |
+
## 🧩 Projects that use Real-ESRGAN
|
274 |
+
|
275 |
+
If you develop/use Real-ESRGAN in your projects, welcome to let me know.
|
276 |
+
|
277 |
+
- NCNN-Android: [RealSR-NCNN-Android](https://github.com/tumuyan/RealSR-NCNN-Android) by [tumuyan](https://github.com/tumuyan)
|
278 |
+
- VapourSynth: [vs-realesrgan](https://github.com/HolyWu/vs-realesrgan) by [HolyWu](https://github.com/HolyWu)
|
279 |
+
- NCNN: [Real-ESRGAN-ncnn-vulkan](https://github.com/xinntao/Real-ESRGAN-ncnn-vulkan)
|
280 |
+
|
281 |
+
**GUI**
|
282 |
+
|
283 |
+
- [Waifu2x-Extension-GUI](https://github.com/AaronFeng753/Waifu2x-Extension-GUI) by [AaronFeng753](https://github.com/AaronFeng753)
|
284 |
+
- [Squirrel-RIFE](https://github.com/Justin62628/Squirrel-RIFE) by [Justin62628](https://github.com/Justin62628)
|
285 |
+
- [Real-GUI](https://github.com/scifx/Real-GUI) by [scifx](https://github.com/scifx)
|
286 |
+
- [Real-ESRGAN_GUI](https://github.com/net2cn/Real-ESRGAN_GUI) by [net2cn](https://github.com/net2cn)
|
287 |
+
- [Real-ESRGAN-EGUI](https://github.com/WGzeyu/Real-ESRGAN-EGUI) by [WGzeyu](https://github.com/WGzeyu)
|
288 |
+
- [anime_upscaler](https://github.com/shangar21/anime_upscaler) by [shangar21](https://github.com/shangar21)
|
289 |
+
- [Upscayl](https://github.com/upscayl/upscayl) by [Nayam Amarshe](https://github.com/NayamAmarshe) and [TGS963](https://github.com/TGS963)
|
290 |
+
|
291 |
+
## 🤗 Acknowledgement
|
292 |
+
|
293 |
+
Thanks for all the contributors.
|
294 |
+
|
295 |
+
- [AK391](https://github.com/AK391): Integrate RealESRGAN to [Huggingface Spaces](https://huggingface.co/spaces) with [Gradio](https://github.com/gradio-app/gradio). See [Gradio Web Demo](https://huggingface.co/spaces/akhaliq/Real-ESRGAN).
|
296 |
+
- [Asiimoviet](https://github.com/Asiimoviet): Translate the README.md to Chinese (中文).
|
297 |
+
- [2ji3150](https://github.com/2ji3150): Thanks for the [detailed and valuable feedbacks/suggestions](https://github.com/xinntao/Real-ESRGAN/issues/131).
|
298 |
+
- [Jared-02](https://github.com/Jared-02): Translate the Training.md to Chinese (中文).
|
realesrgan.egg-info/SOURCES.txt
ADDED
@@ -0,0 +1,21 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
README.md
|
2 |
+
setup.py
|
3 |
+
realesrgan/__init__.py
|
4 |
+
realesrgan/train.py
|
5 |
+
realesrgan/utils.py
|
6 |
+
realesrgan/version.py
|
7 |
+
realesrgan.egg-info/PKG-INFO
|
8 |
+
realesrgan.egg-info/SOURCES.txt
|
9 |
+
realesrgan.egg-info/dependency_links.txt
|
10 |
+
realesrgan.egg-info/not-zip-safe
|
11 |
+
realesrgan.egg-info/requires.txt
|
12 |
+
realesrgan.egg-info/top_level.txt
|
13 |
+
realesrgan/archs/__init__.py
|
14 |
+
realesrgan/archs/discriminator_arch.py
|
15 |
+
realesrgan/archs/srvgg_arch.py
|
16 |
+
realesrgan/data/__init__.py
|
17 |
+
realesrgan/data/realesrgan_dataset.py
|
18 |
+
realesrgan/data/realesrgan_paired_dataset.py
|
19 |
+
realesrgan/models/__init__.py
|
20 |
+
realesrgan/models/realesrgan_model.py
|
21 |
+
realesrgan/models/realesrnet_model.py
|
realesrgan.egg-info/dependency_links.txt
ADDED
@@ -0,0 +1 @@
|
|
|
|
|
1 |
+
|
realesrgan.egg-info/not-zip-safe
ADDED
@@ -0,0 +1 @@
|
|
|
|
|
1 |
+
|
realesrgan.egg-info/requires.txt
ADDED
@@ -0,0 +1,9 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
basicsr>=1.4.2
|
2 |
+
facexlib>=0.2.5
|
3 |
+
gfpgan>=1.3.5
|
4 |
+
numpy
|
5 |
+
opencv-python
|
6 |
+
Pillow
|
7 |
+
torch>=1.7
|
8 |
+
torchvision
|
9 |
+
tqdm
|
realesrgan.egg-info/top_level.txt
ADDED
@@ -0,0 +1 @@
|
|
|
|
|
1 |
+
realesrgan
|
realesrgan/__init__.py
ADDED
@@ -0,0 +1,6 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
# flake8: noqa
|
2 |
+
from .archs import *
|
3 |
+
from .data import *
|
4 |
+
from .models import *
|
5 |
+
from .utils import *
|
6 |
+
from .version import *
|
realesrgan/__pycache__/__init__.cpython-38.pyc
ADDED
Binary file (234 Bytes). View file
|
|
realesrgan/__pycache__/utils.cpython-38.pyc
ADDED
Binary file (9.09 kB). View file
|
|
realesrgan/__pycache__/version.cpython-38.pyc
ADDED
Binary file (226 Bytes). View file
|
|
realesrgan/archs/__init__.py
ADDED
@@ -0,0 +1,10 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
import importlib
|
2 |
+
from basicsr.utils import scandir
|
3 |
+
from os import path as osp
|
4 |
+
|
5 |
+
# automatically scan and import arch modules for registry
|
6 |
+
# scan all the files that end with '_arch.py' under the archs folder
|
7 |
+
arch_folder = osp.dirname(osp.abspath(__file__))
|
8 |
+
arch_filenames = [osp.splitext(osp.basename(v))[0] for v in scandir(arch_folder) if v.endswith('_arch.py')]
|
9 |
+
# import all the arch modules
|
10 |
+
_arch_modules = [importlib.import_module(f'realesrgan.archs.{file_name}') for file_name in arch_filenames]
|
realesrgan/archs/__pycache__/__init__.cpython-38.pyc
ADDED
Binary file (710 Bytes). View file
|
|
realesrgan/archs/__pycache__/discriminator_arch.cpython-38.pyc
ADDED
Binary file (2.41 kB). View file
|
|
realesrgan/archs/__pycache__/srvgg_arch.cpython-38.pyc
ADDED
Binary file (2.38 kB). View file
|
|
realesrgan/archs/discriminator_arch.py
ADDED
@@ -0,0 +1,67 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
from basicsr.utils.registry import ARCH_REGISTRY
|
2 |
+
from torch import nn as nn
|
3 |
+
from torch.nn import functional as F
|
4 |
+
from torch.nn.utils import spectral_norm
|
5 |
+
|
6 |
+
|
7 |
+
@ARCH_REGISTRY.register()
|
8 |
+
class UNetDiscriminatorSN(nn.Module):
|
9 |
+
"""Defines a U-Net discriminator with spectral normalization (SN)
|
10 |
+
|
11 |
+
It is used in Real-ESRGAN: Training Real-World Blind Super-Resolution with Pure Synthetic Data.
|
12 |
+
|
13 |
+
Arg:
|
14 |
+
num_in_ch (int): Channel number of inputs. Default: 3.
|
15 |
+
num_feat (int): Channel number of base intermediate features. Default: 64.
|
16 |
+
skip_connection (bool): Whether to use skip connections between U-Net. Default: True.
|
17 |
+
"""
|
18 |
+
|
19 |
+
def __init__(self, num_in_ch, num_feat=64, skip_connection=True):
|
20 |
+
super(UNetDiscriminatorSN, self).__init__()
|
21 |
+
self.skip_connection = skip_connection
|
22 |
+
norm = spectral_norm
|
23 |
+
# the first convolution
|
24 |
+
self.conv0 = nn.Conv2d(num_in_ch, num_feat, kernel_size=3, stride=1, padding=1)
|
25 |
+
# downsample
|
26 |
+
self.conv1 = norm(nn.Conv2d(num_feat, num_feat * 2, 4, 2, 1, bias=False))
|
27 |
+
self.conv2 = norm(nn.Conv2d(num_feat * 2, num_feat * 4, 4, 2, 1, bias=False))
|
28 |
+
self.conv3 = norm(nn.Conv2d(num_feat * 4, num_feat * 8, 4, 2, 1, bias=False))
|
29 |
+
# upsample
|
30 |
+
self.conv4 = norm(nn.Conv2d(num_feat * 8, num_feat * 4, 3, 1, 1, bias=False))
|
31 |
+
self.conv5 = norm(nn.Conv2d(num_feat * 4, num_feat * 2, 3, 1, 1, bias=False))
|
32 |
+
self.conv6 = norm(nn.Conv2d(num_feat * 2, num_feat, 3, 1, 1, bias=False))
|
33 |
+
# extra convolutions
|
34 |
+
self.conv7 = norm(nn.Conv2d(num_feat, num_feat, 3, 1, 1, bias=False))
|
35 |
+
self.conv8 = norm(nn.Conv2d(num_feat, num_feat, 3, 1, 1, bias=False))
|
36 |
+
self.conv9 = nn.Conv2d(num_feat, 1, 3, 1, 1)
|
37 |
+
|
38 |
+
def forward(self, x):
|
39 |
+
# downsample
|
40 |
+
x0 = F.leaky_relu(self.conv0(x), negative_slope=0.2, inplace=True)
|
41 |
+
x1 = F.leaky_relu(self.conv1(x0), negative_slope=0.2, inplace=True)
|
42 |
+
x2 = F.leaky_relu(self.conv2(x1), negative_slope=0.2, inplace=True)
|
43 |
+
x3 = F.leaky_relu(self.conv3(x2), negative_slope=0.2, inplace=True)
|
44 |
+
|
45 |
+
# upsample
|
46 |
+
x3 = F.interpolate(x3, scale_factor=2, mode='bilinear', align_corners=False)
|
47 |
+
x4 = F.leaky_relu(self.conv4(x3), negative_slope=0.2, inplace=True)
|
48 |
+
|
49 |
+
if self.skip_connection:
|
50 |
+
x4 = x4 + x2
|
51 |
+
x4 = F.interpolate(x4, scale_factor=2, mode='bilinear', align_corners=False)
|
52 |
+
x5 = F.leaky_relu(self.conv5(x4), negative_slope=0.2, inplace=True)
|
53 |
+
|
54 |
+
if self.skip_connection:
|
55 |
+
x5 = x5 + x1
|
56 |
+
x5 = F.interpolate(x5, scale_factor=2, mode='bilinear', align_corners=False)
|
57 |
+
x6 = F.leaky_relu(self.conv6(x5), negative_slope=0.2, inplace=True)
|
58 |
+
|
59 |
+
if self.skip_connection:
|
60 |
+
x6 = x6 + x0
|
61 |
+
|
62 |
+
# extra convolutions
|
63 |
+
out = F.leaky_relu(self.conv7(x6), negative_slope=0.2, inplace=True)
|
64 |
+
out = F.leaky_relu(self.conv8(out), negative_slope=0.2, inplace=True)
|
65 |
+
out = self.conv9(out)
|
66 |
+
|
67 |
+
return out
|
realesrgan/archs/srvgg_arch.py
ADDED
@@ -0,0 +1,69 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
from basicsr.utils.registry import ARCH_REGISTRY
|
2 |
+
from torch import nn as nn
|
3 |
+
from torch.nn import functional as F
|
4 |
+
|
5 |
+
|
6 |
+
@ARCH_REGISTRY.register()
|
7 |
+
class SRVGGNetCompact(nn.Module):
|
8 |
+
"""A compact VGG-style network structure for super-resolution.
|
9 |
+
|
10 |
+
It is a compact network structure, which performs upsampling in the last layer and no convolution is
|
11 |
+
conducted on the HR feature space.
|
12 |
+
|
13 |
+
Args:
|
14 |
+
num_in_ch (int): Channel number of inputs. Default: 3.
|
15 |
+
num_out_ch (int): Channel number of outputs. Default: 3.
|
16 |
+
num_feat (int): Channel number of intermediate features. Default: 64.
|
17 |
+
num_conv (int): Number of convolution layers in the body network. Default: 16.
|
18 |
+
upscale (int): Upsampling factor. Default: 4.
|
19 |
+
act_type (str): Activation type, options: 'relu', 'prelu', 'leakyrelu'. Default: prelu.
|
20 |
+
"""
|
21 |
+
|
22 |
+
def __init__(self, num_in_ch=3, num_out_ch=3, num_feat=64, num_conv=16, upscale=4, act_type='prelu'):
|
23 |
+
super(SRVGGNetCompact, self).__init__()
|
24 |
+
self.num_in_ch = num_in_ch
|
25 |
+
self.num_out_ch = num_out_ch
|
26 |
+
self.num_feat = num_feat
|
27 |
+
self.num_conv = num_conv
|
28 |
+
self.upscale = upscale
|
29 |
+
self.act_type = act_type
|
30 |
+
|
31 |
+
self.body = nn.ModuleList()
|
32 |
+
# the first conv
|
33 |
+
self.body.append(nn.Conv2d(num_in_ch, num_feat, 3, 1, 1))
|
34 |
+
# the first activation
|
35 |
+
if act_type == 'relu':
|
36 |
+
activation = nn.ReLU(inplace=True)
|
37 |
+
elif act_type == 'prelu':
|
38 |
+
activation = nn.PReLU(num_parameters=num_feat)
|
39 |
+
elif act_type == 'leakyrelu':
|
40 |
+
activation = nn.LeakyReLU(negative_slope=0.1, inplace=True)
|
41 |
+
self.body.append(activation)
|
42 |
+
|
43 |
+
# the body structure
|
44 |
+
for _ in range(num_conv):
|
45 |
+
self.body.append(nn.Conv2d(num_feat, num_feat, 3, 1, 1))
|
46 |
+
# activation
|
47 |
+
if act_type == 'relu':
|
48 |
+
activation = nn.ReLU(inplace=True)
|
49 |
+
elif act_type == 'prelu':
|
50 |
+
activation = nn.PReLU(num_parameters=num_feat)
|
51 |
+
elif act_type == 'leakyrelu':
|
52 |
+
activation = nn.LeakyReLU(negative_slope=0.1, inplace=True)
|
53 |
+
self.body.append(activation)
|
54 |
+
|
55 |
+
# the last conv
|
56 |
+
self.body.append(nn.Conv2d(num_feat, num_out_ch * upscale * upscale, 3, 1, 1))
|
57 |
+
# upsample
|
58 |
+
self.upsampler = nn.PixelShuffle(upscale)
|
59 |
+
|
60 |
+
def forward(self, x):
|
61 |
+
out = x
|
62 |
+
for i in range(0, len(self.body)):
|
63 |
+
out = self.body[i](out)
|
64 |
+
|
65 |
+
out = self.upsampler(out)
|
66 |
+
# add the nearest upsampled image, so that the network learns the residual
|
67 |
+
base = F.interpolate(x, scale_factor=self.upscale, mode='nearest')
|
68 |
+
out += base
|
69 |
+
return out
|
realesrgan/data/__init__.py
ADDED
@@ -0,0 +1,10 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
import importlib
|
2 |
+
from basicsr.utils import scandir
|
3 |
+
from os import path as osp
|
4 |
+
|
5 |
+
# automatically scan and import dataset modules for registry
|
6 |
+
# scan all the files that end with '_dataset.py' under the data folder
|
7 |
+
data_folder = osp.dirname(osp.abspath(__file__))
|
8 |
+
dataset_filenames = [osp.splitext(osp.basename(v))[0] for v in scandir(data_folder) if v.endswith('_dataset.py')]
|
9 |
+
# import all the dataset modules
|
10 |
+
_dataset_modules = [importlib.import_module(f'realesrgan.data.{file_name}') for file_name in dataset_filenames]
|
realesrgan/data/__pycache__/__init__.cpython-38.pyc
ADDED
Binary file (717 Bytes). View file
|
|
realesrgan/data/__pycache__/realesrgan_dataset.cpython-38.pyc
ADDED
Binary file (5.63 kB). View file
|
|
realesrgan/data/__pycache__/realesrgan_paired_dataset.cpython-38.pyc
ADDED
Binary file (4.06 kB). View file
|
|
realesrgan/data/realesrgan_dataset.py
ADDED
@@ -0,0 +1,192 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
import cv2
|
2 |
+
import math
|
3 |
+
import numpy as np
|
4 |
+
import os
|
5 |
+
import os.path as osp
|
6 |
+
import random
|
7 |
+
import time
|
8 |
+
import torch
|
9 |
+
from basicsr.data.degradations import circular_lowpass_kernel, random_mixed_kernels
|
10 |
+
from basicsr.data.transforms import augment
|
11 |
+
from basicsr.utils import FileClient, get_root_logger, imfrombytes, img2tensor
|
12 |
+
from basicsr.utils.registry import DATASET_REGISTRY
|
13 |
+
from torch.utils import data as data
|
14 |
+
|
15 |
+
|
16 |
+
@DATASET_REGISTRY.register()
|
17 |
+
class RealESRGANDataset(data.Dataset):
|
18 |
+
"""Dataset used for Real-ESRGAN model:
|
19 |
+
Real-ESRGAN: Training Real-World Blind Super-Resolution with Pure Synthetic Data.
|
20 |
+
|
21 |
+
It loads gt (Ground-Truth) images, and augments them.
|
22 |
+
It also generates blur kernels and sinc kernels for generating low-quality images.
|
23 |
+
Note that the low-quality images are processed in tensors on GPUS for faster processing.
|
24 |
+
|
25 |
+
Args:
|
26 |
+
opt (dict): Config for train datasets. It contains the following keys:
|
27 |
+
dataroot_gt (str): Data root path for gt.
|
28 |
+
meta_info (str): Path for meta information file.
|
29 |
+
io_backend (dict): IO backend type and other kwarg.
|
30 |
+
use_hflip (bool): Use horizontal flips.
|
31 |
+
use_rot (bool): Use rotation (use vertical flip and transposing h and w for implementation).
|
32 |
+
Please see more options in the codes.
|
33 |
+
"""
|
34 |
+
|
35 |
+
def __init__(self, opt):
|
36 |
+
super(RealESRGANDataset, self).__init__()
|
37 |
+
self.opt = opt
|
38 |
+
self.file_client = None
|
39 |
+
self.io_backend_opt = opt['io_backend']
|
40 |
+
self.gt_folder = opt['dataroot_gt']
|
41 |
+
|
42 |
+
# file client (lmdb io backend)
|
43 |
+
if self.io_backend_opt['type'] == 'lmdb':
|
44 |
+
self.io_backend_opt['db_paths'] = [self.gt_folder]
|
45 |
+
self.io_backend_opt['client_keys'] = ['gt']
|
46 |
+
if not self.gt_folder.endswith('.lmdb'):
|
47 |
+
raise ValueError(f"'dataroot_gt' should end with '.lmdb', but received {self.gt_folder}")
|
48 |
+
with open(osp.join(self.gt_folder, 'meta_info.txt')) as fin:
|
49 |
+
self.paths = [line.split('.')[0] for line in fin]
|
50 |
+
else:
|
51 |
+
# disk backend with meta_info
|
52 |
+
# Each line in the meta_info describes the relative path to an image
|
53 |
+
with open(self.opt['meta_info']) as fin:
|
54 |
+
paths = [line.strip().split(' ')[0] for line in fin]
|
55 |
+
self.paths = [os.path.join(self.gt_folder, v) for v in paths]
|
56 |
+
|
57 |
+
# blur settings for the first degradation
|
58 |
+
self.blur_kernel_size = opt['blur_kernel_size']
|
59 |
+
self.kernel_list = opt['kernel_list']
|
60 |
+
self.kernel_prob = opt['kernel_prob'] # a list for each kernel probability
|
61 |
+
self.blur_sigma = opt['blur_sigma']
|
62 |
+
self.betag_range = opt['betag_range'] # betag used in generalized Gaussian blur kernels
|
63 |
+
self.betap_range = opt['betap_range'] # betap used in plateau blur kernels
|
64 |
+
self.sinc_prob = opt['sinc_prob'] # the probability for sinc filters
|
65 |
+
|
66 |
+
# blur settings for the second degradation
|
67 |
+
self.blur_kernel_size2 = opt['blur_kernel_size2']
|
68 |
+
self.kernel_list2 = opt['kernel_list2']
|
69 |
+
self.kernel_prob2 = opt['kernel_prob2']
|
70 |
+
self.blur_sigma2 = opt['blur_sigma2']
|
71 |
+
self.betag_range2 = opt['betag_range2']
|
72 |
+
self.betap_range2 = opt['betap_range2']
|
73 |
+
self.sinc_prob2 = opt['sinc_prob2']
|
74 |
+
|
75 |
+
# a final sinc filter
|
76 |
+
self.final_sinc_prob = opt['final_sinc_prob']
|
77 |
+
|
78 |
+
self.kernel_range = [2 * v + 1 for v in range(3, 11)] # kernel size ranges from 7 to 21
|
79 |
+
# TODO: kernel range is now hard-coded, should be in the configure file
|
80 |
+
self.pulse_tensor = torch.zeros(21, 21).float() # convolving with pulse tensor brings no blurry effect
|
81 |
+
self.pulse_tensor[10, 10] = 1
|
82 |
+
|
83 |
+
def __getitem__(self, index):
|
84 |
+
if self.file_client is None:
|
85 |
+
self.file_client = FileClient(self.io_backend_opt.pop('type'), **self.io_backend_opt)
|
86 |
+
|
87 |
+
# -------------------------------- Load gt images -------------------------------- #
|
88 |
+
# Shape: (h, w, c); channel order: BGR; image range: [0, 1], float32.
|
89 |
+
gt_path = self.paths[index]
|
90 |
+
# avoid errors caused by high latency in reading files
|
91 |
+
retry = 3
|
92 |
+
while retry > 0:
|
93 |
+
try:
|
94 |
+
img_bytes = self.file_client.get(gt_path, 'gt')
|
95 |
+
except (IOError, OSError) as e:
|
96 |
+
logger = get_root_logger()
|
97 |
+
logger.warn(f'File client error: {e}, remaining retry times: {retry - 1}')
|
98 |
+
# change another file to read
|
99 |
+
index = random.randint(0, self.__len__())
|
100 |
+
gt_path = self.paths[index]
|
101 |
+
time.sleep(1) # sleep 1s for occasional server congestion
|
102 |
+
else:
|
103 |
+
break
|
104 |
+
finally:
|
105 |
+
retry -= 1
|
106 |
+
img_gt = imfrombytes(img_bytes, float32=True)
|
107 |
+
|
108 |
+
# -------------------- Do augmentation for training: flip, rotation -------------------- #
|
109 |
+
img_gt = augment(img_gt, self.opt['use_hflip'], self.opt['use_rot'])
|
110 |
+
|
111 |
+
# crop or pad to 400
|
112 |
+
# TODO: 400 is hard-coded. You may change it accordingly
|
113 |
+
h, w = img_gt.shape[0:2]
|
114 |
+
crop_pad_size = 400
|
115 |
+
# pad
|
116 |
+
if h < crop_pad_size or w < crop_pad_size:
|
117 |
+
pad_h = max(0, crop_pad_size - h)
|
118 |
+
pad_w = max(0, crop_pad_size - w)
|
119 |
+
img_gt = cv2.copyMakeBorder(img_gt, 0, pad_h, 0, pad_w, cv2.BORDER_REFLECT_101)
|
120 |
+
# crop
|
121 |
+
if img_gt.shape[0] > crop_pad_size or img_gt.shape[1] > crop_pad_size:
|
122 |
+
h, w = img_gt.shape[0:2]
|
123 |
+
# randomly choose top and left coordinates
|
124 |
+
top = random.randint(0, h - crop_pad_size)
|
125 |
+
left = random.randint(0, w - crop_pad_size)
|
126 |
+
img_gt = img_gt[top:top + crop_pad_size, left:left + crop_pad_size, ...]
|
127 |
+
|
128 |
+
# ------------------------ Generate kernels (used in the first degradation) ------------------------ #
|
129 |
+
kernel_size = random.choice(self.kernel_range)
|
130 |
+
if np.random.uniform() < self.opt['sinc_prob']:
|
131 |
+
# this sinc filter setting is for kernels ranging from [7, 21]
|
132 |
+
if kernel_size < 13:
|
133 |
+
omega_c = np.random.uniform(np.pi / 3, np.pi)
|
134 |
+
else:
|
135 |
+
omega_c = np.random.uniform(np.pi / 5, np.pi)
|
136 |
+
kernel = circular_lowpass_kernel(omega_c, kernel_size, pad_to=False)
|
137 |
+
else:
|
138 |
+
kernel = random_mixed_kernels(
|
139 |
+
self.kernel_list,
|
140 |
+
self.kernel_prob,
|
141 |
+
kernel_size,
|
142 |
+
self.blur_sigma,
|
143 |
+
self.blur_sigma, [-math.pi, math.pi],
|
144 |
+
self.betag_range,
|
145 |
+
self.betap_range,
|
146 |
+
noise_range=None)
|
147 |
+
# pad kernel
|
148 |
+
pad_size = (21 - kernel_size) // 2
|
149 |
+
kernel = np.pad(kernel, ((pad_size, pad_size), (pad_size, pad_size)))
|
150 |
+
|
151 |
+
# ------------------------ Generate kernels (used in the second degradation) ------------------------ #
|
152 |
+
kernel_size = random.choice(self.kernel_range)
|
153 |
+
if np.random.uniform() < self.opt['sinc_prob2']:
|
154 |
+
if kernel_size < 13:
|
155 |
+
omega_c = np.random.uniform(np.pi / 3, np.pi)
|
156 |
+
else:
|
157 |
+
omega_c = np.random.uniform(np.pi / 5, np.pi)
|
158 |
+
kernel2 = circular_lowpass_kernel(omega_c, kernel_size, pad_to=False)
|
159 |
+
else:
|
160 |
+
kernel2 = random_mixed_kernels(
|
161 |
+
self.kernel_list2,
|
162 |
+
self.kernel_prob2,
|
163 |
+
kernel_size,
|
164 |
+
self.blur_sigma2,
|
165 |
+
self.blur_sigma2, [-math.pi, math.pi],
|
166 |
+
self.betag_range2,
|
167 |
+
self.betap_range2,
|
168 |
+
noise_range=None)
|
169 |
+
|
170 |
+
# pad kernel
|
171 |
+
pad_size = (21 - kernel_size) // 2
|
172 |
+
kernel2 = np.pad(kernel2, ((pad_size, pad_size), (pad_size, pad_size)))
|
173 |
+
|
174 |
+
# ------------------------------------- the final sinc kernel ------------------------------------- #
|
175 |
+
if np.random.uniform() < self.opt['final_sinc_prob']:
|
176 |
+
kernel_size = random.choice(self.kernel_range)
|
177 |
+
omega_c = np.random.uniform(np.pi / 3, np.pi)
|
178 |
+
sinc_kernel = circular_lowpass_kernel(omega_c, kernel_size, pad_to=21)
|
179 |
+
sinc_kernel = torch.FloatTensor(sinc_kernel)
|
180 |
+
else:
|
181 |
+
sinc_kernel = self.pulse_tensor
|
182 |
+
|
183 |
+
# BGR to RGB, HWC to CHW, numpy to tensor
|
184 |
+
img_gt = img2tensor([img_gt], bgr2rgb=True, float32=True)[0]
|
185 |
+
kernel = torch.FloatTensor(kernel)
|
186 |
+
kernel2 = torch.FloatTensor(kernel2)
|
187 |
+
|
188 |
+
return_d = {'gt': img_gt, 'kernel1': kernel, 'kernel2': kernel2, 'sinc_kernel': sinc_kernel, 'gt_path': gt_path}
|
189 |
+
return return_d
|
190 |
+
|
191 |
+
def __len__(self):
|
192 |
+
return len(self.paths)
|
realesrgan/data/realesrgan_paired_dataset.py
ADDED
@@ -0,0 +1,108 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
import os
|
2 |
+
from basicsr.data.data_util import paired_paths_from_folder, paired_paths_from_lmdb
|
3 |
+
from basicsr.data.transforms import augment, paired_random_crop
|
4 |
+
from basicsr.utils import FileClient, imfrombytes, img2tensor
|
5 |
+
from basicsr.utils.registry import DATASET_REGISTRY
|
6 |
+
from torch.utils import data as data
|
7 |
+
from torchvision.transforms.functional import normalize
|
8 |
+
|
9 |
+
|
10 |
+
@DATASET_REGISTRY.register()
|
11 |
+
class RealESRGANPairedDataset(data.Dataset):
|
12 |
+
"""Paired image dataset for image restoration.
|
13 |
+
|
14 |
+
Read LQ (Low Quality, e.g. LR (Low Resolution), blurry, noisy, etc) and GT image pairs.
|
15 |
+
|
16 |
+
There are three modes:
|
17 |
+
1. 'lmdb': Use lmdb files.
|
18 |
+
If opt['io_backend'] == lmdb.
|
19 |
+
2. 'meta_info': Use meta information file to generate paths.
|
20 |
+
If opt['io_backend'] != lmdb and opt['meta_info'] is not None.
|
21 |
+
3. 'folder': Scan folders to generate paths.
|
22 |
+
The rest.
|
23 |
+
|
24 |
+
Args:
|
25 |
+
opt (dict): Config for train datasets. It contains the following keys:
|
26 |
+
dataroot_gt (str): Data root path for gt.
|
27 |
+
dataroot_lq (str): Data root path for lq.
|
28 |
+
meta_info (str): Path for meta information file.
|
29 |
+
io_backend (dict): IO backend type and other kwarg.
|
30 |
+
filename_tmpl (str): Template for each filename. Note that the template excludes the file extension.
|
31 |
+
Default: '{}'.
|
32 |
+
gt_size (int): Cropped patched size for gt patches.
|
33 |
+
use_hflip (bool): Use horizontal flips.
|
34 |
+
use_rot (bool): Use rotation (use vertical flip and transposing h
|
35 |
+
and w for implementation).
|
36 |
+
|
37 |
+
scale (bool): Scale, which will be added automatically.
|
38 |
+
phase (str): 'train' or 'val'.
|
39 |
+
"""
|
40 |
+
|
41 |
+
def __init__(self, opt):
|
42 |
+
super(RealESRGANPairedDataset, self).__init__()
|
43 |
+
self.opt = opt
|
44 |
+
self.file_client = None
|
45 |
+
self.io_backend_opt = opt['io_backend']
|
46 |
+
# mean and std for normalizing the input images
|
47 |
+
self.mean = opt['mean'] if 'mean' in opt else None
|
48 |
+
self.std = opt['std'] if 'std' in opt else None
|
49 |
+
|
50 |
+
self.gt_folder, self.lq_folder = opt['dataroot_gt'], opt['dataroot_lq']
|
51 |
+
self.filename_tmpl = opt['filename_tmpl'] if 'filename_tmpl' in opt else '{}'
|
52 |
+
|
53 |
+
# file client (lmdb io backend)
|
54 |
+
if self.io_backend_opt['type'] == 'lmdb':
|
55 |
+
self.io_backend_opt['db_paths'] = [self.lq_folder, self.gt_folder]
|
56 |
+
self.io_backend_opt['client_keys'] = ['lq', 'gt']
|
57 |
+
self.paths = paired_paths_from_lmdb([self.lq_folder, self.gt_folder], ['lq', 'gt'])
|
58 |
+
elif 'meta_info' in self.opt and self.opt['meta_info'] is not None:
|
59 |
+
# disk backend with meta_info
|
60 |
+
# Each line in the meta_info describes the relative path to an image
|
61 |
+
with open(self.opt['meta_info']) as fin:
|
62 |
+
paths = [line.strip() for line in fin]
|
63 |
+
self.paths = []
|
64 |
+
for path in paths:
|
65 |
+
gt_path, lq_path = path.split(', ')
|
66 |
+
gt_path = os.path.join(self.gt_folder, gt_path)
|
67 |
+
lq_path = os.path.join(self.lq_folder, lq_path)
|
68 |
+
self.paths.append(dict([('gt_path', gt_path), ('lq_path', lq_path)]))
|
69 |
+
else:
|
70 |
+
# disk backend
|
71 |
+
# it will scan the whole folder to get meta info
|
72 |
+
# it will be time-consuming for folders with too many files. It is recommended using an extra meta txt file
|
73 |
+
self.paths = paired_paths_from_folder([self.lq_folder, self.gt_folder], ['lq', 'gt'], self.filename_tmpl)
|
74 |
+
|
75 |
+
def __getitem__(self, index):
|
76 |
+
if self.file_client is None:
|
77 |
+
self.file_client = FileClient(self.io_backend_opt.pop('type'), **self.io_backend_opt)
|
78 |
+
|
79 |
+
scale = self.opt['scale']
|
80 |
+
|
81 |
+
# Load gt and lq images. Dimension order: HWC; channel order: BGR;
|
82 |
+
# image range: [0, 1], float32.
|
83 |
+
gt_path = self.paths[index]['gt_path']
|
84 |
+
img_bytes = self.file_client.get(gt_path, 'gt')
|
85 |
+
img_gt = imfrombytes(img_bytes, float32=True)
|
86 |
+
lq_path = self.paths[index]['lq_path']
|
87 |
+
img_bytes = self.file_client.get(lq_path, 'lq')
|
88 |
+
img_lq = imfrombytes(img_bytes, float32=True)
|
89 |
+
|
90 |
+
# augmentation for training
|
91 |
+
if self.opt['phase'] == 'train':
|
92 |
+
gt_size = self.opt['gt_size']
|
93 |
+
# random crop
|
94 |
+
img_gt, img_lq = paired_random_crop(img_gt, img_lq, gt_size, scale, gt_path)
|
95 |
+
# flip, rotation
|
96 |
+
img_gt, img_lq = augment([img_gt, img_lq], self.opt['use_hflip'], self.opt['use_rot'])
|
97 |
+
|
98 |
+
# BGR to RGB, HWC to CHW, numpy to tensor
|
99 |
+
img_gt, img_lq = img2tensor([img_gt, img_lq], bgr2rgb=True, float32=True)
|
100 |
+
# normalize
|
101 |
+
if self.mean is not None or self.std is not None:
|
102 |
+
normalize(img_lq, self.mean, self.std, inplace=True)
|
103 |
+
normalize(img_gt, self.mean, self.std, inplace=True)
|
104 |
+
|
105 |
+
return {'lq': img_lq, 'gt': img_gt, 'lq_path': lq_path, 'gt_path': gt_path}
|
106 |
+
|
107 |
+
def __len__(self):
|
108 |
+
return len(self.paths)
|
realesrgan/models/__init__.py
ADDED
@@ -0,0 +1,10 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
import importlib
|
2 |
+
from basicsr.utils import scandir
|
3 |
+
from os import path as osp
|
4 |
+
|
5 |
+
# automatically scan and import model modules for registry
|
6 |
+
# scan all the files that end with '_model.py' under the model folder
|
7 |
+
model_folder = osp.dirname(osp.abspath(__file__))
|
8 |
+
model_filenames = [osp.splitext(osp.basename(v))[0] for v in scandir(model_folder) if v.endswith('_model.py')]
|
9 |
+
# import all the model modules
|
10 |
+
_model_modules = [importlib.import_module(f'realesrgan.models.{file_name}') for file_name in model_filenames]
|
realesrgan/models/__pycache__/__init__.cpython-38.pyc
ADDED
Binary file (716 Bytes). View file
|
|
realesrgan/models/__pycache__/realesrgan_model.cpython-38.pyc
ADDED
Binary file (6.7 kB). View file
|
|
realesrgan/models/__pycache__/realesrnet_model.cpython-38.pyc
ADDED
Binary file (5.35 kB). View file
|
|
realesrgan/models/realesrgan_model.py
ADDED
@@ -0,0 +1,258 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
import numpy as np
|
2 |
+
import random
|
3 |
+
import torch
|
4 |
+
from basicsr.data.degradations import random_add_gaussian_noise_pt, random_add_poisson_noise_pt
|
5 |
+
from basicsr.data.transforms import paired_random_crop
|
6 |
+
from basicsr.models.srgan_model import SRGANModel
|
7 |
+
from basicsr.utils import DiffJPEG, USMSharp
|
8 |
+
from basicsr.utils.img_process_util import filter2D
|
9 |
+
from basicsr.utils.registry import MODEL_REGISTRY
|
10 |
+
from collections import OrderedDict
|
11 |
+
from torch.nn import functional as F
|
12 |
+
|
13 |
+
|
14 |
+
@MODEL_REGISTRY.register()
|
15 |
+
class RealESRGANModel(SRGANModel):
|
16 |
+
"""RealESRGAN Model for Real-ESRGAN: Training Real-World Blind Super-Resolution with Pure Synthetic Data.
|
17 |
+
|
18 |
+
It mainly performs:
|
19 |
+
1. randomly synthesize LQ images in GPU tensors
|
20 |
+
2. optimize the networks with GAN training.
|
21 |
+
"""
|
22 |
+
|
23 |
+
def __init__(self, opt):
|
24 |
+
super(RealESRGANModel, self).__init__(opt)
|
25 |
+
self.jpeger = DiffJPEG(differentiable=False).cuda() # simulate JPEG compression artifacts
|
26 |
+
self.usm_sharpener = USMSharp().cuda() # do usm sharpening
|
27 |
+
self.queue_size = opt.get('queue_size', 180)
|
28 |
+
|
29 |
+
@torch.no_grad()
|
30 |
+
def _dequeue_and_enqueue(self):
|
31 |
+
"""It is the training pair pool for increasing the diversity in a batch.
|
32 |
+
|
33 |
+
Batch processing limits the diversity of synthetic degradations in a batch. For example, samples in a
|
34 |
+
batch could not have different resize scaling factors. Therefore, we employ this training pair pool
|
35 |
+
to increase the degradation diversity in a batch.
|
36 |
+
"""
|
37 |
+
# initialize
|
38 |
+
b, c, h, w = self.lq.size()
|
39 |
+
if not hasattr(self, 'queue_lr'):
|
40 |
+
assert self.queue_size % b == 0, f'queue size {self.queue_size} should be divisible by batch size {b}'
|
41 |
+
self.queue_lr = torch.zeros(self.queue_size, c, h, w).cuda()
|
42 |
+
_, c, h, w = self.gt.size()
|
43 |
+
self.queue_gt = torch.zeros(self.queue_size, c, h, w).cuda()
|
44 |
+
self.queue_ptr = 0
|
45 |
+
if self.queue_ptr == self.queue_size: # the pool is full
|
46 |
+
# do dequeue and enqueue
|
47 |
+
# shuffle
|
48 |
+
idx = torch.randperm(self.queue_size)
|
49 |
+
self.queue_lr = self.queue_lr[idx]
|
50 |
+
self.queue_gt = self.queue_gt[idx]
|
51 |
+
# get first b samples
|
52 |
+
lq_dequeue = self.queue_lr[0:b, :, :, :].clone()
|
53 |
+
gt_dequeue = self.queue_gt[0:b, :, :, :].clone()
|
54 |
+
# update the queue
|
55 |
+
self.queue_lr[0:b, :, :, :] = self.lq.clone()
|
56 |
+
self.queue_gt[0:b, :, :, :] = self.gt.clone()
|
57 |
+
|
58 |
+
self.lq = lq_dequeue
|
59 |
+
self.gt = gt_dequeue
|
60 |
+
else:
|
61 |
+
# only do enqueue
|
62 |
+
self.queue_lr[self.queue_ptr:self.queue_ptr + b, :, :, :] = self.lq.clone()
|
63 |
+
self.queue_gt[self.queue_ptr:self.queue_ptr + b, :, :, :] = self.gt.clone()
|
64 |
+
self.queue_ptr = self.queue_ptr + b
|
65 |
+
|
66 |
+
@torch.no_grad()
|
67 |
+
def feed_data(self, data):
|
68 |
+
"""Accept data from dataloader, and then add two-order degradations to obtain LQ images.
|
69 |
+
"""
|
70 |
+
if self.is_train and self.opt.get('high_order_degradation', True):
|
71 |
+
# training data synthesis
|
72 |
+
self.gt = data['gt'].to(self.device)
|
73 |
+
self.gt_usm = self.usm_sharpener(self.gt)
|
74 |
+
|
75 |
+
self.kernel1 = data['kernel1'].to(self.device)
|
76 |
+
self.kernel2 = data['kernel2'].to(self.device)
|
77 |
+
self.sinc_kernel = data['sinc_kernel'].to(self.device)
|
78 |
+
|
79 |
+
ori_h, ori_w = self.gt.size()[2:4]
|
80 |
+
|
81 |
+
# ----------------------- The first degradation process ----------------------- #
|
82 |
+
# blur
|
83 |
+
out = filter2D(self.gt_usm, self.kernel1)
|
84 |
+
# random resize
|
85 |
+
updown_type = random.choices(['up', 'down', 'keep'], self.opt['resize_prob'])[0]
|
86 |
+
if updown_type == 'up':
|
87 |
+
scale = np.random.uniform(1, self.opt['resize_range'][1])
|
88 |
+
elif updown_type == 'down':
|
89 |
+
scale = np.random.uniform(self.opt['resize_range'][0], 1)
|
90 |
+
else:
|
91 |
+
scale = 1
|
92 |
+
mode = random.choice(['area', 'bilinear', 'bicubic'])
|
93 |
+
out = F.interpolate(out, scale_factor=scale, mode=mode)
|
94 |
+
# add noise
|
95 |
+
gray_noise_prob = self.opt['gray_noise_prob']
|
96 |
+
if np.random.uniform() < self.opt['gaussian_noise_prob']:
|
97 |
+
out = random_add_gaussian_noise_pt(
|
98 |
+
out, sigma_range=self.opt['noise_range'], clip=True, rounds=False, gray_prob=gray_noise_prob)
|
99 |
+
else:
|
100 |
+
out = random_add_poisson_noise_pt(
|
101 |
+
out,
|
102 |
+
scale_range=self.opt['poisson_scale_range'],
|
103 |
+
gray_prob=gray_noise_prob,
|
104 |
+
clip=True,
|
105 |
+
rounds=False)
|
106 |
+
# JPEG compression
|
107 |
+
jpeg_p = out.new_zeros(out.size(0)).uniform_(*self.opt['jpeg_range'])
|
108 |
+
out = torch.clamp(out, 0, 1) # clamp to [0, 1], otherwise JPEGer will result in unpleasant artifacts
|
109 |
+
out = self.jpeger(out, quality=jpeg_p)
|
110 |
+
|
111 |
+
# ----------------------- The second degradation process ----------------------- #
|
112 |
+
# blur
|
113 |
+
if np.random.uniform() < self.opt['second_blur_prob']:
|
114 |
+
out = filter2D(out, self.kernel2)
|
115 |
+
# random resize
|
116 |
+
updown_type = random.choices(['up', 'down', 'keep'], self.opt['resize_prob2'])[0]
|
117 |
+
if updown_type == 'up':
|
118 |
+
scale = np.random.uniform(1, self.opt['resize_range2'][1])
|
119 |
+
elif updown_type == 'down':
|
120 |
+
scale = np.random.uniform(self.opt['resize_range2'][0], 1)
|
121 |
+
else:
|
122 |
+
scale = 1
|
123 |
+
mode = random.choice(['area', 'bilinear', 'bicubic'])
|
124 |
+
out = F.interpolate(
|
125 |
+
out, size=(int(ori_h / self.opt['scale'] * scale), int(ori_w / self.opt['scale'] * scale)), mode=mode)
|
126 |
+
# add noise
|
127 |
+
gray_noise_prob = self.opt['gray_noise_prob2']
|
128 |
+
if np.random.uniform() < self.opt['gaussian_noise_prob2']:
|
129 |
+
out = random_add_gaussian_noise_pt(
|
130 |
+
out, sigma_range=self.opt['noise_range2'], clip=True, rounds=False, gray_prob=gray_noise_prob)
|
131 |
+
else:
|
132 |
+
out = random_add_poisson_noise_pt(
|
133 |
+
out,
|
134 |
+
scale_range=self.opt['poisson_scale_range2'],
|
135 |
+
gray_prob=gray_noise_prob,
|
136 |
+
clip=True,
|
137 |
+
rounds=False)
|
138 |
+
|
139 |
+
# JPEG compression + the final sinc filter
|
140 |
+
# We also need to resize images to desired sizes. We group [resize back + sinc filter] together
|
141 |
+
# as one operation.
|
142 |
+
# We consider two orders:
|
143 |
+
# 1. [resize back + sinc filter] + JPEG compression
|
144 |
+
# 2. JPEG compression + [resize back + sinc filter]
|
145 |
+
# Empirically, we find other combinations (sinc + JPEG + Resize) will introduce twisted lines.
|
146 |
+
if np.random.uniform() < 0.5:
|
147 |
+
# resize back + the final sinc filter
|
148 |
+
mode = random.choice(['area', 'bilinear', 'bicubic'])
|
149 |
+
out = F.interpolate(out, size=(ori_h // self.opt['scale'], ori_w // self.opt['scale']), mode=mode)
|
150 |
+
out = filter2D(out, self.sinc_kernel)
|
151 |
+
# JPEG compression
|
152 |
+
jpeg_p = out.new_zeros(out.size(0)).uniform_(*self.opt['jpeg_range2'])
|
153 |
+
out = torch.clamp(out, 0, 1)
|
154 |
+
out = self.jpeger(out, quality=jpeg_p)
|
155 |
+
else:
|
156 |
+
# JPEG compression
|
157 |
+
jpeg_p = out.new_zeros(out.size(0)).uniform_(*self.opt['jpeg_range2'])
|
158 |
+
out = torch.clamp(out, 0, 1)
|
159 |
+
out = self.jpeger(out, quality=jpeg_p)
|
160 |
+
# resize back + the final sinc filter
|
161 |
+
mode = random.choice(['area', 'bilinear', 'bicubic'])
|
162 |
+
out = F.interpolate(out, size=(ori_h // self.opt['scale'], ori_w // self.opt['scale']), mode=mode)
|
163 |
+
out = filter2D(out, self.sinc_kernel)
|
164 |
+
|
165 |
+
# clamp and round
|
166 |
+
self.lq = torch.clamp((out * 255.0).round(), 0, 255) / 255.
|
167 |
+
|
168 |
+
# random crop
|
169 |
+
gt_size = self.opt['gt_size']
|
170 |
+
(self.gt, self.gt_usm), self.lq = paired_random_crop([self.gt, self.gt_usm], self.lq, gt_size,
|
171 |
+
self.opt['scale'])
|
172 |
+
|
173 |
+
# training pair pool
|
174 |
+
self._dequeue_and_enqueue()
|
175 |
+
# sharpen self.gt again, as we have changed the self.gt with self._dequeue_and_enqueue
|
176 |
+
self.gt_usm = self.usm_sharpener(self.gt)
|
177 |
+
self.lq = self.lq.contiguous() # for the warning: grad and param do not obey the gradient layout contract
|
178 |
+
else:
|
179 |
+
# for paired training or validation
|
180 |
+
self.lq = data['lq'].to(self.device)
|
181 |
+
if 'gt' in data:
|
182 |
+
self.gt = data['gt'].to(self.device)
|
183 |
+
self.gt_usm = self.usm_sharpener(self.gt)
|
184 |
+
|
185 |
+
def nondist_validation(self, dataloader, current_iter, tb_logger, save_img):
|
186 |
+
# do not use the synthetic process during validation
|
187 |
+
self.is_train = False
|
188 |
+
super(RealESRGANModel, self).nondist_validation(dataloader, current_iter, tb_logger, save_img)
|
189 |
+
self.is_train = True
|
190 |
+
|
191 |
+
def optimize_parameters(self, current_iter):
|
192 |
+
# usm sharpening
|
193 |
+
l1_gt = self.gt_usm
|
194 |
+
percep_gt = self.gt_usm
|
195 |
+
gan_gt = self.gt_usm
|
196 |
+
if self.opt['l1_gt_usm'] is False:
|
197 |
+
l1_gt = self.gt
|
198 |
+
if self.opt['percep_gt_usm'] is False:
|
199 |
+
percep_gt = self.gt
|
200 |
+
if self.opt['gan_gt_usm'] is False:
|
201 |
+
gan_gt = self.gt
|
202 |
+
|
203 |
+
# optimize net_g
|
204 |
+
for p in self.net_d.parameters():
|
205 |
+
p.requires_grad = False
|
206 |
+
|
207 |
+
self.optimizer_g.zero_grad()
|
208 |
+
self.output = self.net_g(self.lq)
|
209 |
+
|
210 |
+
l_g_total = 0
|
211 |
+
loss_dict = OrderedDict()
|
212 |
+
if (current_iter % self.net_d_iters == 0 and current_iter > self.net_d_init_iters):
|
213 |
+
# pixel loss
|
214 |
+
if self.cri_pix:
|
215 |
+
l_g_pix = self.cri_pix(self.output, l1_gt)
|
216 |
+
l_g_total += l_g_pix
|
217 |
+
loss_dict['l_g_pix'] = l_g_pix
|
218 |
+
# perceptual loss
|
219 |
+
if self.cri_perceptual:
|
220 |
+
l_g_percep, l_g_style = self.cri_perceptual(self.output, percep_gt)
|
221 |
+
if l_g_percep is not None:
|
222 |
+
l_g_total += l_g_percep
|
223 |
+
loss_dict['l_g_percep'] = l_g_percep
|
224 |
+
if l_g_style is not None:
|
225 |
+
l_g_total += l_g_style
|
226 |
+
loss_dict['l_g_style'] = l_g_style
|
227 |
+
# gan loss
|
228 |
+
fake_g_pred = self.net_d(self.output)
|
229 |
+
l_g_gan = self.cri_gan(fake_g_pred, True, is_disc=False)
|
230 |
+
l_g_total += l_g_gan
|
231 |
+
loss_dict['l_g_gan'] = l_g_gan
|
232 |
+
|
233 |
+
l_g_total.backward()
|
234 |
+
self.optimizer_g.step()
|
235 |
+
|
236 |
+
# optimize net_d
|
237 |
+
for p in self.net_d.parameters():
|
238 |
+
p.requires_grad = True
|
239 |
+
|
240 |
+
self.optimizer_d.zero_grad()
|
241 |
+
# real
|
242 |
+
real_d_pred = self.net_d(gan_gt)
|
243 |
+
l_d_real = self.cri_gan(real_d_pred, True, is_disc=True)
|
244 |
+
loss_dict['l_d_real'] = l_d_real
|
245 |
+
loss_dict['out_d_real'] = torch.mean(real_d_pred.detach())
|
246 |
+
l_d_real.backward()
|
247 |
+
# fake
|
248 |
+
fake_d_pred = self.net_d(self.output.detach().clone()) # clone for pt1.9
|
249 |
+
l_d_fake = self.cri_gan(fake_d_pred, False, is_disc=True)
|
250 |
+
loss_dict['l_d_fake'] = l_d_fake
|
251 |
+
loss_dict['out_d_fake'] = torch.mean(fake_d_pred.detach())
|
252 |
+
l_d_fake.backward()
|
253 |
+
self.optimizer_d.step()
|
254 |
+
|
255 |
+
if self.ema_decay > 0:
|
256 |
+
self.model_ema(decay=self.ema_decay)
|
257 |
+
|
258 |
+
self.log_dict = self.reduce_loss_dict(loss_dict)
|
realesrgan/models/realesrnet_model.py
ADDED
@@ -0,0 +1,188 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
import numpy as np
|
2 |
+
import random
|
3 |
+
import torch
|
4 |
+
from basicsr.data.degradations import random_add_gaussian_noise_pt, random_add_poisson_noise_pt
|
5 |
+
from basicsr.data.transforms import paired_random_crop
|
6 |
+
from basicsr.models.sr_model import SRModel
|
7 |
+
from basicsr.utils import DiffJPEG, USMSharp
|
8 |
+
from basicsr.utils.img_process_util import filter2D
|
9 |
+
from basicsr.utils.registry import MODEL_REGISTRY
|
10 |
+
from torch.nn import functional as F
|
11 |
+
|
12 |
+
|
13 |
+
@MODEL_REGISTRY.register()
|
14 |
+
class RealESRNetModel(SRModel):
|
15 |
+
"""RealESRNet Model for Real-ESRGAN: Training Real-World Blind Super-Resolution with Pure Synthetic Data.
|
16 |
+
|
17 |
+
It is trained without GAN losses.
|
18 |
+
It mainly performs:
|
19 |
+
1. randomly synthesize LQ images in GPU tensors
|
20 |
+
2. optimize the networks with GAN training.
|
21 |
+
"""
|
22 |
+
|
23 |
+
def __init__(self, opt):
|
24 |
+
super(RealESRNetModel, self).__init__(opt)
|
25 |
+
self.jpeger = DiffJPEG(differentiable=False).cuda() # simulate JPEG compression artifacts
|
26 |
+
self.usm_sharpener = USMSharp().cuda() # do usm sharpening
|
27 |
+
self.queue_size = opt.get('queue_size', 180)
|
28 |
+
|
29 |
+
@torch.no_grad()
|
30 |
+
def _dequeue_and_enqueue(self):
|
31 |
+
"""It is the training pair pool for increasing the diversity in a batch.
|
32 |
+
|
33 |
+
Batch processing limits the diversity of synthetic degradations in a batch. For example, samples in a
|
34 |
+
batch could not have different resize scaling factors. Therefore, we employ this training pair pool
|
35 |
+
to increase the degradation diversity in a batch.
|
36 |
+
"""
|
37 |
+
# initialize
|
38 |
+
b, c, h, w = self.lq.size()
|
39 |
+
if not hasattr(self, 'queue_lr'):
|
40 |
+
assert self.queue_size % b == 0, f'queue size {self.queue_size} should be divisible by batch size {b}'
|
41 |
+
self.queue_lr = torch.zeros(self.queue_size, c, h, w).cuda()
|
42 |
+
_, c, h, w = self.gt.size()
|
43 |
+
self.queue_gt = torch.zeros(self.queue_size, c, h, w).cuda()
|
44 |
+
self.queue_ptr = 0
|
45 |
+
if self.queue_ptr == self.queue_size: # the pool is full
|
46 |
+
# do dequeue and enqueue
|
47 |
+
# shuffle
|
48 |
+
idx = torch.randperm(self.queue_size)
|
49 |
+
self.queue_lr = self.queue_lr[idx]
|
50 |
+
self.queue_gt = self.queue_gt[idx]
|
51 |
+
# get first b samples
|
52 |
+
lq_dequeue = self.queue_lr[0:b, :, :, :].clone()
|
53 |
+
gt_dequeue = self.queue_gt[0:b, :, :, :].clone()
|
54 |
+
# update the queue
|
55 |
+
self.queue_lr[0:b, :, :, :] = self.lq.clone()
|
56 |
+
self.queue_gt[0:b, :, :, :] = self.gt.clone()
|
57 |
+
|
58 |
+
self.lq = lq_dequeue
|
59 |
+
self.gt = gt_dequeue
|
60 |
+
else:
|
61 |
+
# only do enqueue
|
62 |
+
self.queue_lr[self.queue_ptr:self.queue_ptr + b, :, :, :] = self.lq.clone()
|
63 |
+
self.queue_gt[self.queue_ptr:self.queue_ptr + b, :, :, :] = self.gt.clone()
|
64 |
+
self.queue_ptr = self.queue_ptr + b
|
65 |
+
|
66 |
+
@torch.no_grad()
|
67 |
+
def feed_data(self, data):
|
68 |
+
"""Accept data from dataloader, and then add two-order degradations to obtain LQ images.
|
69 |
+
"""
|
70 |
+
if self.is_train and self.opt.get('high_order_degradation', True):
|
71 |
+
# training data synthesis
|
72 |
+
self.gt = data['gt'].to(self.device)
|
73 |
+
# USM sharpen the GT images
|
74 |
+
if self.opt['gt_usm'] is True:
|
75 |
+
self.gt = self.usm_sharpener(self.gt)
|
76 |
+
|
77 |
+
self.kernel1 = data['kernel1'].to(self.device)
|
78 |
+
self.kernel2 = data['kernel2'].to(self.device)
|
79 |
+
self.sinc_kernel = data['sinc_kernel'].to(self.device)
|
80 |
+
|
81 |
+
ori_h, ori_w = self.gt.size()[2:4]
|
82 |
+
|
83 |
+
# ----------------------- The first degradation process ----------------------- #
|
84 |
+
# blur
|
85 |
+
out = filter2D(self.gt, self.kernel1)
|
86 |
+
# random resize
|
87 |
+
updown_type = random.choices(['up', 'down', 'keep'], self.opt['resize_prob'])[0]
|
88 |
+
if updown_type == 'up':
|
89 |
+
scale = np.random.uniform(1, self.opt['resize_range'][1])
|
90 |
+
elif updown_type == 'down':
|
91 |
+
scale = np.random.uniform(self.opt['resize_range'][0], 1)
|
92 |
+
else:
|
93 |
+
scale = 1
|
94 |
+
mode = random.choice(['area', 'bilinear', 'bicubic'])
|
95 |
+
out = F.interpolate(out, scale_factor=scale, mode=mode)
|
96 |
+
# add noise
|
97 |
+
gray_noise_prob = self.opt['gray_noise_prob']
|
98 |
+
if np.random.uniform() < self.opt['gaussian_noise_prob']:
|
99 |
+
out = random_add_gaussian_noise_pt(
|
100 |
+
out, sigma_range=self.opt['noise_range'], clip=True, rounds=False, gray_prob=gray_noise_prob)
|
101 |
+
else:
|
102 |
+
out = random_add_poisson_noise_pt(
|
103 |
+
out,
|
104 |
+
scale_range=self.opt['poisson_scale_range'],
|
105 |
+
gray_prob=gray_noise_prob,
|
106 |
+
clip=True,
|
107 |
+
rounds=False)
|
108 |
+
# JPEG compression
|
109 |
+
jpeg_p = out.new_zeros(out.size(0)).uniform_(*self.opt['jpeg_range'])
|
110 |
+
out = torch.clamp(out, 0, 1) # clamp to [0, 1], otherwise JPEGer will result in unpleasant artifacts
|
111 |
+
out = self.jpeger(out, quality=jpeg_p)
|
112 |
+
|
113 |
+
# ----------------------- The second degradation process ----------------------- #
|
114 |
+
# blur
|
115 |
+
if np.random.uniform() < self.opt['second_blur_prob']:
|
116 |
+
out = filter2D(out, self.kernel2)
|
117 |
+
# random resize
|
118 |
+
updown_type = random.choices(['up', 'down', 'keep'], self.opt['resize_prob2'])[0]
|
119 |
+
if updown_type == 'up':
|
120 |
+
scale = np.random.uniform(1, self.opt['resize_range2'][1])
|
121 |
+
elif updown_type == 'down':
|
122 |
+
scale = np.random.uniform(self.opt['resize_range2'][0], 1)
|
123 |
+
else:
|
124 |
+
scale = 1
|
125 |
+
mode = random.choice(['area', 'bilinear', 'bicubic'])
|
126 |
+
out = F.interpolate(
|
127 |
+
out, size=(int(ori_h / self.opt['scale'] * scale), int(ori_w / self.opt['scale'] * scale)), mode=mode)
|
128 |
+
# add noise
|
129 |
+
gray_noise_prob = self.opt['gray_noise_prob2']
|
130 |
+
if np.random.uniform() < self.opt['gaussian_noise_prob2']:
|
131 |
+
out = random_add_gaussian_noise_pt(
|
132 |
+
out, sigma_range=self.opt['noise_range2'], clip=True, rounds=False, gray_prob=gray_noise_prob)
|
133 |
+
else:
|
134 |
+
out = random_add_poisson_noise_pt(
|
135 |
+
out,
|
136 |
+
scale_range=self.opt['poisson_scale_range2'],
|
137 |
+
gray_prob=gray_noise_prob,
|
138 |
+
clip=True,
|
139 |
+
rounds=False)
|
140 |
+
|
141 |
+
# JPEG compression + the final sinc filter
|
142 |
+
# We also need to resize images to desired sizes. We group [resize back + sinc filter] together
|
143 |
+
# as one operation.
|
144 |
+
# We consider two orders:
|
145 |
+
# 1. [resize back + sinc filter] + JPEG compression
|
146 |
+
# 2. JPEG compression + [resize back + sinc filter]
|
147 |
+
# Empirically, we find other combinations (sinc + JPEG + Resize) will introduce twisted lines.
|
148 |
+
if np.random.uniform() < 0.5:
|
149 |
+
# resize back + the final sinc filter
|
150 |
+
mode = random.choice(['area', 'bilinear', 'bicubic'])
|
151 |
+
out = F.interpolate(out, size=(ori_h // self.opt['scale'], ori_w // self.opt['scale']), mode=mode)
|
152 |
+
out = filter2D(out, self.sinc_kernel)
|
153 |
+
# JPEG compression
|
154 |
+
jpeg_p = out.new_zeros(out.size(0)).uniform_(*self.opt['jpeg_range2'])
|
155 |
+
out = torch.clamp(out, 0, 1)
|
156 |
+
out = self.jpeger(out, quality=jpeg_p)
|
157 |
+
else:
|
158 |
+
# JPEG compression
|
159 |
+
jpeg_p = out.new_zeros(out.size(0)).uniform_(*self.opt['jpeg_range2'])
|
160 |
+
out = torch.clamp(out, 0, 1)
|
161 |
+
out = self.jpeger(out, quality=jpeg_p)
|
162 |
+
# resize back + the final sinc filter
|
163 |
+
mode = random.choice(['area', 'bilinear', 'bicubic'])
|
164 |
+
out = F.interpolate(out, size=(ori_h // self.opt['scale'], ori_w // self.opt['scale']), mode=mode)
|
165 |
+
out = filter2D(out, self.sinc_kernel)
|
166 |
+
|
167 |
+
# clamp and round
|
168 |
+
self.lq = torch.clamp((out * 255.0).round(), 0, 255) / 255.
|
169 |
+
|
170 |
+
# random crop
|
171 |
+
gt_size = self.opt['gt_size']
|
172 |
+
self.gt, self.lq = paired_random_crop(self.gt, self.lq, gt_size, self.opt['scale'])
|
173 |
+
|
174 |
+
# training pair pool
|
175 |
+
self._dequeue_and_enqueue()
|
176 |
+
self.lq = self.lq.contiguous() # for the warning: grad and param do not obey the gradient layout contract
|
177 |
+
else:
|
178 |
+
# for paired training or validation
|
179 |
+
self.lq = data['lq'].to(self.device)
|
180 |
+
if 'gt' in data:
|
181 |
+
self.gt = data['gt'].to(self.device)
|
182 |
+
self.gt_usm = self.usm_sharpener(self.gt)
|
183 |
+
|
184 |
+
def nondist_validation(self, dataloader, current_iter, tb_logger, save_img):
|
185 |
+
# do not use the synthetic process during validation
|
186 |
+
self.is_train = False
|
187 |
+
super(RealESRNetModel, self).nondist_validation(dataloader, current_iter, tb_logger, save_img)
|
188 |
+
self.is_train = True
|
realesrgan/train.py
ADDED
@@ -0,0 +1,11 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
# flake8: noqa
|
2 |
+
import os.path as osp
|
3 |
+
from basicsr.train import train_pipeline
|
4 |
+
|
5 |
+
import realesrgan.archs
|
6 |
+
import realesrgan.data
|
7 |
+
import realesrgan.models
|
8 |
+
|
9 |
+
if __name__ == '__main__':
|
10 |
+
root_path = osp.abspath(osp.join(__file__, osp.pardir, osp.pardir))
|
11 |
+
train_pipeline(root_path)
|
realesrgan/utils.py
ADDED
@@ -0,0 +1,313 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
import cv2
|
2 |
+
import math
|
3 |
+
import numpy as np
|
4 |
+
import os
|
5 |
+
import queue
|
6 |
+
import threading
|
7 |
+
import torch
|
8 |
+
from basicsr.utils.download_util import load_file_from_url
|
9 |
+
from torch.nn import functional as F
|
10 |
+
|
11 |
+
ROOT_DIR = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
|
12 |
+
|
13 |
+
|
14 |
+
class RealESRGANer():
|
15 |
+
"""A helper class for upsampling images with RealESRGAN.
|
16 |
+
|
17 |
+
Args:
|
18 |
+
scale (int): Upsampling scale factor used in the networks. It is usually 2 or 4.
|
19 |
+
model_path (str): The path to the pretrained model. It can be urls (will first download it automatically).
|
20 |
+
model (nn.Module): The defined network. Default: None.
|
21 |
+
tile (int): As too large images result in the out of GPU memory issue, so this tile option will first crop
|
22 |
+
input images into tiles, and then process each of them. Finally, they will be merged into one image.
|
23 |
+
0 denotes for do not use tile. Default: 0.
|
24 |
+
tile_pad (int): The pad size for each tile, to remove border artifacts. Default: 10.
|
25 |
+
pre_pad (int): Pad the input images to avoid border artifacts. Default: 10.
|
26 |
+
half (float): Whether to use half precision during inference. Default: False.
|
27 |
+
"""
|
28 |
+
|
29 |
+
def __init__(self,
|
30 |
+
scale,
|
31 |
+
model_path,
|
32 |
+
dni_weight=None,
|
33 |
+
model=None,
|
34 |
+
tile=0,
|
35 |
+
tile_pad=10,
|
36 |
+
pre_pad=10,
|
37 |
+
half=False,
|
38 |
+
device=None,
|
39 |
+
gpu_id=None):
|
40 |
+
self.scale = scale
|
41 |
+
self.tile_size = tile
|
42 |
+
self.tile_pad = tile_pad
|
43 |
+
self.pre_pad = pre_pad
|
44 |
+
self.mod_scale = None
|
45 |
+
self.half = half
|
46 |
+
|
47 |
+
# initialize model
|
48 |
+
if gpu_id:
|
49 |
+
self.device = torch.device(
|
50 |
+
f'cuda:{gpu_id}' if torch.cuda.is_available() else 'cpu') if device is None else device
|
51 |
+
else:
|
52 |
+
self.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') if device is None else device
|
53 |
+
|
54 |
+
if isinstance(model_path, list):
|
55 |
+
# dni
|
56 |
+
assert len(model_path) == len(dni_weight), 'model_path and dni_weight should have the save length.'
|
57 |
+
loadnet = self.dni(model_path[0], model_path[1], dni_weight)
|
58 |
+
else:
|
59 |
+
# if the model_path starts with https, it will first download models to the folder: weights
|
60 |
+
if model_path.startswith('https://'):
|
61 |
+
model_path = load_file_from_url(
|
62 |
+
url=model_path, model_dir=os.path.join(ROOT_DIR, 'weights'), progress=True, file_name=None)
|
63 |
+
loadnet = torch.load(model_path, map_location=torch.device('cpu'))
|
64 |
+
|
65 |
+
# prefer to use params_ema
|
66 |
+
if 'params_ema' in loadnet:
|
67 |
+
keyname = 'params_ema'
|
68 |
+
else:
|
69 |
+
keyname = 'params'
|
70 |
+
model.load_state_dict(loadnet[keyname], strict=True)
|
71 |
+
|
72 |
+
model.eval()
|
73 |
+
self.model = model.to(self.device)
|
74 |
+
if self.half:
|
75 |
+
self.model = self.model.half()
|
76 |
+
|
77 |
+
def dni(self, net_a, net_b, dni_weight, key='params', loc='cpu'):
|
78 |
+
"""Deep network interpolation.
|
79 |
+
|
80 |
+
``Paper: Deep Network Interpolation for Continuous Imagery Effect Transition``
|
81 |
+
"""
|
82 |
+
net_a = torch.load(net_a, map_location=torch.device(loc))
|
83 |
+
net_b = torch.load(net_b, map_location=torch.device(loc))
|
84 |
+
for k, v_a in net_a[key].items():
|
85 |
+
net_a[key][k] = dni_weight[0] * v_a + dni_weight[1] * net_b[key][k]
|
86 |
+
return net_a
|
87 |
+
|
88 |
+
def pre_process(self, img):
|
89 |
+
"""Pre-process, such as pre-pad and mod pad, so that the images can be divisible
|
90 |
+
"""
|
91 |
+
img = torch.from_numpy(np.transpose(img, (2, 0, 1))).float()
|
92 |
+
self.img = img.unsqueeze(0).to(self.device)
|
93 |
+
if self.half:
|
94 |
+
self.img = self.img.half()
|
95 |
+
|
96 |
+
# pre_pad
|
97 |
+
if self.pre_pad != 0:
|
98 |
+
self.img = F.pad(self.img, (0, self.pre_pad, 0, self.pre_pad), 'reflect')
|
99 |
+
# mod pad for divisible borders
|
100 |
+
if self.scale == 2:
|
101 |
+
self.mod_scale = 2
|
102 |
+
elif self.scale == 1:
|
103 |
+
self.mod_scale = 4
|
104 |
+
if self.mod_scale is not None:
|
105 |
+
self.mod_pad_h, self.mod_pad_w = 0, 0
|
106 |
+
_, _, h, w = self.img.size()
|
107 |
+
if (h % self.mod_scale != 0):
|
108 |
+
self.mod_pad_h = (self.mod_scale - h % self.mod_scale)
|
109 |
+
if (w % self.mod_scale != 0):
|
110 |
+
self.mod_pad_w = (self.mod_scale - w % self.mod_scale)
|
111 |
+
self.img = F.pad(self.img, (0, self.mod_pad_w, 0, self.mod_pad_h), 'reflect')
|
112 |
+
|
113 |
+
def process(self):
|
114 |
+
# model inference
|
115 |
+
self.output = self.model(self.img)
|
116 |
+
|
117 |
+
def tile_process(self):
|
118 |
+
"""It will first crop input images to tiles, and then process each tile.
|
119 |
+
Finally, all the processed tiles are merged into one images.
|
120 |
+
|
121 |
+
Modified from: https://github.com/ata4/esrgan-launcher
|
122 |
+
"""
|
123 |
+
batch, channel, height, width = self.img.shape
|
124 |
+
output_height = height * self.scale
|
125 |
+
output_width = width * self.scale
|
126 |
+
output_shape = (batch, channel, output_height, output_width)
|
127 |
+
|
128 |
+
# start with black image
|
129 |
+
self.output = self.img.new_zeros(output_shape)
|
130 |
+
tiles_x = math.ceil(width / self.tile_size)
|
131 |
+
tiles_y = math.ceil(height / self.tile_size)
|
132 |
+
|
133 |
+
# loop over all tiles
|
134 |
+
for y in range(tiles_y):
|
135 |
+
for x in range(tiles_x):
|
136 |
+
# extract tile from input image
|
137 |
+
ofs_x = x * self.tile_size
|
138 |
+
ofs_y = y * self.tile_size
|
139 |
+
# input tile area on total image
|
140 |
+
input_start_x = ofs_x
|
141 |
+
input_end_x = min(ofs_x + self.tile_size, width)
|
142 |
+
input_start_y = ofs_y
|
143 |
+
input_end_y = min(ofs_y + self.tile_size, height)
|
144 |
+
|
145 |
+
# input tile area on total image with padding
|
146 |
+
input_start_x_pad = max(input_start_x - self.tile_pad, 0)
|
147 |
+
input_end_x_pad = min(input_end_x + self.tile_pad, width)
|
148 |
+
input_start_y_pad = max(input_start_y - self.tile_pad, 0)
|
149 |
+
input_end_y_pad = min(input_end_y + self.tile_pad, height)
|
150 |
+
|
151 |
+
# input tile dimensions
|
152 |
+
input_tile_width = input_end_x - input_start_x
|
153 |
+
input_tile_height = input_end_y - input_start_y
|
154 |
+
tile_idx = y * tiles_x + x + 1
|
155 |
+
input_tile = self.img[:, :, input_start_y_pad:input_end_y_pad, input_start_x_pad:input_end_x_pad]
|
156 |
+
|
157 |
+
# upscale tile
|
158 |
+
try:
|
159 |
+
with torch.no_grad():
|
160 |
+
output_tile = self.model(input_tile)
|
161 |
+
except RuntimeError as error:
|
162 |
+
print('Error', error)
|
163 |
+
print(f'\tTile {tile_idx}/{tiles_x * tiles_y}')
|
164 |
+
|
165 |
+
# output tile area on total image
|
166 |
+
output_start_x = input_start_x * self.scale
|
167 |
+
output_end_x = input_end_x * self.scale
|
168 |
+
output_start_y = input_start_y * self.scale
|
169 |
+
output_end_y = input_end_y * self.scale
|
170 |
+
|
171 |
+
# output tile area without padding
|
172 |
+
output_start_x_tile = (input_start_x - input_start_x_pad) * self.scale
|
173 |
+
output_end_x_tile = output_start_x_tile + input_tile_width * self.scale
|
174 |
+
output_start_y_tile = (input_start_y - input_start_y_pad) * self.scale
|
175 |
+
output_end_y_tile = output_start_y_tile + input_tile_height * self.scale
|
176 |
+
|
177 |
+
# put tile into output image
|
178 |
+
self.output[:, :, output_start_y:output_end_y,
|
179 |
+
output_start_x:output_end_x] = output_tile[:, :, output_start_y_tile:output_end_y_tile,
|
180 |
+
output_start_x_tile:output_end_x_tile]
|
181 |
+
|
182 |
+
def post_process(self):
|
183 |
+
# remove extra pad
|
184 |
+
if self.mod_scale is not None:
|
185 |
+
_, _, h, w = self.output.size()
|
186 |
+
self.output = self.output[:, :, 0:h - self.mod_pad_h * self.scale, 0:w - self.mod_pad_w * self.scale]
|
187 |
+
# remove prepad
|
188 |
+
if self.pre_pad != 0:
|
189 |
+
_, _, h, w = self.output.size()
|
190 |
+
self.output = self.output[:, :, 0:h - self.pre_pad * self.scale, 0:w - self.pre_pad * self.scale]
|
191 |
+
return self.output
|
192 |
+
|
193 |
+
@torch.no_grad()
|
194 |
+
def enhance(self, img, outscale=None, alpha_upsampler='realesrgan'):
|
195 |
+
h_input, w_input = img.shape[0:2]
|
196 |
+
# img: numpy
|
197 |
+
img = img.astype(np.float32)
|
198 |
+
if np.max(img) > 256: # 16-bit image
|
199 |
+
max_range = 65535
|
200 |
+
print('\tInput is a 16-bit image')
|
201 |
+
else:
|
202 |
+
max_range = 255
|
203 |
+
img = img / max_range
|
204 |
+
if len(img.shape) == 2: # gray image
|
205 |
+
img_mode = 'L'
|
206 |
+
img = cv2.cvtColor(img, cv2.COLOR_GRAY2RGB)
|
207 |
+
elif img.shape[2] == 4: # RGBA image with alpha channel
|
208 |
+
img_mode = 'RGBA'
|
209 |
+
alpha = img[:, :, 3]
|
210 |
+
img = img[:, :, 0:3]
|
211 |
+
img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
|
212 |
+
if alpha_upsampler == 'realesrgan':
|
213 |
+
alpha = cv2.cvtColor(alpha, cv2.COLOR_GRAY2RGB)
|
214 |
+
else:
|
215 |
+
img_mode = 'RGB'
|
216 |
+
img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
|
217 |
+
|
218 |
+
# ------------------- process image (without the alpha channel) ------------------- #
|
219 |
+
self.pre_process(img)
|
220 |
+
if self.tile_size > 0:
|
221 |
+
self.tile_process()
|
222 |
+
else:
|
223 |
+
self.process()
|
224 |
+
output_img = self.post_process()
|
225 |
+
output_img = output_img.data.squeeze().float().cpu().clamp_(0, 1).numpy()
|
226 |
+
output_img = np.transpose(output_img[[2, 1, 0], :, :], (1, 2, 0))
|
227 |
+
if img_mode == 'L':
|
228 |
+
output_img = cv2.cvtColor(output_img, cv2.COLOR_BGR2GRAY)
|
229 |
+
|
230 |
+
# ------------------- process the alpha channel if necessary ------------------- #
|
231 |
+
if img_mode == 'RGBA':
|
232 |
+
if alpha_upsampler == 'realesrgan':
|
233 |
+
self.pre_process(alpha)
|
234 |
+
if self.tile_size > 0:
|
235 |
+
self.tile_process()
|
236 |
+
else:
|
237 |
+
self.process()
|
238 |
+
output_alpha = self.post_process()
|
239 |
+
output_alpha = output_alpha.data.squeeze().float().cpu().clamp_(0, 1).numpy()
|
240 |
+
output_alpha = np.transpose(output_alpha[[2, 1, 0], :, :], (1, 2, 0))
|
241 |
+
output_alpha = cv2.cvtColor(output_alpha, cv2.COLOR_BGR2GRAY)
|
242 |
+
else: # use the cv2 resize for alpha channel
|
243 |
+
h, w = alpha.shape[0:2]
|
244 |
+
output_alpha = cv2.resize(alpha, (w * self.scale, h * self.scale), interpolation=cv2.INTER_LINEAR)
|
245 |
+
|
246 |
+
# merge the alpha channel
|
247 |
+
output_img = cv2.cvtColor(output_img, cv2.COLOR_BGR2BGRA)
|
248 |
+
output_img[:, :, 3] = output_alpha
|
249 |
+
|
250 |
+
# ------------------------------ return ------------------------------ #
|
251 |
+
if max_range == 65535: # 16-bit image
|
252 |
+
output = (output_img * 65535.0).round().astype(np.uint16)
|
253 |
+
else:
|
254 |
+
output = (output_img * 255.0).round().astype(np.uint8)
|
255 |
+
|
256 |
+
if outscale is not None and outscale != float(self.scale):
|
257 |
+
output = cv2.resize(
|
258 |
+
output, (
|
259 |
+
int(w_input * outscale),
|
260 |
+
int(h_input * outscale),
|
261 |
+
), interpolation=cv2.INTER_LANCZOS4)
|
262 |
+
|
263 |
+
return output, img_mode
|
264 |
+
|
265 |
+
|
266 |
+
class PrefetchReader(threading.Thread):
|
267 |
+
"""Prefetch images.
|
268 |
+
|
269 |
+
Args:
|
270 |
+
img_list (list[str]): A image list of image paths to be read.
|
271 |
+
num_prefetch_queue (int): Number of prefetch queue.
|
272 |
+
"""
|
273 |
+
|
274 |
+
def __init__(self, img_list, num_prefetch_queue):
|
275 |
+
super().__init__()
|
276 |
+
self.que = queue.Queue(num_prefetch_queue)
|
277 |
+
self.img_list = img_list
|
278 |
+
|
279 |
+
def run(self):
|
280 |
+
for img_path in self.img_list:
|
281 |
+
img = cv2.imread(img_path, cv2.IMREAD_UNCHANGED)
|
282 |
+
self.que.put(img)
|
283 |
+
|
284 |
+
self.que.put(None)
|
285 |
+
|
286 |
+
def __next__(self):
|
287 |
+
next_item = self.que.get()
|
288 |
+
if next_item is None:
|
289 |
+
raise StopIteration
|
290 |
+
return next_item
|
291 |
+
|
292 |
+
def __iter__(self):
|
293 |
+
return self
|
294 |
+
|
295 |
+
|
296 |
+
class IOConsumer(threading.Thread):
|
297 |
+
|
298 |
+
def __init__(self, opt, que, qid):
|
299 |
+
super().__init__()
|
300 |
+
self._queue = que
|
301 |
+
self.qid = qid
|
302 |
+
self.opt = opt
|
303 |
+
|
304 |
+
def run(self):
|
305 |
+
while True:
|
306 |
+
msg = self._queue.get()
|
307 |
+
if isinstance(msg, str) and msg == 'quit':
|
308 |
+
break
|
309 |
+
|
310 |
+
output = msg['output']
|
311 |
+
save_path = msg['save_path']
|
312 |
+
cv2.imwrite(save_path, output)
|
313 |
+
print(f'IO worker {self.qid} is done.')
|
realesrgan/version.py
ADDED
@@ -0,0 +1,5 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
# GENERATED VERSION FILE
|
2 |
+
# TIME: Fri May 24 21:09:51 2024
|
3 |
+
__version__ = '0.3.0'
|
4 |
+
__gitsha__ = 'a4abfb2'
|
5 |
+
version_info = (0, 3, 0)
|
requirements.txt
ADDED
@@ -0,0 +1,11 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
new-basicsr
|
2 |
+
facexlib>=0.2.5
|
3 |
+
gfpgan>=1.3.5
|
4 |
+
numpy
|
5 |
+
opencv-python
|
6 |
+
Pillow
|
7 |
+
torch>=1.7
|
8 |
+
torchvision
|
9 |
+
tqdm
|
10 |
+
streamlit
|
11 |
+
opencv-python
|
run.bat
ADDED
@@ -0,0 +1,7 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
@echo off
|
2 |
+
REM Activate the cuda environment
|
3 |
+
call "%USERPROFILE%\anaconda3\Scripts\activate.bat" cuda
|
4 |
+
REM Change directory to Real-ESRGAN-Web-App
|
5 |
+
cd /d %USERPROFILE%\Real-ESRGAN-Web-App
|
6 |
+
REM Run Streamlit app
|
7 |
+
streamlit run app.py
|
setup.py
ADDED
@@ -0,0 +1,107 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
#!/usr/bin/env python
|
2 |
+
|
3 |
+
from setuptools import find_packages, setup
|
4 |
+
|
5 |
+
import os
|
6 |
+
import subprocess
|
7 |
+
import time
|
8 |
+
|
9 |
+
version_file = 'realesrgan/version.py'
|
10 |
+
|
11 |
+
|
12 |
+
def readme():
|
13 |
+
with open('README.md', encoding='utf-8') as f:
|
14 |
+
content = f.read()
|
15 |
+
return content
|
16 |
+
|
17 |
+
|
18 |
+
def get_git_hash():
|
19 |
+
|
20 |
+
def _minimal_ext_cmd(cmd):
|
21 |
+
# construct minimal environment
|
22 |
+
env = {}
|
23 |
+
for k in ['SYSTEMROOT', 'PATH', 'HOME']:
|
24 |
+
v = os.environ.get(k)
|
25 |
+
if v is not None:
|
26 |
+
env[k] = v
|
27 |
+
# LANGUAGE is used on win32
|
28 |
+
env['LANGUAGE'] = 'C'
|
29 |
+
env['LANG'] = 'C'
|
30 |
+
env['LC_ALL'] = 'C'
|
31 |
+
out = subprocess.Popen(cmd, stdout=subprocess.PIPE, env=env).communicate()[0]
|
32 |
+
return out
|
33 |
+
|
34 |
+
try:
|
35 |
+
out = _minimal_ext_cmd(['git', 'rev-parse', 'HEAD'])
|
36 |
+
sha = out.strip().decode('ascii')
|
37 |
+
except OSError:
|
38 |
+
sha = 'unknown'
|
39 |
+
|
40 |
+
return sha
|
41 |
+
|
42 |
+
|
43 |
+
def get_hash():
|
44 |
+
if os.path.exists('.git'):
|
45 |
+
sha = get_git_hash()[:7]
|
46 |
+
else:
|
47 |
+
sha = 'unknown'
|
48 |
+
|
49 |
+
return sha
|
50 |
+
|
51 |
+
|
52 |
+
def write_version_py():
|
53 |
+
content = """# GENERATED VERSION FILE
|
54 |
+
# TIME: {}
|
55 |
+
__version__ = '{}'
|
56 |
+
__gitsha__ = '{}'
|
57 |
+
version_info = ({})
|
58 |
+
"""
|
59 |
+
sha = get_hash()
|
60 |
+
with open('VERSION', 'r') as f:
|
61 |
+
SHORT_VERSION = f.read().strip()
|
62 |
+
VERSION_INFO = ', '.join([x if x.isdigit() else f'"{x}"' for x in SHORT_VERSION.split('.')])
|
63 |
+
|
64 |
+
version_file_str = content.format(time.asctime(), SHORT_VERSION, sha, VERSION_INFO)
|
65 |
+
with open(version_file, 'w') as f:
|
66 |
+
f.write(version_file_str)
|
67 |
+
|
68 |
+
|
69 |
+
def get_version():
|
70 |
+
with open(version_file, 'r') as f:
|
71 |
+
exec(compile(f.read(), version_file, 'exec'))
|
72 |
+
return locals()['__version__']
|
73 |
+
|
74 |
+
|
75 |
+
def get_requirements(filename='requirements.txt'):
|
76 |
+
here = os.path.dirname(os.path.realpath(__file__))
|
77 |
+
with open(os.path.join(here, filename), 'r') as f:
|
78 |
+
requires = [line.replace('\n', '') for line in f.readlines()]
|
79 |
+
return requires
|
80 |
+
|
81 |
+
|
82 |
+
if __name__ == '__main__':
|
83 |
+
write_version_py()
|
84 |
+
setup(
|
85 |
+
name='realesrgan',
|
86 |
+
version=get_version(),
|
87 |
+
description='Real-ESRGAN aims at developing Practical Algorithms for General Image Restoration',
|
88 |
+
long_description=readme(),
|
89 |
+
long_description_content_type='text/markdown',
|
90 |
+
author='Xintao Wang',
|
91 |
+
author_email='[email protected]',
|
92 |
+
keywords='computer vision, pytorch, image restoration, super-resolution, esrgan, real-esrgan',
|
93 |
+
url='https://github.com/xinntao/Real-ESRGAN',
|
94 |
+
include_package_data=True,
|
95 |
+
packages=find_packages(exclude=('options', 'datasets', 'experiments', 'results', 'tb_logger', 'wandb')),
|
96 |
+
classifiers=[
|
97 |
+
'Development Status :: 4 - Beta',
|
98 |
+
'License :: OSI Approved :: Apache Software License',
|
99 |
+
'Operating System :: OS Independent',
|
100 |
+
'Programming Language :: Python :: 3',
|
101 |
+
'Programming Language :: Python :: 3.7',
|
102 |
+
'Programming Language :: Python :: 3.8',
|
103 |
+
],
|
104 |
+
license='BSD-3-Clause License',
|
105 |
+
setup_requires=['cython', 'numpy'],
|
106 |
+
install_requires=get_requirements(),
|
107 |
+
zip_safe=False)
|