Compare commits
32 Commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
9976a34454 | ||
|
|
424a09457b | ||
|
|
52f77e74a8 | ||
|
|
bfa4678bef | ||
|
|
68f9f2445e | ||
|
|
7840a3d16a | ||
|
|
b28958cdf2 | ||
|
|
667e34e7ca | ||
|
|
978def19a6 | ||
|
|
a7153c7fce | ||
|
|
00116244cb | ||
|
|
571b89257a | ||
|
|
bed7df7d99 | ||
|
|
fb3ff055e4 | ||
|
|
9ef97853f9 | ||
|
|
58fea8db69 | ||
|
|
3c6cf5290e | ||
|
|
64ad194dda | ||
|
|
5745599813 | ||
|
|
3ce0c97e89 | ||
|
|
13186ac2c2 | ||
|
|
4356ba0578 | ||
|
|
32a4fa1772 | ||
|
|
18ebf723f2 | ||
|
|
1f83ce5432 | ||
|
|
bef5e3cabd | ||
|
|
064df9956b | ||
|
|
52eab16d11 | ||
|
|
94ae626008 | ||
|
|
9baa0b3d00 | ||
|
|
f932289af1 | ||
|
|
f59a0c66ec |
34
.github/workflows/no-response.yml
vendored
Normal file
34
.github/workflows/no-response.yml
vendored
Normal file
@@ -0,0 +1,34 @@
|
||||
name: No Response
|
||||
|
||||
# Modified from: https://raw.githubusercontent.com/github/docs/main/.github/workflows/no-response.yaml
|
||||
|
||||
# **What it does**: Closes issues that don't have enough information to be
|
||||
# actionable.
|
||||
# **Why we have it**: To remove the need for maintainers to remember to check
|
||||
# back on issues periodically to see if contributors have
|
||||
# responded.
|
||||
# **Who does it impact**: Everyone that works on docs or docs-internal.
|
||||
|
||||
on:
|
||||
issue_comment:
|
||||
types: [created]
|
||||
|
||||
schedule:
|
||||
# Schedule for five minutes after the hour every hour
|
||||
- cron: '5 * * * *'
|
||||
|
||||
jobs:
|
||||
noResponse:
|
||||
runs-on: ubuntu-latest
|
||||
steps:
|
||||
- uses: lee-dohm/no-response@v0.5.0
|
||||
with:
|
||||
token: ${{ github.token }}
|
||||
closeComment: >
|
||||
This issue has been automatically closed because there has been no response
|
||||
to our request for more information from the original author. With only the
|
||||
information that is currently in the issue, we don't have enough information
|
||||
to take action. Please reach out if you have or find the answers we need so
|
||||
that we can investigate further.
|
||||
If you still have questions, please improve your description and re-open it.
|
||||
Thanks :-)
|
||||
33
.github/workflows/publish-pip.yml
vendored
Normal file
33
.github/workflows/publish-pip.yml
vendored
Normal file
@@ -0,0 +1,33 @@
|
||||
name: PyPI Publish
|
||||
|
||||
on: push
|
||||
|
||||
jobs:
|
||||
build-n-publish:
|
||||
runs-on: ubuntu-latest
|
||||
if: startsWith(github.event.ref, 'refs/tags')
|
||||
|
||||
steps:
|
||||
- uses: actions/checkout@v2
|
||||
- name: Set up Python 3.8
|
||||
uses: actions/setup-python@v1
|
||||
with:
|
||||
python-version: 3.8
|
||||
- name: Upgrade pip
|
||||
run: pip install pip --upgrade
|
||||
- name: Install PyTorch (cpu)
|
||||
run: pip install torch==1.7.0+cpu torchvision==0.8.1+cpu -f https://download.pytorch.org/whl/torch_stable.html
|
||||
- name: Install dependencies
|
||||
run: |
|
||||
pip install basicsr
|
||||
pip install facexlib
|
||||
pip install gfpgan
|
||||
pip install -r requirements.txt
|
||||
- name: Build and install
|
||||
run: rm -rf .eggs && pip install -e .
|
||||
- name: Build for distribution
|
||||
run: python setup.py sdist bdist_wheel
|
||||
- name: Publish distribution to PyPI
|
||||
uses: pypa/gh-action-pypi-publish@master
|
||||
with:
|
||||
password: ${{ secrets.PYPI_API_TOKEN }}
|
||||
4
.github/workflows/pylint.yml
vendored
4
.github/workflows/pylint.yml
vendored
@@ -26,5 +26,5 @@ jobs:
|
||||
- name: Lint
|
||||
run: |
|
||||
flake8 .
|
||||
isort --check-only --diff data/ models/ inference_realesrgan.py
|
||||
yapf -r -d data/ models/ inference_realesrgan.py
|
||||
isort --check-only --diff realesrgan/ scripts/ inference_realesrgan.py setup.py
|
||||
yapf -r -d realesrgan/ scripts/ inference_realesrgan.py setup.py
|
||||
|
||||
11
.gitignore
vendored
11
.gitignore
vendored
@@ -1,4 +1,13 @@
|
||||
.vscode
|
||||
# ignored folders
|
||||
datasets/*
|
||||
experiments/*
|
||||
results/*
|
||||
tb_logger/*
|
||||
wandb/*
|
||||
tmp/*
|
||||
realesrgan/weights/*
|
||||
|
||||
version.py
|
||||
|
||||
# Byte-compiled / optimized / DLL files
|
||||
__pycache__/
|
||||
|
||||
19
.vscode/settings.json
vendored
Normal file
19
.vscode/settings.json
vendored
Normal file
@@ -0,0 +1,19 @@
|
||||
{
|
||||
"files.trimTrailingWhitespace": true,
|
||||
"editor.wordWrap": "on",
|
||||
"editor.rulers": [
|
||||
80,
|
||||
120
|
||||
],
|
||||
"editor.renderWhitespace": "all",
|
||||
"editor.renderControlCharacters": true,
|
||||
"python.formatting.provider": "yapf",
|
||||
"python.formatting.yapfArgs": [
|
||||
"--style",
|
||||
"{BASED_ON_STYLE = pep8, BLANK_LINE_BEFORE_NESTED_CLASS_OR_DEF = true, SPLIT_BEFORE_EXPRESSION_AFTER_OPENING_PAREN = true, COLUMN_LIMIT = 120}"
|
||||
],
|
||||
"python.linting.flake8Enabled": true,
|
||||
"python.linting.flake8Args": [
|
||||
"max-line-length=120"
|
||||
],
|
||||
}
|
||||
9
FAQ.md
Normal file
9
FAQ.md
Normal file
@@ -0,0 +1,9 @@
|
||||
# FAQ
|
||||
|
||||
1. **What is the difference of `--netscale` and `outscale`?**
|
||||
|
||||
A: TODO.
|
||||
|
||||
1. **How to select models?**
|
||||
|
||||
A: TODO.
|
||||
8
MANIFEST.in
Normal file
8
MANIFEST.in
Normal file
@@ -0,0 +1,8 @@
|
||||
include assets/*
|
||||
include inputs/*
|
||||
include scripts/*.py
|
||||
include inference_realesrgan.py
|
||||
include VERSION
|
||||
include LICENSE
|
||||
include requirements.txt
|
||||
include realesrgan/weights/README.md
|
||||
28
README.md
28
README.md
@@ -1,21 +1,32 @@
|
||||
# Real-ESRGAN
|
||||
|
||||
[](https://github.com/xinntao/Real-ESRGAN/releases)
|
||||
[](https://pypi.org/project/realesrgan/)
|
||||
[](https://github.com/xinntao/Real-ESRGAN/issues)
|
||||
[](https://github.com/xinntao/Real-ESRGAN/blob/master/LICENSE)
|
||||
[](https://github.com/xinntao/Real-ESRGAN/blob/master/.github/workflows/pylint.yml)
|
||||
[](https://github.com/xinntao/Real-ESRGAN/blob/master/.github/workflows/publish-pip.yml)
|
||||
|
||||
1. [Colab Demo](https://colab.research.google.com/drive/1k2Zod6kSHEvraybHl50Lys0LerhyTMCo?usp=sharing) for Real-ESRGAN <a href="https://colab.research.google.com/drive/1k2Zod6kSHEvraybHl50Lys0LerhyTMCo?usp=sharing"><img src="https://colab.research.google.com/assets/colab-badge.svg" alt="google colab logo"></a>.
|
||||
2. [Portable Windows executable file](https://github.com/xinntao/Real-ESRGAN/releases). You can find more information [here](#Portable-executable-files).
|
||||
2. Portable [Windows](https://github.com/xinntao/Real-ESRGAN/releases/download/v0.1.2/realesrgan-ncnn-vulkan-20210801-windows.zip) / [Linux](https://github.com/xinntao/Real-ESRGAN/releases/download/v0.1.2/realesrgan-ncnn-vulkan-20210801-ubuntu.zip) / [MacOS](https://github.com/xinntao/Real-ESRGAN/releases/download/v0.1.2/realesrgan-ncnn-vulkan-20210801-macos.zip) **executable files for Intel/AMD/Nvidia GPU**. You can find more information [here](#Portable-executable-files).
|
||||
|
||||
Real-ESRGAN aims at developing **Practical Algorithms for General Image Restoration**.<br>
|
||||
We extend the powerful ESRGAN to a practical restoration application (namely, Real-ESRGAN), which is trained with pure synthetic data.
|
||||
|
||||
:triangular_flag_on_post: **Updates**
|
||||
:question: Frequently Asked Questions can be found in [FAQ.md](FAQ.md).
|
||||
|
||||
:triangular_flag_on_post: **Updates**
|
||||
- :white_check_mark: Integrate [GFPGAN](https://github.com/TencentARC/GFPGAN) to support **face enhancement**.
|
||||
- :white_check_mark: 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).
|
||||
- :white_check_mark: Support arbitrary scale with `--outscale` (It actually further resizes outputs with `LANCZOS4`). Add *RealESRGAN_x2plus.pth* model.
|
||||
- :white_check_mark: [The inference code](inference_realesrgan.py) supports: 1) **tile** options; 2) images with **alpha channel**; 3) **gray** images; 4) **16-bit** images.
|
||||
- :white_check_mark: The training codes have been released. A detailed guide can be found in [Training.md](Training.md).
|
||||
|
||||
If Real-ESRGAN is helpful in your photos/projects, please help to :star: this repo. Thanks:blush: <br>
|
||||
Other recommended projects:   :arrow_forward: [GFPGAN](https://github.com/TencentARC/GFPGAN)   :arrow_forward: [BasicSR](https://github.com/xinntao/BasicSR)   :arrow_forward: [facexlib](https://github.com/xinntao/facexlib)
|
||||
|
||||
---
|
||||
|
||||
### :book: Real-ESRGAN: Training Real-World Blind Super-Resolution with Pure Synthetic Data
|
||||
|
||||
> [[Paper](https://arxiv.org/abs/2107.10833)]   [Project Page]   [Demo] <br>
|
||||
@@ -41,7 +52,7 @@ Here is a TODO list in the near future:
|
||||
|
||||
- [ ] optimize for human faces
|
||||
- [ ] optimize for texts
|
||||
- [ ] optimize for animation images
|
||||
- [ ] optimize for anime images [in progress]
|
||||
- [ ] support more scales
|
||||
- [ ] support controllable restoration strength
|
||||
|
||||
@@ -52,11 +63,11 @@ If you have some images that Real-ESRGAN could not well restored, please also op
|
||||
|
||||
### Portable executable files
|
||||
|
||||
You can download **Windows executable files** from https://github.com/xinntao/Real-ESRGAN/releases/download/v0.1.1/RealESRGAN-ncnn-vulkan-20210725-windows.zip
|
||||
You can download [Windows](https://github.com/xinntao/Real-ESRGAN/releases/download/v0.1.2/realesrgan-ncnn-vulkan-20210801-windows.zip) / [Linux](https://github.com/xinntao/Real-ESRGAN/releases/download/v0.1.2/realesrgan-ncnn-vulkan-20210801-ubuntu.zip) / [MacOS](https://github.com/xinntao/Real-ESRGAN/releases/download/v0.1.2/realesrgan-ncnn-vulkan-20210801-macos.zip) **executable files for Intel/AMD/Nvidia GPU**.
|
||||
|
||||
This executable file is **portable** and includes all the binaries and models required. No CUDA or PyTorch environment is needed.<br>
|
||||
|
||||
You can simply run the following command:
|
||||
You can simply run the following command (the Windows example, more information is in the README.md of each executable files):
|
||||
|
||||
```bash
|
||||
./realesrgan-ncnn-vulkan.exe -i input.jpg -o output.png
|
||||
@@ -96,7 +107,11 @@ This executable file is based on the wonderful [Tencent/ncnn](https://github.com
|
||||
# Install basicsr - https://github.com/xinntao/BasicSR
|
||||
# We use BasicSR for both training and inference
|
||||
pip install basicsr
|
||||
# facexlib and gfpgan are for face enhancement
|
||||
pip install facexlib
|
||||
pip install gfpgan
|
||||
pip install -r requirements.txt
|
||||
python setup.py develop
|
||||
```
|
||||
|
||||
## :zap: Quick Inference
|
||||
@@ -112,7 +127,7 @@ wget https://github.com/xinntao/Real-ESRGAN/releases/download/v0.1.0/RealESRGAN_
|
||||
Inference!
|
||||
|
||||
```bash
|
||||
python inference_realesrgan.py --model_path experiments/pretrained_models/RealESRGAN_x4plus.pth --input inputs
|
||||
python inference_realesrgan.py --model_path experiments/pretrained_models/RealESRGAN_x4plus.pth --input inputs --face_enhance
|
||||
```
|
||||
|
||||
Results are in the `results` folder
|
||||
@@ -121,6 +136,7 @@ Results are in the `results` folder
|
||||
|
||||
- [RealESRGAN-x4plus](https://github.com/xinntao/Real-ESRGAN/releases/download/v0.1.0/RealESRGAN_x4plus.pth)
|
||||
- [RealESRNet-x4plus](https://github.com/xinntao/Real-ESRGAN/releases/download/v0.1.1/RealESRNet_x4plus.pth)
|
||||
- [RealESRGAN-x2plus](https://github.com/xinntao/Real-ESRGAN/releases/download/v0.2.1/RealESRGAN_x2plus.pth)
|
||||
- [official ESRGAN-x4](https://github.com/xinntao/Real-ESRGAN/releases/download/v0.1.1/ESRGAN_SRx4_DF2KOST_official-ff704c30.pth)
|
||||
|
||||
## :computer: Training
|
||||
|
||||
12
Training.md
12
Training.md
@@ -1,7 +1,7 @@
|
||||
# :computer: How to Train Real-ESRGAN
|
||||
|
||||
The training codes have been released. <br>
|
||||
Note that the codes have a lot of refactoring. So there may be some bugs/performance drops. Welcome to report issues and I will also retrain the models.
|
||||
Note that the codes have a lot of refactoring. So there may be some bugs/performance drops. Welcome to report bugs/issues.
|
||||
|
||||
## Overview
|
||||
|
||||
@@ -44,7 +44,7 @@ DF2K_HR_sub/000001_s003.png
|
||||
name: DF2K+OST
|
||||
type: RealESRGANDataset
|
||||
dataroot_gt: datasets/DF2K # modify to the root path of your folder
|
||||
meta_info: data/meta_info/meta_info_DF2Kmultiscale+OST_sub.txt # modify to your own generate meta info txt
|
||||
meta_info: realesrgan/meta_info/meta_info_DF2Kmultiscale+OST_sub.txt # modify to your own generate meta info txt
|
||||
io_backend:
|
||||
type: disk
|
||||
```
|
||||
@@ -76,12 +76,12 @@ DF2K_HR_sub/000001_s003.png
|
||||
1. Before the formal training, you may run in the `--debug` mode to see whether everything is OK. We use four GPUs for training:
|
||||
```bash
|
||||
CUDA_VISIBLE_DEVICES=0,1,2,3 \
|
||||
python -m torch.distributed.launch --nproc_per_node=4 --master_port=4321 train.py -opt options/train_realesrnet_x4plus.yml --launcher pytorch --debug
|
||||
python -m torch.distributed.launch --nproc_per_node=4 --master_port=4321 realesrgan/train.py -opt options/train_realesrnet_x4plus.yml --launcher pytorch --debug
|
||||
```
|
||||
1. The formal training. We use four GPUs for training. We use the `--auto_resume` argument to automatically resume the training if necessary.
|
||||
```bash
|
||||
CUDA_VISIBLE_DEVICES=0,1,2,3 \
|
||||
python -m torch.distributed.launch --nproc_per_node=4 --master_port=4321 train.py -opt options/train_realesrnet_x4plus.yml --launcher pytorch --auto_resume
|
||||
python -m torch.distributed.launch --nproc_per_node=4 --master_port=4321 realesrgan/train.py -opt options/train_realesrnet_x4plus.yml --launcher pytorch --auto_resume
|
||||
```
|
||||
|
||||
## Train Real-ESRGAN
|
||||
@@ -91,10 +91,10 @@ DF2K_HR_sub/000001_s003.png
|
||||
1. Before the formal training, you may run in the `--debug` mode to see whether everything is OK. We use four GPUs for training:
|
||||
```bash
|
||||
CUDA_VISIBLE_DEVICES=0,1,2,3 \
|
||||
python -m torch.distributed.launch --nproc_per_node=4 --master_port=4321 train.py -opt options/train_realesrgan_x4plus.yml --launcher pytorch --debug
|
||||
python -m torch.distributed.launch --nproc_per_node=4 --master_port=4321 realesrgan/train.py -opt options/train_realesrgan_x4plus.yml --launcher pytorch --debug
|
||||
```
|
||||
1. The formal training. We use four GPUs for training. We use the `--auto_resume` argument to automatically resume the training if necessary.
|
||||
```bash
|
||||
CUDA_VISIBLE_DEVICES=0,1,2,3 \
|
||||
python -m torch.distributed.launch --nproc_per_node=4 --master_port=4321 train.py -opt options/train_realesrgan_x4plus.yml --launcher pytorch --auto_resume
|
||||
python -m torch.distributed.launch --nproc_per_node=4 --master_port=4321 realesrgan/train.py -opt options/train_realesrgan_x4plus.yml --launcher pytorch --auto_resume
|
||||
```
|
||||
|
||||
BIN
assets/teaser-text.png
Normal file
BIN
assets/teaser-text.png
Normal file
Binary file not shown.
|
After Width: | Height: | Size: 546 KiB |
@@ -1,12 +1,9 @@
|
||||
import argparse
|
||||
import cv2
|
||||
import glob
|
||||
import math
|
||||
import numpy as np
|
||||
import os
|
||||
import torch
|
||||
from basicsr.archs.rrdbnet_arch import RRDBNet
|
||||
from torch.nn import functional as F
|
||||
|
||||
from realesrgan import RealESRGANer
|
||||
|
||||
|
||||
def main():
|
||||
@@ -17,11 +14,15 @@ def main():
|
||||
type=str,
|
||||
default='experiments/pretrained_models/RealESRGAN_x4plus.pth',
|
||||
help='Path to the pre-trained model')
|
||||
parser.add_argument('--scale', type=int, default=4, help='Upsample scale factor')
|
||||
parser.add_argument('--output', type=str, default='results', help='Output folder')
|
||||
parser.add_argument('--netscale', type=int, default=4, help='Upsample scale factor of the network')
|
||||
parser.add_argument('--outscale', type=float, default=4, help='The final upsampling scale of the image')
|
||||
parser.add_argument('--suffix', type=str, default='out', help='Suffix of the restored image')
|
||||
parser.add_argument('--tile', type=int, default=0, help='Tile size, 0 for no tile during testing')
|
||||
parser.add_argument('--tile', type=int, default=800, help='Tile size, 0 for no tile during testing')
|
||||
parser.add_argument('--tile_pad', type=int, default=10, help='Tile padding')
|
||||
parser.add_argument('--pre_pad', type=int, default=0, help='Pre padding size at each border')
|
||||
parser.add_argument('--face_enhance', action='store_true', help='Use GFPGAN to enhance face')
|
||||
parser.add_argument('--half', action='store_true', help='Use half precision during inference')
|
||||
parser.add_argument(
|
||||
'--alpha_upsampler',
|
||||
type=str,
|
||||
@@ -35,8 +36,23 @@ def main():
|
||||
args = parser.parse_args()
|
||||
|
||||
upsampler = RealESRGANer(
|
||||
scale=args.scale, model_path=args.model_path, tile=args.tile, tile_pad=args.tile_pad, pre_pad=args.pre_pad)
|
||||
os.makedirs('results/', exist_ok=True)
|
||||
scale=args.netscale,
|
||||
model_path=args.model_path,
|
||||
tile=args.tile,
|
||||
tile_pad=args.tile_pad,
|
||||
pre_pad=args.pre_pad,
|
||||
half=args.half)
|
||||
|
||||
if args.face_enhance:
|
||||
from gfpgan import GFPGANer
|
||||
face_enhancer = GFPGANer(
|
||||
model_path='https://github.com/TencentARC/GFPGAN/releases/download/v0.2.0/GFPGANCleanv1-NoCE-C2.pth',
|
||||
upscale=args.outscale,
|
||||
arch='clean',
|
||||
channel_multiplier=2,
|
||||
bg_upsampler=upsampler)
|
||||
os.makedirs(args.output, exist_ok=True)
|
||||
|
||||
if os.path.isfile(args.input):
|
||||
paths = [args.input]
|
||||
else:
|
||||
@@ -46,198 +62,38 @@ def main():
|
||||
imgname, extension = os.path.splitext(os.path.basename(path))
|
||||
print('Testing', idx, imgname)
|
||||
|
||||
# ------------------------------ read image ------------------------------ #
|
||||
img = cv2.imread(path, cv2.IMREAD_UNCHANGED).astype(np.float32)
|
||||
if np.max(img) > 255: # 16-bit image
|
||||
max_range = 65535
|
||||
print('\tInput is a 16-bit image')
|
||||
else:
|
||||
max_range = 255
|
||||
img = img / max_range
|
||||
if len(img.shape) == 2: # gray image
|
||||
img_mode = 'L'
|
||||
img = cv2.cvtColor(img, cv2.COLOR_GRAY2RGB)
|
||||
elif img.shape[2] == 4: # RGBA image with alpha channel
|
||||
img = cv2.imread(path, cv2.IMREAD_UNCHANGED)
|
||||
if len(img.shape) == 3 and img.shape[2] == 4:
|
||||
img_mode = 'RGBA'
|
||||
alpha = img[:, :, 3]
|
||||
img = img[:, :, 0:3]
|
||||
img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
|
||||
if args.alpha_upsampler == 'realesrgan':
|
||||
alpha = cv2.cvtColor(alpha, cv2.COLOR_GRAY2RGB)
|
||||
else:
|
||||
img_mode = 'RGB'
|
||||
img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
|
||||
img_mode = None
|
||||
|
||||
# ------------------- process image (without the alpha channel) ------------------- #
|
||||
upsampler.pre_process(img)
|
||||
if args.tile:
|
||||
upsampler.tile_process()
|
||||
h, w = img.shape[0:2]
|
||||
if max(h, w) > 1000 and args.netscale == 4:
|
||||
import warnings
|
||||
warnings.warn('The input image is large, try X2 model for better performace.')
|
||||
if max(h, w) < 500 and args.netscale == 2:
|
||||
import warnings
|
||||
warnings.warn('The input image is small, try X4 model for better performace.')
|
||||
|
||||
try:
|
||||
if args.face_enhance:
|
||||
_, _, output = face_enhancer.enhance(img, has_aligned=False, only_center_face=False, paste_back=True)
|
||||
else:
|
||||
upsampler.process()
|
||||
output_img = upsampler.post_process()
|
||||
output_img = output_img.data.squeeze().float().cpu().clamp_(0, 1).numpy()
|
||||
output_img = np.transpose(output_img[[2, 1, 0], :, :], (1, 2, 0))
|
||||
if img_mode == 'L':
|
||||
output_img = cv2.cvtColor(output_img, cv2.COLOR_BGR2GRAY)
|
||||
|
||||
# ------------------- process the alpha channel if necessary ------------------- #
|
||||
if img_mode == 'RGBA':
|
||||
if args.alpha_upsampler == 'realesrgan':
|
||||
upsampler.pre_process(alpha)
|
||||
if args.tile:
|
||||
upsampler.tile_process()
|
||||
output, _ = upsampler.enhance(img, outscale=args.outscale)
|
||||
except Exception as error:
|
||||
print('Error', error)
|
||||
print('If you encounter CUDA out of memory, try to set --tile with a smaller number.')
|
||||
else:
|
||||
upsampler.process()
|
||||
output_alpha = upsampler.post_process()
|
||||
output_alpha = output_alpha.data.squeeze().float().cpu().clamp_(0, 1).numpy()
|
||||
output_alpha = np.transpose(output_alpha[[2, 1, 0], :, :], (1, 2, 0))
|
||||
output_alpha = cv2.cvtColor(output_alpha, cv2.COLOR_BGR2GRAY)
|
||||
else:
|
||||
h, w = alpha.shape[0:2]
|
||||
output_alpha = cv2.resize(alpha, (w * args.scale, h * args.scale), interpolation=cv2.INTER_LINEAR)
|
||||
|
||||
# merge the alpha channel
|
||||
output_img = cv2.cvtColor(output_img, cv2.COLOR_BGR2BGRA)
|
||||
output_img[:, :, 3] = output_alpha
|
||||
|
||||
# ------------------------------ save image ------------------------------ #
|
||||
if args.ext == 'auto':
|
||||
extension = extension[1:]
|
||||
else:
|
||||
extension = args.ext
|
||||
if img_mode == 'RGBA': # RGBA images should be saved in png format
|
||||
extension = 'png'
|
||||
save_path = f'results/{imgname}_{args.suffix}.{extension}'
|
||||
if max_range == 65535: # 16-bit image
|
||||
output = (output_img * 65535.0).round().astype(np.uint16)
|
||||
else:
|
||||
output = (output_img * 255.0).round().astype(np.uint8)
|
||||
save_path = os.path.join(args.output, f'{imgname}_{args.suffix}.{extension}')
|
||||
cv2.imwrite(save_path, output)
|
||||
|
||||
|
||||
class RealESRGANer():
|
||||
|
||||
def __init__(self, scale, model_path, tile=0, tile_pad=10, pre_pad=10):
|
||||
self.scale = scale
|
||||
self.tile_size = tile
|
||||
self.tile_pad = tile_pad
|
||||
self.pre_pad = pre_pad
|
||||
self.mod_scale = None
|
||||
|
||||
# initialize model
|
||||
self.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
|
||||
model = RRDBNet(num_in_ch=3, num_out_ch=3, num_feat=64, num_block=23, num_grow_ch=32, scale=scale)
|
||||
loadnet = torch.load(model_path)
|
||||
if 'params_ema' in loadnet:
|
||||
keyname = 'params_ema'
|
||||
else:
|
||||
keyname = 'params'
|
||||
model.load_state_dict(loadnet[keyname], strict=True)
|
||||
model.eval()
|
||||
self.model = model.to(self.device)
|
||||
|
||||
def pre_process(self, img):
|
||||
img = torch.from_numpy(np.transpose(img, (2, 0, 1))).float()
|
||||
self.img = img.unsqueeze(0).to(self.device)
|
||||
|
||||
# pre_pad
|
||||
if self.pre_pad != 0:
|
||||
self.img = F.pad(self.img, (0, self.pre_pad, 0, self.pre_pad), 'reflect')
|
||||
# mod pad
|
||||
if self.scale == 2:
|
||||
self.mod_scale = 2
|
||||
elif self.scale == 1:
|
||||
self.mod_scale = 4
|
||||
if self.mod_scale is not None:
|
||||
self.mod_pad_h, self.mod_pad_w = 0, 0
|
||||
_, _, h, w = self.img.size()
|
||||
if (h % self.mod_scale != 0):
|
||||
self.mod_pad_h = (self.mod_scale - h % self.mod_scale)
|
||||
if (w % self.mod_scale != 0):
|
||||
self.mod_pad_w = (self.mod_scale - w % self.mod_scale)
|
||||
self.img = F.pad(self.img, (0, self.mod_pad_w, 0, self.mod_pad_h), 'reflect')
|
||||
|
||||
def process(self):
|
||||
try:
|
||||
# inference
|
||||
with torch.no_grad():
|
||||
self.output = self.model(self.img)
|
||||
except Exception as error:
|
||||
print('Error', error)
|
||||
|
||||
def tile_process(self):
|
||||
"""Modified from: https://github.com/ata4/esrgan-launcher
|
||||
"""
|
||||
batch, channel, height, width = self.img.shape
|
||||
output_height = height * self.scale
|
||||
output_width = width * self.scale
|
||||
output_shape = (batch, channel, output_height, output_width)
|
||||
|
||||
# start with black image
|
||||
self.output = self.img.new_zeros(output_shape)
|
||||
tiles_x = math.ceil(width / self.tile_size)
|
||||
tiles_y = math.ceil(height / self.tile_size)
|
||||
|
||||
# loop over all tiles
|
||||
for y in range(tiles_y):
|
||||
for x in range(tiles_x):
|
||||
# extract tile from input image
|
||||
ofs_x = x * self.tile_size
|
||||
ofs_y = y * self.tile_size
|
||||
# input tile area on total image
|
||||
input_start_x = ofs_x
|
||||
input_end_x = min(ofs_x + self.tile_size, width)
|
||||
input_start_y = ofs_y
|
||||
input_end_y = min(ofs_y + self.tile_size, height)
|
||||
|
||||
# input tile area on total image with padding
|
||||
input_start_x_pad = max(input_start_x - self.tile_pad, 0)
|
||||
input_end_x_pad = min(input_end_x + self.tile_pad, width)
|
||||
input_start_y_pad = max(input_start_y - self.tile_pad, 0)
|
||||
input_end_y_pad = min(input_end_y + self.tile_pad, height)
|
||||
|
||||
# input tile dimensions
|
||||
input_tile_width = input_end_x - input_start_x
|
||||
input_tile_height = input_end_y - input_start_y
|
||||
tile_idx = y * tiles_x + x + 1
|
||||
input_tile = self.img[:, :, input_start_y_pad:input_end_y_pad, input_start_x_pad:input_end_x_pad]
|
||||
|
||||
# upscale tile
|
||||
try:
|
||||
with torch.no_grad():
|
||||
output_tile = self.model(input_tile)
|
||||
except Exception as error:
|
||||
print('Error', error)
|
||||
print(f'\tTile {tile_idx}/{tiles_x * tiles_y}')
|
||||
|
||||
# output tile area on total image
|
||||
output_start_x = input_start_x * self.scale
|
||||
output_end_x = input_end_x * self.scale
|
||||
output_start_y = input_start_y * self.scale
|
||||
output_end_y = input_end_y * self.scale
|
||||
|
||||
# output tile area without padding
|
||||
output_start_x_tile = (input_start_x - input_start_x_pad) * self.scale
|
||||
output_end_x_tile = output_start_x_tile + input_tile_width * self.scale
|
||||
output_start_y_tile = (input_start_y - input_start_y_pad) * self.scale
|
||||
output_end_y_tile = output_start_y_tile + input_tile_height * self.scale
|
||||
|
||||
# put tile into output image
|
||||
self.output[:, :, output_start_y:output_end_y,
|
||||
output_start_x:output_end_x] = output_tile[:, :, output_start_y_tile:output_end_y_tile,
|
||||
output_start_x_tile:output_end_x_tile]
|
||||
|
||||
def post_process(self):
|
||||
# remove extra pad
|
||||
if self.mod_scale is not None:
|
||||
_, _, h, w = self.output.size()
|
||||
self.output = self.output[:, :, 0:h - self.mod_pad_h * self.scale, 0:w - self.mod_pad_w * self.scale]
|
||||
# remove prepad
|
||||
if self.pre_pad != 0:
|
||||
_, _, h, w = self.output.size()
|
||||
self.output = self.output[:, :, 0:h - self.pre_pad * self.scale, 0:w - self.pre_pad * self.scale]
|
||||
return self.output
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
main()
|
||||
|
||||
189
options/finetune_realesrgan_x4plus.yml
Normal file
189
options/finetune_realesrgan_x4plus.yml
Normal file
@@ -0,0 +1,189 @@
|
||||
# general settings
|
||||
name: finetune_RealESRGANx4plus_400k_B12G4
|
||||
model_type: RealESRGANModel
|
||||
scale: 4
|
||||
num_gpu: 4
|
||||
manual_seed: 0
|
||||
|
||||
# ----------------- options for synthesizing training data in RealESRGANModel ----------------- #
|
||||
# USM the ground-truth
|
||||
l1_gt_usm: True
|
||||
percep_gt_usm: True
|
||||
gan_gt_usm: False
|
||||
|
||||
# the first degradation process
|
||||
resize_prob: [0.2, 0.7, 0.1] # up, down, keep
|
||||
resize_range: [0.15, 1.5]
|
||||
gaussian_noise_prob: 0.5
|
||||
noise_range: [1, 30]
|
||||
poisson_scale_range: [0.05, 3]
|
||||
gray_noise_prob: 0.4
|
||||
jpeg_range: [30, 95]
|
||||
|
||||
# the second degradation process
|
||||
second_blur_prob: 0.8
|
||||
resize_prob2: [0.3, 0.4, 0.3] # up, down, keep
|
||||
resize_range2: [0.3, 1.2]
|
||||
gaussian_noise_prob2: 0.5
|
||||
noise_range2: [1, 25]
|
||||
poisson_scale_range2: [0.05, 2.5]
|
||||
gray_noise_prob2: 0.4
|
||||
jpeg_range2: [30, 95]
|
||||
|
||||
gt_size: 256
|
||||
queue_size: 180
|
||||
|
||||
# dataset and data loader settings
|
||||
datasets:
|
||||
train:
|
||||
name: DF2K+OST
|
||||
type: RealESRGANDataset
|
||||
dataroot_gt: datasets/DF2K
|
||||
meta_info: datasets/DF2K/meta_info/meta_info_DF2Kmultiscale+OST_sub.txt
|
||||
io_backend:
|
||||
type: disk
|
||||
|
||||
blur_kernel_size: 21
|
||||
kernel_list: ['iso', 'aniso', 'generalized_iso', 'generalized_aniso', 'plateau_iso', 'plateau_aniso']
|
||||
kernel_prob: [0.45, 0.25, 0.12, 0.03, 0.12, 0.03]
|
||||
sinc_prob: 0.1
|
||||
blur_sigma: [0.2, 3]
|
||||
betag_range: [0.5, 4]
|
||||
betap_range: [1, 2]
|
||||
|
||||
blur_kernel_size2: 21
|
||||
kernel_list2: ['iso', 'aniso', 'generalized_iso', 'generalized_aniso', 'plateau_iso', 'plateau_aniso']
|
||||
kernel_prob2: [0.45, 0.25, 0.12, 0.03, 0.12, 0.03]
|
||||
sinc_prob2: 0.1
|
||||
blur_sigma2: [0.2, 1.5]
|
||||
betag_range2: [0.5, 4]
|
||||
betap_range2: [1, 2]
|
||||
|
||||
final_sinc_prob: 0.8
|
||||
|
||||
gt_size: 256
|
||||
use_hflip: True
|
||||
use_rot: False
|
||||
|
||||
# data loader
|
||||
use_shuffle: true
|
||||
num_worker_per_gpu: 5
|
||||
batch_size_per_gpu: 12
|
||||
dataset_enlarge_ratio: 1
|
||||
prefetch_mode: ~
|
||||
|
||||
# Uncomment these for validation
|
||||
# val:
|
||||
# name: validation
|
||||
# type: PairedImageDataset
|
||||
# dataroot_gt: path_to_gt
|
||||
# dataroot_lq: path_to_lq
|
||||
# io_backend:
|
||||
# type: disk
|
||||
|
||||
# network structures
|
||||
network_g:
|
||||
type: RRDBNet
|
||||
num_in_ch: 3
|
||||
num_out_ch: 3
|
||||
num_feat: 64
|
||||
num_block: 23
|
||||
num_grow_ch: 32
|
||||
|
||||
|
||||
network_d:
|
||||
type: UNetDiscriminatorSN
|
||||
num_in_ch: 3
|
||||
num_feat: 64
|
||||
skip_connection: True
|
||||
|
||||
# path
|
||||
path:
|
||||
# use the pre-trained Real-ESRNet model
|
||||
pretrain_network_g: experiments/pretrained_models/RealESRNet_x4plus.pth
|
||||
param_key_g: params_ema
|
||||
strict_load_g: true
|
||||
pretrain_network_d: experiments/pretrained_models/RealESRGAN_x4plus_netD.pth
|
||||
param_key_d: params
|
||||
strict_load_d: true
|
||||
resume_state: ~
|
||||
|
||||
# training settings
|
||||
train:
|
||||
ema_decay: 0.999
|
||||
optim_g:
|
||||
type: Adam
|
||||
lr: !!float 1e-4
|
||||
weight_decay: 0
|
||||
betas: [0.9, 0.99]
|
||||
optim_d:
|
||||
type: Adam
|
||||
lr: !!float 1e-4
|
||||
weight_decay: 0
|
||||
betas: [0.9, 0.99]
|
||||
|
||||
scheduler:
|
||||
type: MultiStepLR
|
||||
milestones: [400000]
|
||||
gamma: 0.5
|
||||
|
||||
total_iter: 400000
|
||||
warmup_iter: -1 # no warm up
|
||||
|
||||
# losses
|
||||
pixel_opt:
|
||||
type: L1Loss
|
||||
loss_weight: 1.0
|
||||
reduction: mean
|
||||
# perceptual loss (content and style losses)
|
||||
perceptual_opt:
|
||||
type: PerceptualLoss
|
||||
layer_weights:
|
||||
# before relu
|
||||
'conv1_2': 0.1
|
||||
'conv2_2': 0.1
|
||||
'conv3_4': 1
|
||||
'conv4_4': 1
|
||||
'conv5_4': 1
|
||||
vgg_type: vgg19
|
||||
use_input_norm: true
|
||||
perceptual_weight: !!float 1.0
|
||||
style_weight: 0
|
||||
range_norm: false
|
||||
criterion: l1
|
||||
# gan loss
|
||||
gan_opt:
|
||||
type: GANLoss
|
||||
gan_type: vanilla
|
||||
real_label_val: 1.0
|
||||
fake_label_val: 0.0
|
||||
loss_weight: !!float 1e-1
|
||||
|
||||
net_d_iters: 1
|
||||
net_d_init_iters: 0
|
||||
|
||||
# Uncomment these for validation
|
||||
# validation settings
|
||||
# val:
|
||||
# val_freq: !!float 5e3
|
||||
# save_img: True
|
||||
|
||||
# metrics:
|
||||
# psnr: # metric name, can be arbitrary
|
||||
# type: calculate_psnr
|
||||
# crop_border: 4
|
||||
# test_y_channel: false
|
||||
|
||||
# logging settings
|
||||
logger:
|
||||
print_freq: 100
|
||||
save_checkpoint_freq: !!float 5e3
|
||||
use_tb_logger: true
|
||||
wandb:
|
||||
project: ~
|
||||
resume_id: ~
|
||||
|
||||
# dist training settings
|
||||
dist_params:
|
||||
backend: nccl
|
||||
port: 29500
|
||||
187
options/train_realesrgan_x2plus.yml
Normal file
187
options/train_realesrgan_x2plus.yml
Normal file
@@ -0,0 +1,187 @@
|
||||
# general settings
|
||||
name: train_RealESRGANx2plus_400k_B12G4
|
||||
model_type: RealESRGANModel
|
||||
scale: 2
|
||||
num_gpu: 4
|
||||
manual_seed: 0
|
||||
|
||||
# ----------------- options for synthesizing training data in RealESRGANModel ----------------- #
|
||||
# USM the ground-truth
|
||||
l1_gt_usm: True
|
||||
percep_gt_usm: True
|
||||
gan_gt_usm: False
|
||||
|
||||
# the first degradation process
|
||||
resize_prob: [0.2, 0.7, 0.1] # up, down, keep
|
||||
resize_range: [0.15, 1.5]
|
||||
gaussian_noise_prob: 0.5
|
||||
noise_range: [1, 30]
|
||||
poisson_scale_range: [0.05, 3]
|
||||
gray_noise_prob: 0.4
|
||||
jpeg_range: [30, 95]
|
||||
|
||||
# the second degradation process
|
||||
second_blur_prob: 0.8
|
||||
resize_prob2: [0.3, 0.4, 0.3] # up, down, keep
|
||||
resize_range2: [0.3, 1.2]
|
||||
gaussian_noise_prob2: 0.5
|
||||
noise_range2: [1, 25]
|
||||
poisson_scale_range2: [0.05, 2.5]
|
||||
gray_noise_prob2: 0.4
|
||||
jpeg_range2: [30, 95]
|
||||
|
||||
gt_size: 256
|
||||
queue_size: 180
|
||||
|
||||
# dataset and data loader settings
|
||||
datasets:
|
||||
train:
|
||||
name: DF2K+OST
|
||||
type: RealESRGANDataset
|
||||
dataroot_gt: datasets/DF2K
|
||||
meta_info: datasets/DF2K/meta_info/meta_info_DF2Kmultiscale+OST_sub.txt
|
||||
io_backend:
|
||||
type: disk
|
||||
|
||||
blur_kernel_size: 21
|
||||
kernel_list: ['iso', 'aniso', 'generalized_iso', 'generalized_aniso', 'plateau_iso', 'plateau_aniso']
|
||||
kernel_prob: [0.45, 0.25, 0.12, 0.03, 0.12, 0.03]
|
||||
sinc_prob: 0.1
|
||||
blur_sigma: [0.2, 3]
|
||||
betag_range: [0.5, 4]
|
||||
betap_range: [1, 2]
|
||||
|
||||
blur_kernel_size2: 21
|
||||
kernel_list2: ['iso', 'aniso', 'generalized_iso', 'generalized_aniso', 'plateau_iso', 'plateau_aniso']
|
||||
kernel_prob2: [0.45, 0.25, 0.12, 0.03, 0.12, 0.03]
|
||||
sinc_prob2: 0.1
|
||||
blur_sigma2: [0.2, 1.5]
|
||||
betag_range2: [0.5, 4]
|
||||
betap_range2: [1, 2]
|
||||
|
||||
final_sinc_prob: 0.8
|
||||
|
||||
gt_size: 256
|
||||
use_hflip: True
|
||||
use_rot: False
|
||||
|
||||
# data loader
|
||||
use_shuffle: true
|
||||
num_worker_per_gpu: 5
|
||||
batch_size_per_gpu: 12
|
||||
dataset_enlarge_ratio: 1
|
||||
prefetch_mode: ~
|
||||
|
||||
# Uncomment these for validation
|
||||
# val:
|
||||
# name: validation
|
||||
# type: PairedImageDataset
|
||||
# dataroot_gt: path_to_gt
|
||||
# dataroot_lq: path_to_lq
|
||||
# io_backend:
|
||||
# type: disk
|
||||
|
||||
# network structures
|
||||
network_g:
|
||||
type: RRDBNet
|
||||
num_in_ch: 3
|
||||
num_out_ch: 3
|
||||
num_feat: 64
|
||||
num_block: 23
|
||||
num_grow_ch: 32
|
||||
scale: 2
|
||||
|
||||
|
||||
network_d:
|
||||
type: UNetDiscriminatorSN
|
||||
num_in_ch: 3
|
||||
num_feat: 64
|
||||
skip_connection: True
|
||||
|
||||
# path
|
||||
path:
|
||||
# use the pre-trained Real-ESRNet model
|
||||
pretrain_network_g: experiments/pretrained_models/RealESRNet_x2plus.pth
|
||||
param_key_g: params_ema
|
||||
strict_load_g: true
|
||||
resume_state: ~
|
||||
|
||||
# training settings
|
||||
train:
|
||||
ema_decay: 0.999
|
||||
optim_g:
|
||||
type: Adam
|
||||
lr: !!float 1e-4
|
||||
weight_decay: 0
|
||||
betas: [0.9, 0.99]
|
||||
optim_d:
|
||||
type: Adam
|
||||
lr: !!float 1e-4
|
||||
weight_decay: 0
|
||||
betas: [0.9, 0.99]
|
||||
|
||||
scheduler:
|
||||
type: MultiStepLR
|
||||
milestones: [400000]
|
||||
gamma: 0.5
|
||||
|
||||
total_iter: 400000
|
||||
warmup_iter: -1 # no warm up
|
||||
|
||||
# losses
|
||||
pixel_opt:
|
||||
type: L1Loss
|
||||
loss_weight: 1.0
|
||||
reduction: mean
|
||||
# perceptual loss (content and style losses)
|
||||
perceptual_opt:
|
||||
type: PerceptualLoss
|
||||
layer_weights:
|
||||
# before relu
|
||||
'conv1_2': 0.1
|
||||
'conv2_2': 0.1
|
||||
'conv3_4': 1
|
||||
'conv4_4': 1
|
||||
'conv5_4': 1
|
||||
vgg_type: vgg19
|
||||
use_input_norm: true
|
||||
perceptual_weight: !!float 1.0
|
||||
style_weight: 0
|
||||
range_norm: false
|
||||
criterion: l1
|
||||
# gan loss
|
||||
gan_opt:
|
||||
type: GANLoss
|
||||
gan_type: vanilla
|
||||
real_label_val: 1.0
|
||||
fake_label_val: 0.0
|
||||
loss_weight: !!float 1e-1
|
||||
|
||||
net_d_iters: 1
|
||||
net_d_init_iters: 0
|
||||
|
||||
# Uncomment these for validation
|
||||
# validation settings
|
||||
# val:
|
||||
# val_freq: !!float 5e3
|
||||
# save_img: True
|
||||
|
||||
# metrics:
|
||||
# psnr: # metric name, can be arbitrary
|
||||
# type: calculate_psnr
|
||||
# crop_border: 4
|
||||
# test_y_channel: false
|
||||
|
||||
# logging settings
|
||||
logger:
|
||||
print_freq: 100
|
||||
save_checkpoint_freq: !!float 5e3
|
||||
use_tb_logger: true
|
||||
wandb:
|
||||
project: ~
|
||||
resume_id: ~
|
||||
|
||||
# dist training settings
|
||||
dist_params:
|
||||
backend: nccl
|
||||
port: 29500
|
||||
@@ -1,5 +1,5 @@
|
||||
# general settings
|
||||
name: train_RealESRGANx4plus_400k_B12G4_fromRealESRNet
|
||||
name: train_RealESRGANx4plus_400k_B12G4
|
||||
model_type: RealESRGANModel
|
||||
scale: 4
|
||||
num_gpu: 4
|
||||
@@ -39,7 +39,7 @@ datasets:
|
||||
name: DF2K+OST
|
||||
type: RealESRGANDataset
|
||||
dataroot_gt: datasets/DF2K
|
||||
meta_info: data/meta_info/meta_info_DF2Kmultiscale+OST_sub.txt
|
||||
meta_info: datasets/DF2K/meta_info/meta_info_DF2Kmultiscale+OST_sub.txt
|
||||
io_backend:
|
||||
type: disk
|
||||
|
||||
@@ -100,7 +100,7 @@ network_d:
|
||||
# path
|
||||
path:
|
||||
# use the pre-trained Real-ESRNet model
|
||||
pretrain_network_g: experiments/train_RealESRNetx4plus_1000k_B12G4_fromESRGAN/model/net_g_1000000.pth
|
||||
pretrain_network_g: experiments/pretrained_models/RealESRNet_x4plus.pth
|
||||
param_key_g: params_ema
|
||||
strict_load_g: true
|
||||
resume_state: ~
|
||||
|
||||
145
options/train_realesrnet_x2plus.yml
Normal file
145
options/train_realesrnet_x2plus.yml
Normal file
@@ -0,0 +1,145 @@
|
||||
# general settings
|
||||
name: train_RealESRNetx2plus_1000k_B12G4
|
||||
model_type: RealESRNetModel
|
||||
scale: 2
|
||||
num_gpu: 4
|
||||
manual_seed: 0
|
||||
|
||||
# ----------------- options for synthesizing training data in RealESRNetModel ----------------- #
|
||||
gt_usm: True # USM the ground-truth
|
||||
|
||||
# the first degradation process
|
||||
resize_prob: [0.2, 0.7, 0.1] # up, down, keep
|
||||
resize_range: [0.15, 1.5]
|
||||
gaussian_noise_prob: 0.5
|
||||
noise_range: [1, 30]
|
||||
poisson_scale_range: [0.05, 3]
|
||||
gray_noise_prob: 0.4
|
||||
jpeg_range: [30, 95]
|
||||
|
||||
# the second degradation process
|
||||
second_blur_prob: 0.8
|
||||
resize_prob2: [0.3, 0.4, 0.3] # up, down, keep
|
||||
resize_range2: [0.3, 1.2]
|
||||
gaussian_noise_prob2: 0.5
|
||||
noise_range2: [1, 25]
|
||||
poisson_scale_range2: [0.05, 2.5]
|
||||
gray_noise_prob2: 0.4
|
||||
jpeg_range2: [30, 95]
|
||||
|
||||
gt_size: 256
|
||||
queue_size: 180
|
||||
|
||||
# dataset and data loader settings
|
||||
datasets:
|
||||
train:
|
||||
name: DF2K+OST
|
||||
type: RealESRGANDataset
|
||||
dataroot_gt: datasets/DF2K
|
||||
meta_info: datasets/DF2K/meta_info/meta_info_DF2Kmultiscale+OST_sub.txt
|
||||
io_backend:
|
||||
type: disk
|
||||
|
||||
blur_kernel_size: 21
|
||||
kernel_list: ['iso', 'aniso', 'generalized_iso', 'generalized_aniso', 'plateau_iso', 'plateau_aniso']
|
||||
kernel_prob: [0.45, 0.25, 0.12, 0.03, 0.12, 0.03]
|
||||
sinc_prob: 0.1
|
||||
blur_sigma: [0.2, 3]
|
||||
betag_range: [0.5, 4]
|
||||
betap_range: [1, 2]
|
||||
|
||||
blur_kernel_size2: 21
|
||||
kernel_list2: ['iso', 'aniso', 'generalized_iso', 'generalized_aniso', 'plateau_iso', 'plateau_aniso']
|
||||
kernel_prob2: [0.45, 0.25, 0.12, 0.03, 0.12, 0.03]
|
||||
sinc_prob2: 0.1
|
||||
blur_sigma2: [0.2, 1.5]
|
||||
betag_range2: [0.5, 4]
|
||||
betap_range2: [1, 2]
|
||||
|
||||
final_sinc_prob: 0.8
|
||||
|
||||
gt_size: 256
|
||||
use_hflip: True
|
||||
use_rot: False
|
||||
|
||||
# data loader
|
||||
use_shuffle: true
|
||||
num_worker_per_gpu: 5
|
||||
batch_size_per_gpu: 12
|
||||
dataset_enlarge_ratio: 1
|
||||
prefetch_mode: ~
|
||||
|
||||
# Uncomment these for validation
|
||||
# val:
|
||||
# name: validation
|
||||
# type: PairedImageDataset
|
||||
# dataroot_gt: path_to_gt
|
||||
# dataroot_lq: path_to_lq
|
||||
# io_backend:
|
||||
# type: disk
|
||||
|
||||
# network structures
|
||||
network_g:
|
||||
type: RRDBNet
|
||||
num_in_ch: 3
|
||||
num_out_ch: 3
|
||||
num_feat: 64
|
||||
num_block: 23
|
||||
num_grow_ch: 32
|
||||
scale: 2
|
||||
|
||||
# path
|
||||
path:
|
||||
pretrain_network_g: experiments/pretrained_models/RealESRGAN_x4plus.pth
|
||||
param_key_g: params_ema
|
||||
strict_load_g: False
|
||||
resume_state: ~
|
||||
|
||||
# training settings
|
||||
train:
|
||||
ema_decay: 0.999
|
||||
optim_g:
|
||||
type: Adam
|
||||
lr: !!float 2e-4
|
||||
weight_decay: 0
|
||||
betas: [0.9, 0.99]
|
||||
|
||||
scheduler:
|
||||
type: MultiStepLR
|
||||
milestones: [1000000]
|
||||
gamma: 0.5
|
||||
|
||||
total_iter: 1000000
|
||||
warmup_iter: -1 # no warm up
|
||||
|
||||
# losses
|
||||
pixel_opt:
|
||||
type: L1Loss
|
||||
loss_weight: 1.0
|
||||
reduction: mean
|
||||
|
||||
# Uncomment these for validation
|
||||
# validation settings
|
||||
# val:
|
||||
# val_freq: !!float 5e3
|
||||
# save_img: True
|
||||
|
||||
# metrics:
|
||||
# psnr: # metric name, can be arbitrary
|
||||
# type: calculate_psnr
|
||||
# crop_border: 4
|
||||
# test_y_channel: false
|
||||
|
||||
# logging settings
|
||||
logger:
|
||||
print_freq: 100
|
||||
save_checkpoint_freq: !!float 5e3
|
||||
use_tb_logger: true
|
||||
wandb:
|
||||
project: ~
|
||||
resume_id: ~
|
||||
|
||||
# dist training settings
|
||||
dist_params:
|
||||
backend: nccl
|
||||
port: 29500
|
||||
@@ -1,5 +1,5 @@
|
||||
# general settings
|
||||
name: train_RealESRNetx4plus_1000k_B12G4_fromESRGAN
|
||||
name: train_RealESRNetx4plus_1000k_B12G4
|
||||
model_type: RealESRNetModel
|
||||
scale: 4
|
||||
num_gpu: 4
|
||||
@@ -36,7 +36,7 @@ datasets:
|
||||
name: DF2K+OST
|
||||
type: RealESRGANDataset
|
||||
dataroot_gt: datasets/DF2K
|
||||
meta_info: data/meta_info/meta_info_DF2Kmultiscale+OST_sub.txt
|
||||
meta_info: datasets/DF2K/meta_info/meta_info_DF2Kmultiscale+OST_sub.txt
|
||||
io_backend:
|
||||
type: disk
|
||||
|
||||
|
||||
6
realesrgan/__init__.py
Normal file
6
realesrgan/__init__.py
Normal file
@@ -0,0 +1,6 @@
|
||||
# flake8: noqa
|
||||
from .archs import *
|
||||
from .data import *
|
||||
from .models import *
|
||||
from .utils import *
|
||||
from .version import __gitsha__, __version__
|
||||
@@ -7,4 +7,4 @@ from os import path as osp
|
||||
arch_folder = osp.dirname(osp.abspath(__file__))
|
||||
arch_filenames = [osp.splitext(osp.basename(v))[0] for v in scandir(arch_folder) if v.endswith('_arch.py')]
|
||||
# import all the arch modules
|
||||
_arch_modules = [importlib.import_module(f'archs.{file_name}') for file_name in arch_filenames]
|
||||
_arch_modules = [importlib.import_module(f'realesrgan.archs.{file_name}') for file_name in arch_filenames]
|
||||
@@ -7,4 +7,4 @@ from os import path as osp
|
||||
data_folder = osp.dirname(osp.abspath(__file__))
|
||||
dataset_filenames = [osp.splitext(osp.basename(v))[0] for v in scandir(data_folder) if v.endswith('_dataset.py')]
|
||||
# import all the dataset modules
|
||||
_dataset_modules = [importlib.import_module(f'data.{file_name}') for file_name in dataset_filenames]
|
||||
_dataset_modules = [importlib.import_module(f'realesrgan.data.{file_name}') for file_name in dataset_filenames]
|
||||
@@ -7,4 +7,4 @@ from os import path as osp
|
||||
model_folder = osp.dirname(osp.abspath(__file__))
|
||||
model_filenames = [osp.splitext(osp.basename(v))[0] for v in scandir(model_folder) if v.endswith('_model.py')]
|
||||
# import all the model modules
|
||||
_model_modules = [importlib.import_module(f'models.{file_name}') for file_name in model_filenames]
|
||||
_model_modules = [importlib.import_module(f'realesrgan.models.{file_name}') for file_name in model_filenames]
|
||||
11
realesrgan/train.py
Normal file
11
realesrgan/train.py
Normal file
@@ -0,0 +1,11 @@
|
||||
# flake8: noqa
|
||||
import os.path as osp
|
||||
from basicsr.train import train_pipeline
|
||||
|
||||
import realesrgan.archs
|
||||
import realesrgan.data
|
||||
import realesrgan.models
|
||||
|
||||
if __name__ == '__main__':
|
||||
root_path = osp.abspath(osp.join(__file__, osp.pardir, osp.pardir))
|
||||
train_pipeline(root_path)
|
||||
231
realesrgan/utils.py
Normal file
231
realesrgan/utils.py
Normal file
@@ -0,0 +1,231 @@
|
||||
import cv2
|
||||
import math
|
||||
import numpy as np
|
||||
import os
|
||||
import torch
|
||||
from basicsr.archs.rrdbnet_arch import RRDBNet
|
||||
from torch.hub import download_url_to_file, get_dir
|
||||
from torch.nn import functional as F
|
||||
from urllib.parse import urlparse
|
||||
|
||||
ROOT_DIR = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
|
||||
|
||||
|
||||
class RealESRGANer():
|
||||
|
||||
def __init__(self, scale, model_path, tile=0, tile_pad=10, pre_pad=10, half=False):
|
||||
self.scale = scale
|
||||
self.tile_size = tile
|
||||
self.tile_pad = tile_pad
|
||||
self.pre_pad = pre_pad
|
||||
self.mod_scale = None
|
||||
self.half = half
|
||||
|
||||
# initialize model
|
||||
self.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
|
||||
model = RRDBNet(num_in_ch=3, num_out_ch=3, num_feat=64, num_block=23, num_grow_ch=32, scale=scale)
|
||||
|
||||
if model_path.startswith('https://'):
|
||||
model_path = load_file_from_url(
|
||||
url=model_path, model_dir='realesrgan/weights', progress=True, file_name=None)
|
||||
loadnet = torch.load(model_path)
|
||||
if 'params_ema' in loadnet:
|
||||
keyname = 'params_ema'
|
||||
else:
|
||||
keyname = 'params'
|
||||
model.load_state_dict(loadnet[keyname], strict=True)
|
||||
model.eval()
|
||||
self.model = model.to(self.device)
|
||||
if self.half:
|
||||
self.model = self.model.half()
|
||||
|
||||
def pre_process(self, img):
|
||||
img = torch.from_numpy(np.transpose(img, (2, 0, 1))).float()
|
||||
self.img = img.unsqueeze(0).to(self.device)
|
||||
if self.half:
|
||||
self.img = self.img.half()
|
||||
|
||||
# pre_pad
|
||||
if self.pre_pad != 0:
|
||||
self.img = F.pad(self.img, (0, self.pre_pad, 0, self.pre_pad), 'reflect')
|
||||
# mod pad
|
||||
if self.scale == 2:
|
||||
self.mod_scale = 2
|
||||
elif self.scale == 1:
|
||||
self.mod_scale = 4
|
||||
if self.mod_scale is not None:
|
||||
self.mod_pad_h, self.mod_pad_w = 0, 0
|
||||
_, _, h, w = self.img.size()
|
||||
if (h % self.mod_scale != 0):
|
||||
self.mod_pad_h = (self.mod_scale - h % self.mod_scale)
|
||||
if (w % self.mod_scale != 0):
|
||||
self.mod_pad_w = (self.mod_scale - w % self.mod_scale)
|
||||
self.img = F.pad(self.img, (0, self.mod_pad_w, 0, self.mod_pad_h), 'reflect')
|
||||
|
||||
def process(self):
|
||||
self.output = self.model(self.img)
|
||||
|
||||
def tile_process(self):
|
||||
"""Modified from: https://github.com/ata4/esrgan-launcher
|
||||
"""
|
||||
batch, channel, height, width = self.img.shape
|
||||
output_height = height * self.scale
|
||||
output_width = width * self.scale
|
||||
output_shape = (batch, channel, output_height, output_width)
|
||||
|
||||
# start with black image
|
||||
self.output = self.img.new_zeros(output_shape)
|
||||
tiles_x = math.ceil(width / self.tile_size)
|
||||
tiles_y = math.ceil(height / self.tile_size)
|
||||
|
||||
# loop over all tiles
|
||||
for y in range(tiles_y):
|
||||
for x in range(tiles_x):
|
||||
# extract tile from input image
|
||||
ofs_x = x * self.tile_size
|
||||
ofs_y = y * self.tile_size
|
||||
# input tile area on total image
|
||||
input_start_x = ofs_x
|
||||
input_end_x = min(ofs_x + self.tile_size, width)
|
||||
input_start_y = ofs_y
|
||||
input_end_y = min(ofs_y + self.tile_size, height)
|
||||
|
||||
# input tile area on total image with padding
|
||||
input_start_x_pad = max(input_start_x - self.tile_pad, 0)
|
||||
input_end_x_pad = min(input_end_x + self.tile_pad, width)
|
||||
input_start_y_pad = max(input_start_y - self.tile_pad, 0)
|
||||
input_end_y_pad = min(input_end_y + self.tile_pad, height)
|
||||
|
||||
# input tile dimensions
|
||||
input_tile_width = input_end_x - input_start_x
|
||||
input_tile_height = input_end_y - input_start_y
|
||||
tile_idx = y * tiles_x + x + 1
|
||||
input_tile = self.img[:, :, input_start_y_pad:input_end_y_pad, input_start_x_pad:input_end_x_pad]
|
||||
|
||||
# upscale tile
|
||||
try:
|
||||
with torch.no_grad():
|
||||
output_tile = self.model(input_tile)
|
||||
except Exception as error:
|
||||
print('Error', error)
|
||||
print(f'\tTile {tile_idx}/{tiles_x * tiles_y}')
|
||||
|
||||
# output tile area on total image
|
||||
output_start_x = input_start_x * self.scale
|
||||
output_end_x = input_end_x * self.scale
|
||||
output_start_y = input_start_y * self.scale
|
||||
output_end_y = input_end_y * self.scale
|
||||
|
||||
# output tile area without padding
|
||||
output_start_x_tile = (input_start_x - input_start_x_pad) * self.scale
|
||||
output_end_x_tile = output_start_x_tile + input_tile_width * self.scale
|
||||
output_start_y_tile = (input_start_y - input_start_y_pad) * self.scale
|
||||
output_end_y_tile = output_start_y_tile + input_tile_height * self.scale
|
||||
|
||||
# put tile into output image
|
||||
self.output[:, :, output_start_y:output_end_y,
|
||||
output_start_x:output_end_x] = output_tile[:, :, output_start_y_tile:output_end_y_tile,
|
||||
output_start_x_tile:output_end_x_tile]
|
||||
|
||||
def post_process(self):
|
||||
# remove extra pad
|
||||
if self.mod_scale is not None:
|
||||
_, _, h, w = self.output.size()
|
||||
self.output = self.output[:, :, 0:h - self.mod_pad_h * self.scale, 0:w - self.mod_pad_w * self.scale]
|
||||
# remove prepad
|
||||
if self.pre_pad != 0:
|
||||
_, _, h, w = self.output.size()
|
||||
self.output = self.output[:, :, 0:h - self.pre_pad * self.scale, 0:w - self.pre_pad * self.scale]
|
||||
return self.output
|
||||
|
||||
@torch.no_grad()
|
||||
def enhance(self, img, outscale=None, alpha_upsampler='realesrgan'):
|
||||
h_input, w_input = img.shape[0:2]
|
||||
# img: numpy
|
||||
img = img.astype(np.float32)
|
||||
if np.max(img) > 256: # 16-bit image
|
||||
max_range = 65535
|
||||
print('\tInput is a 16-bit image')
|
||||
else:
|
||||
max_range = 255
|
||||
img = img / max_range
|
||||
if len(img.shape) == 2: # gray image
|
||||
img_mode = 'L'
|
||||
img = cv2.cvtColor(img, cv2.COLOR_GRAY2RGB)
|
||||
elif img.shape[2] == 4: # RGBA image with alpha channel
|
||||
img_mode = 'RGBA'
|
||||
alpha = img[:, :, 3]
|
||||
img = img[:, :, 0:3]
|
||||
img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
|
||||
if alpha_upsampler == 'realesrgan':
|
||||
alpha = cv2.cvtColor(alpha, cv2.COLOR_GRAY2RGB)
|
||||
else:
|
||||
img_mode = 'RGB'
|
||||
img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
|
||||
|
||||
# ------------------- process image (without the alpha channel) ------------------- #
|
||||
self.pre_process(img)
|
||||
if self.tile_size > 0:
|
||||
self.tile_process()
|
||||
else:
|
||||
self.process()
|
||||
output_img = self.post_process()
|
||||
output_img = output_img.data.squeeze().float().cpu().clamp_(0, 1).numpy()
|
||||
output_img = np.transpose(output_img[[2, 1, 0], :, :], (1, 2, 0))
|
||||
if img_mode == 'L':
|
||||
output_img = cv2.cvtColor(output_img, cv2.COLOR_BGR2GRAY)
|
||||
|
||||
# ------------------- process the alpha channel if necessary ------------------- #
|
||||
if img_mode == 'RGBA':
|
||||
if alpha_upsampler == 'realesrgan':
|
||||
self.pre_process(alpha)
|
||||
if self.tile_size > 0:
|
||||
self.tile_process()
|
||||
else:
|
||||
self.process()
|
||||
output_alpha = self.post_process()
|
||||
output_alpha = output_alpha.data.squeeze().float().cpu().clamp_(0, 1).numpy()
|
||||
output_alpha = np.transpose(output_alpha[[2, 1, 0], :, :], (1, 2, 0))
|
||||
output_alpha = cv2.cvtColor(output_alpha, cv2.COLOR_BGR2GRAY)
|
||||
else:
|
||||
h, w = alpha.shape[0:2]
|
||||
output_alpha = cv2.resize(alpha, (w * self.scale, h * self.scale), interpolation=cv2.INTER_LINEAR)
|
||||
|
||||
# merge the alpha channel
|
||||
output_img = cv2.cvtColor(output_img, cv2.COLOR_BGR2BGRA)
|
||||
output_img[:, :, 3] = output_alpha
|
||||
|
||||
# ------------------------------ return ------------------------------ #
|
||||
if max_range == 65535: # 16-bit image
|
||||
output = (output_img * 65535.0).round().astype(np.uint16)
|
||||
else:
|
||||
output = (output_img * 255.0).round().astype(np.uint8)
|
||||
|
||||
if outscale is not None and outscale != float(self.scale):
|
||||
output = cv2.resize(
|
||||
output, (
|
||||
int(w_input * outscale),
|
||||
int(h_input * outscale),
|
||||
), interpolation=cv2.INTER_LANCZOS4)
|
||||
|
||||
return output, img_mode
|
||||
|
||||
|
||||
def load_file_from_url(url, model_dir=None, progress=True, file_name=None):
|
||||
"""Ref:https://github.com/1adrianb/face-alignment/blob/master/face_alignment/utils.py
|
||||
"""
|
||||
if model_dir is None:
|
||||
hub_dir = get_dir()
|
||||
model_dir = os.path.join(hub_dir, 'checkpoints')
|
||||
|
||||
os.makedirs(os.path.join(ROOT_DIR, model_dir), exist_ok=True)
|
||||
|
||||
parts = urlparse(url)
|
||||
filename = os.path.basename(parts.path)
|
||||
if file_name is not None:
|
||||
filename = file_name
|
||||
cached_file = os.path.abspath(os.path.join(ROOT_DIR, model_dir, filename))
|
||||
if not os.path.exists(cached_file):
|
||||
print(f'Downloading: "{url}" to {cached_file}\n')
|
||||
download_url_to_file(url, cached_file, hash_prefix=None, progress=progress)
|
||||
return cached_file
|
||||
3
realesrgan/weights/README.md
Normal file
3
realesrgan/weights/README.md
Normal file
@@ -0,0 +1,3 @@
|
||||
# Weights
|
||||
|
||||
Put the downloaded weights to this folder.
|
||||
@@ -1,4 +1,7 @@
|
||||
basicsr
|
||||
cv2
|
||||
basicsr>=1.3.3.11
|
||||
facexlib>=0.2.0.3
|
||||
gfpgan>=0.2.1
|
||||
numpy
|
||||
opencv-python
|
||||
Pillow
|
||||
torch>=1.7
|
||||
|
||||
40
scripts/generate_meta_info.py
Normal file
40
scripts/generate_meta_info.py
Normal file
@@ -0,0 +1,40 @@
|
||||
import argparse
|
||||
import glob
|
||||
import os
|
||||
|
||||
|
||||
def main(args):
|
||||
txt_file = open(args.meta_info, 'w')
|
||||
for folder, root in zip(args.input, args.root):
|
||||
img_paths = sorted(glob.glob(os.path.join(folder, '*')))
|
||||
for img_path in img_paths:
|
||||
img_name = os.path.relpath(img_path, root)
|
||||
print(img_name)
|
||||
txt_file.write(f'{img_name}\n')
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
"""Generate meta info (txt file) for only Ground-Truth images.
|
||||
|
||||
It can also generate meta info from several folders into one txt file.
|
||||
"""
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument(
|
||||
'--input',
|
||||
nargs='+',
|
||||
default=['datasets/DF2K/DF2K_HR', 'datasets/DF2K/DF2K_multiscale'],
|
||||
help='Input folder, can be a list')
|
||||
parser.add_argument(
|
||||
'--root',
|
||||
nargs='+',
|
||||
default=['datasets/DF2K', 'datasets/DF2K'],
|
||||
help='Folder root, should have the length as input folders')
|
||||
parser.add_argument(
|
||||
'--meta_info',
|
||||
type=str,
|
||||
default='datasets/DF2K/meta_info/meta_info_DF2Kmultiscale.txt',
|
||||
help='txt path for meta info')
|
||||
args = parser.parse_args()
|
||||
assert len(args.input) == len(args.root), ('Input folder and folder root should have the same length, but got '
|
||||
f'{len(args.input)} and {len(args.root)}.')
|
||||
main(args)
|
||||
46
scripts/generate_multiscale_DF2K.py
Normal file
46
scripts/generate_multiscale_DF2K.py
Normal file
@@ -0,0 +1,46 @@
|
||||
import argparse
|
||||
import glob
|
||||
import os
|
||||
from PIL import Image
|
||||
|
||||
|
||||
def main(args):
|
||||
|
||||
# For DF2K, we consider the following three scales,
|
||||
# and the smallest image whose shortest edge is 400
|
||||
scale_list = [0.75, 0.5, 1 / 3]
|
||||
shortest_edge = 400
|
||||
|
||||
path_list = sorted(glob.glob(os.path.join(args.input, '*')))
|
||||
for path in path_list:
|
||||
print(path)
|
||||
basename = os.path.splitext(os.path.basename(path))[0]
|
||||
|
||||
img = Image.open(path)
|
||||
width, height = img.size
|
||||
for idx, scale in enumerate(scale_list):
|
||||
print(f'\t{scale:.2f}')
|
||||
rlt = img.resize((int(width * scale), int(height * scale)), resample=Image.LANCZOS)
|
||||
rlt.save(os.path.join(args.output, f'{basename}T{idx}.png'))
|
||||
|
||||
# save the smallest image which the shortest edge is 400
|
||||
if width < height:
|
||||
ratio = height / width
|
||||
width = shortest_edge
|
||||
height = int(width * ratio)
|
||||
else:
|
||||
ratio = width / height
|
||||
height = shortest_edge
|
||||
width = int(height * ratio)
|
||||
rlt = img.resize((int(width), int(height)), resample=Image.LANCZOS)
|
||||
rlt.save(os.path.join(args.output, f'{basename}T{idx+1}.png'))
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument('--input', type=str, default='datasets/DF2K/DF2K_HR', help='Input folder')
|
||||
parser.add_argument('--output', type=str, default='datasets/DF2K/DF2K_multiscale', help='Output folder')
|
||||
args = parser.parse_args()
|
||||
|
||||
os.makedirs(args.output, exist_ok=True)
|
||||
main(args)
|
||||
@@ -16,7 +16,7 @@ split_before_expression_after_opening_paren = true
|
||||
line_length = 120
|
||||
multi_line_output = 0
|
||||
known_standard_library = pkg_resources,setuptools
|
||||
known_first_party = basicsr # modify it!
|
||||
known_third_party = basicsr,cv2,numpy,torch
|
||||
known_first_party = realesrgan
|
||||
known_third_party = PIL,basicsr,cv2,numpy,torch
|
||||
no_lines_before = STDLIB,LOCALFOLDER
|
||||
default_section = THIRDPARTY
|
||||
|
||||
113
setup.py
Normal file
113
setup.py
Normal file
@@ -0,0 +1,113 @@
|
||||
#!/usr/bin/env python
|
||||
|
||||
from setuptools import find_packages, setup
|
||||
|
||||
import os
|
||||
import subprocess
|
||||
import time
|
||||
|
||||
version_file = 'realesrgan/version.py'
|
||||
|
||||
|
||||
def readme():
|
||||
with open('README.md', encoding='utf-8') as f:
|
||||
content = f.read()
|
||||
return content
|
||||
|
||||
|
||||
def get_git_hash():
|
||||
|
||||
def _minimal_ext_cmd(cmd):
|
||||
# construct minimal environment
|
||||
env = {}
|
||||
for k in ['SYSTEMROOT', 'PATH', 'HOME']:
|
||||
v = os.environ.get(k)
|
||||
if v is not None:
|
||||
env[k] = v
|
||||
# LANGUAGE is used on win32
|
||||
env['LANGUAGE'] = 'C'
|
||||
env['LANG'] = 'C'
|
||||
env['LC_ALL'] = 'C'
|
||||
out = subprocess.Popen(cmd, stdout=subprocess.PIPE, env=env).communicate()[0]
|
||||
return out
|
||||
|
||||
try:
|
||||
out = _minimal_ext_cmd(['git', 'rev-parse', 'HEAD'])
|
||||
sha = out.strip().decode('ascii')
|
||||
except OSError:
|
||||
sha = 'unknown'
|
||||
|
||||
return sha
|
||||
|
||||
|
||||
def get_hash():
|
||||
if os.path.exists('.git'):
|
||||
sha = get_git_hash()[:7]
|
||||
elif os.path.exists(version_file):
|
||||
try:
|
||||
from facexlib.version import __version__
|
||||
sha = __version__.split('+')[-1]
|
||||
except ImportError:
|
||||
raise ImportError('Unable to get git version')
|
||||
else:
|
||||
sha = 'unknown'
|
||||
|
||||
return sha
|
||||
|
||||
|
||||
def write_version_py():
|
||||
content = """# GENERATED VERSION FILE
|
||||
# TIME: {}
|
||||
__version__ = '{}'
|
||||
__gitsha__ = '{}'
|
||||
version_info = ({})
|
||||
"""
|
||||
sha = get_hash()
|
||||
with open('VERSION', 'r') as f:
|
||||
SHORT_VERSION = f.read().strip()
|
||||
VERSION_INFO = ', '.join([x if x.isdigit() else f'"{x}"' for x in SHORT_VERSION.split('.')])
|
||||
|
||||
version_file_str = content.format(time.asctime(), SHORT_VERSION, sha, VERSION_INFO)
|
||||
with open(version_file, 'w') as f:
|
||||
f.write(version_file_str)
|
||||
|
||||
|
||||
def get_version():
|
||||
with open(version_file, 'r') as f:
|
||||
exec(compile(f.read(), version_file, 'exec'))
|
||||
return locals()['__version__']
|
||||
|
||||
|
||||
def get_requirements(filename='requirements.txt'):
|
||||
here = os.path.dirname(os.path.realpath(__file__))
|
||||
with open(os.path.join(here, filename), 'r') as f:
|
||||
requires = [line.replace('\n', '') for line in f.readlines()]
|
||||
return requires
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
write_version_py()
|
||||
setup(
|
||||
name='realesrgan',
|
||||
version=get_version(),
|
||||
description='Real-ESRGAN aims at developing Practical Algorithms for General Image Restoration',
|
||||
long_description=readme(),
|
||||
long_description_content_type='text/markdown',
|
||||
author='Xintao Wang',
|
||||
author_email='xintao.wang@outlook.com',
|
||||
keywords='computer vision, pytorch, image restoration, super-resolution, esrgan, real-esrgan',
|
||||
url='https://github.com/xinntao/Real-ESRGAN',
|
||||
include_package_data=True,
|
||||
packages=find_packages(exclude=('options', 'datasets', 'experiments', 'results', 'tb_logger', 'wandb')),
|
||||
classifiers=[
|
||||
'Development Status :: 4 - Beta',
|
||||
'License :: OSI Approved :: Apache Software License',
|
||||
'Operating System :: OS Independent',
|
||||
'Programming Language :: Python :: 3',
|
||||
'Programming Language :: Python :: 3.7',
|
||||
'Programming Language :: Python :: 3.8',
|
||||
],
|
||||
license='BSD-3-Clause License',
|
||||
setup_requires=['cython', 'numpy'],
|
||||
install_requires=get_requirements(),
|
||||
zip_safe=False)
|
||||
Reference in New Issue
Block a user