(Fork of BasicSR joeyballentine/traiNNer-redux) Open Source Image and Video Restoration Toolbox for Super-resolution, Denoise, Deblurring, etc. Currently, it includes EDSR, RCAN, SRResNet, SRGAN, ESRGAN, EDVR, BasicVSR, SwinIR, ECBSR, OmniSR, HAT, GRL, A-ESRGAN, DAT, WaveMixSR, StarSRGAN, DLGSANet, DITN etc. Also support StyleGAN2, DFDNet.
NEW ADD ARCH SUPPORT
- [ESRGAN 8X]
- Omnisr
- The arch implementation of Omnisr is from Omnisr. The LICENSE of Omnisr is [Apache License 2.0]. The LICENSE is included as LICENSE_Omnisr.
- HAT
- The arch implementation of HAT is from HAT. The LICENSE of HAT is [MIT License]. The LICENSE is included as LICENSE_HAT.
- GRL
- The arch implementation of GRL is from GRL. The LICENSE of GRL is [MIT License].
- ESWT # have bug
- The arch implementation of ESWT is from ESWT. The LICENSE of ESWT is [MIT License].
- SRFormer
- The arch implementation of SRFormer is from SRFormer. The LICENSE of SRFormer is [Apache License 2.0].
- A-ESRGAN
- The arch implementation of A-ESRGAN is from A-ESRGAN. The LICENSE of A-ESRGAN is [BSD 3-Clause "New" or "Revised" License].
- DAT
- The arch implementation of DAT is from DAT. The LICENSE of DAT is [Apache License 2.0].
- WaveMixSR
- The arch implementation of WaveMixSR is from WaveMixSR. The LICENSE of WaveMixSR is [MIT License].
- StarSRGAN
- The arch implementation of StarSRGAN is from StarSRGAN. The LICENSE of StarSRGAN is [Apache License 2.0].
- DLGSANet
- The arch implementation of DLGSANet is from DLGSANet. The LICENSE of DLGSANet is [Apache License 2.0].
- DITN
- The arch implementation of DITN is from DITN.
NEW FEATURE SUPPORT
- ContextualLoss weight
- amp support
try_autoamp_g: True # enable amp Automatic mixed precision for network_g. if loss inf or nan or error just set to False try_autoamp_d: True # enable amp Automatic mixed precision for network_d. if loss inf or nan or error just set to False
TIPS
- PairedImageDataset set high_order_degradation : False
- RealESRGANDataset set high_order_degradation : True
- model_type can use GeneralGANModel or GeneralNetModel for ,ost of archs
- To use Automatic mixed precision, edit the yml
- e.g. if you want use omnisr-gan with amp, just use model_type: GeneralGANModel and edit the .yml 's following part
name: train_OmniSRGANModel_SRx4_scratch_P48W8_DIV2K_500k_B4G8
try_autoamp_g: True # enable amp Automatic mixed precision for network_g. if loss inf or nan or error just set to False
try_autoamp_d: True # enable amp Automatic mixed precision for network_d. if loss inf or nan or error just set to False
network_g:
type: OmniSRNet
num_in_ch: 3
num_out_ch: 3
num_feat: 64
upsampling: 4
res_num: 5
block_num: 1
bias: True
block_script_name: OSA
block_class_name: OSA_Block
window_size: 8
pe: True
ffn_bias: True
- If you want to use aesrgan's network_d for other network_d. you should edit the .yml network_d to multiscale type
- e.g. if you want use omnisr as network_g and multiscale as network_d, just edit the .yml 's following part
# network structures
network_g:
type: OmniSRNet
num_in_ch: 3
num_out_ch: 3
num_feat: 64
upsampling: 4
res_num: 5
block_num: 1
bias: True
block_script_name: OSA
block_class_name: OSA_Block
window_size: 8
pe: True
ffn_bias: True
network_d:
type: multiscale
num_in_ch: 3
num_feat: 64
num_D: 2
for easy use here are examples for network_g
# network structures
network_g:
# DAT-S, need to set batch size >1
type: DAT
upscale: 4
in_chans: 3
img_size: 64
img_range: 1.
split_size: [8,16]
depth: [6,6,6,6,6,6]
embed_dim: 180
num_heads: [6,6,6,6,6,6]
expansion_factor: 2
resi_connection: '1conv'
type: WaveMixSR
scale: 4
depth: 4
mult: 1
ff_channel: 144
final_dim: 144
dropout: 0.3
type: StarSRNet
num_in_ch: 3
num_out_ch: 3
scale: 4
num_feat: 64
num_block: 23
num_grow_ch: 32
drop_out: False
type: DLGSANet
upscale: 4
in_chans: 3
dim: 90
groups: 3
blocks: 2
buildblock_type: 'sparseedge'
window_size: 7
idynamic_num_heads: 6
idynamic_ffn_type: 'GDFN'
idynamic_ffn_expansion_factor: 2.
idynamic: true
restormer_num_heads: 6
restormer_ffn_type: 'GDFN'
restormer_ffn_expansion_factor: 2.
tlc_flag: true
tlc_kernel: 48 # using tlc during validation
activation: 'relu'
body_norm: false
img_range: 1.
upsampler: 'pixelshuffledirect'
type: DITN # need torch >2.0
inp_channels: 3
dim: 60
ITL_blocks: 4
SAL_blocks: 4
UFONE_blocks: 1
ffn_expansion_factor: 2
bias: False
LayerNorm_type: 'WithBias'
patch_size: 8
upscale: 4