Awesome
<div align="center"> <a href=https://github.com/muslll/neosr><b>neosr</b></a> is a framework for training real-world single-image super-resolution networks. <a href=https://github.com/muslll/neosr/wiki><b>wiki</b></a>. <br><br> <a href="https://discord.gg/NN2HGtJ3d6"><img src="https://github.com/muslll/neosr/assets/132400428/4bd54b1d-4639-4940-b9c7-b3f212aea5c8?sanitize=true" width="120"</img></a><br> Join our <a href="https://discord.gg/NN2HGtJ3d6">Discord</a> </div> <details> <summary><b>news</b></summary> 05.09.2024 - Release Real-PLKSR network. <a href=https://github.com/muslll/neosr/wiki/Arch%E2%80%90specific-options#realplksr-realplksr_s>wiki</a><br> 05.08.2024 - Release <b>Nomos-v2</b> dataset. For more details, see <a href=#datasets>datasets</a><br> </details>🤝 support me
💻 installation
Requires Python 3.11 and CUDA =>11.8. Install latest Pytorch (=>2.1) and TorchVision (required).
Clone the repository:
git clone https://github.com/muslll/neosr
cd neosr
Then install other dependencies via pip
:
pip install -e .
Alternatively, use poetry
(recommended on linux):
poetry install
poetry add torch@latest torchvision@latest
Note: You must use poetry shell
to enter the env after installation.
(optional) If you want to convert your models (convert.py), you need the following dependencies:
pip install onnx onnxruntime-gpu onnxconverter-common onnxsim
You can also install using poetry (recommended on linux):
poetry add onnx onnxruntime-gpu onnxconverter-common onnxsim
Please read the wiki tutorial for converting your models.
quick start
Start training by running:
python train.py -opt options.yml
Where options.yml
is a configuration file. Templates can be found in options.
[!TIP] Please read the wiki Configuration Walkthrough for an explanation of each option.
features
Supported archs:
arch | option |
---|---|
Real-ESRGAN | esrgan |
SRVGGNetCompact | compact |
SwinIR | swinir_small , swinir_medium |
HAT | hat_s , hat_m , hat_l |
OmniSR | omnisr |
SRFormer | srformer_light , srformer_medium |
DAT | dat_small , dat_medium , dat_2 |
DITN | ditn |
DCTLSA | dctlsa |
SPAN | span |
Real-CUGAN | cugan |
CRAFT | craft |
SAFMN | safmn , safmn_l |
RGT | rgt , rgt_s |
ATD | atd , atd_light |
PLKSR | plksr , plksr_tiny |
RealPLKSR | realplksr , realplksr_s |
DRCT | drct , drct_l , drct_s |
EFEN | efen |
[!NOTE] For all arch-specific parameters, read the wiki.
Under Testing
arch | option |
---|
Supported Discriminators:
net | option |
---|---|
U-Net w/ SN | unet |
A2-FPN w/ SN | a2fpn |
PatchGAN w/ SN | patchgan |
Supported Optimizers:
optimizer | option |
---|---|
Adam | Adam or adam |
AdamW | AdamW or adamw |
NAdam | NAdam or nadam |
Lion | Lion or lion |
LAMB | Lamb or lamb |
Adan | Adan or adan |
Supported losses:
loss | option |
---|---|
L1 Loss | L1Loss , l1 |
L2 Loss | MSELoss , l2 |
Huber Loss | HuberLoss , huber |
CHC (Clipped Huber with Cosine Similarity Loss) | chc , chc_l2 |
Perceptual Loss | perceptual_opt , PerceptualLoss |
GAN | gan_opt , GANLoss , MultiScaleGANLoss |
YCbCr Color Loss (bt601) | color_opt , colorloss |
Luma Loss (CIE L*) | luma_opt lumaloss |
MS-SSIM | mssim_opt mssim |
LDL Loss | ldl_opt |
Focal Frequency | ff_opt , focalfrequencyloss |
DISTS | dists_opt , dists |
Wavelet Guided | wavelet_guided |
Gradient-Weighted | gw_opt , gw_loss |
Supported Augmentations:
augmentation | option |
---|---|
Rotation | use_rot |
Flip | use_hflip |
MixUp | mixup |
CutMix | cutmix |
ResizeMix | resizemix |
CutBlur | cutblur |
Supported models:
model | description | option |
---|---|---|
Default | Base model, supports both Generator and Discriminator | default |
OTF | Builds on top of default , adding Real-ESRGAN on-the-fly degradations | otf |
Supported dataset loaders:
loader | option |
---|---|
Paired datasets | paired |
Single datasets (for inference, no GT required) | single |
Real-ESRGAN on-the-fly degradation | otf |
📸 datasets
As part of neosr, I have released a dataset series called Nomos. The purpose of these datasets is to distill only the best images from the academic and community datasets. A total of 14 datasets were manually reviewed and processed, including: Adobe-MIT-5k, RAISE, LSDIR, LIU4k-v2, KONIQ-10k, Nikon LL RAW, DIV8k, FFHQ, Flickr2k, ModernAnimation1080_v2, Rawsamples, SignatureEdits, Hasselblad raw samples and Unsplash.
Nomos-v2
(recommended): contains 6000 images, multipurpose. Data distribution:
pie
title Nomos-v2 distribution
"Animal / fur" : 439
"Interiors" : 280
"Exteriors / misc" : 696
"Architecture / geometric" : 1470
"Drawing / painting / anime" : 1076
"Humans" : 598
"Mountain / Rocks" : 317
"Text" : 102
"Textures" : 439
"Vegetation" : 574
nomos_uni
(recommended for lightweight networks): contains 2989 images, multipurpose. Meant to be used on lightweight networks (<800k parameters).hfa2k
: contains 2568 anime images.
dataset download | sha256 |
---|---|
nomosv2 (3GB) | sha256 |
nomosv2.lmdb (3GB) | sha256 |
nomosv2_lq_4x (187MB) | sha256 |
nomosv2_lq_4x.lmdb (187MB) | sha256 |
nomos_uni (1.3GB) | sha256 |
nomos_uni.lmdb (1.3GB) | sha256 |
nomos_uni_lq_4x | sha256 |
nomos_uni_lq_4x.lmdb | sha256 |
hfa2k | sha256 |
community datasets
Datasets made by the upscaling community. More info can be found in the Enhance Everything discord
FaceUp
: Curated version of FFHQSSDIR
: Curated version of LSDIR.kim's 8k Dataset V2
: Video Game Dataset
dataset | download |
---|---|
@Kim2091 8k Dataset V2 | GDrive (33.5GB) |
@Phhofm FaceUp | GDrive (4GB) |
@Phhofm SSDIR | Gdrive (4.5GB) |
resources
- OpenModelDB
- chaiNNer
- Training Guide from @Sirosky
- Training Info from @Kim
📄 license and acknowledgements
Released under the Apache license. All licenses listed on license/readme. This code was originally based on BasicSR.
Thanks to victorca25/traiNNer, styler00dollar/Colab-traiNNer and timm for providing helpful insights into some problems.
Thanks to contributors @Phhofm, @Sirosky, @Kim2091, @terrainer, @Corpsecreate and @umzi2 for helping with tests and bug reporting.