# -*- coding: utf-8 -*-
import os
opt = {}
##################################################### Frequently Changed Setting ###########################################################
opt['description'] = "4x_GRL_paper" # Description to add to the log
opt['architecture'] = "GRL" # "ESRNET" || "ESRGAN" || "GRL" || "GRLGAN" (GRL only support 4x)
# Essential Setting
opt['scale'] = 4 # In default, this is 4x
opt["full_patch_source"] = "../datasets_anime/APISR_dataset" # The HR image without cropping
opt["degrade_hr_dataset_path"] = "datasets/train_hr" # The cropped GT images
opt["train_hr_dataset_path"] = "datasets/train_hr_enhanced" # The cropped Pseudo-GT path (after hand-drawn line enhancement)
################################################################################################################################
# GPU setting
opt['CUDA_VISIBLE_DEVICES'] = '0' # '0' / '1' based on different GPU you have.
os.environ['CUDA_VISIBLE_DEVICES'] = opt['CUDA_VISIBLE_DEVICES']
##################################################### Setting for General Training #############################################
# Dataset Setting
opt["lr_dataset_path"] = "datasets/train_lr" # Where you temporally store the LR synthetic images
opt['hr_size'] = 256
# Loss function
opt['pixel_loss'] = "L1" # Usually it is "L1"
# Adam optimizer setting
opt["adam_beta1"] = 0.9
opt["adam_beta2"] = 0.99
opt['decay_gamma'] = 0.5 # Decay the learning rate per decay_iteration
# Miscellaneous Setting
opt['degradate_generation_freq'] = 1 # How frequent we degradate HR to LR (1: means Real-Time Degrade) [No need to change this]
opt['train_dataloader_workers'] = 5 # Number of workers for DataLoader
opt['checkpoints_freq'] = 50 # frequency to store checkpoints in the folder (unit: epoch)
#################################################################################################################################
# Add setting for different architecture (Please go through the model architecture you want!)
if opt['architecture'] == "ESRNET":
# Setting for ESRNET Training
opt['ESR_blocks_num'] = 6 # How many RRDB blocks you need
opt['train_iterations'] = 500000 # Training Iterations (500K for large resolution large dataset overlap training)
opt['train_batch_size'] = 32 #
# Learning Rate
opt["start_learning_rate"] = 0.0002 # Training Epoch, use the as Real-ESRGAN: 0.0001 - 0.0002 is ok, based on your need
opt['decay_iteration'] = 100000 # Decay iteration
opt['double_milestones'] = [] # Iteration based time you double your learning rate
elif opt['architecture'] == "ESRGAN":
# Setting for ESRGAN Training
opt['ESR_blocks_num'] = 6 # How many RRDB blocks you need
opt['train_iterations'] = 200000 # Training Iterations
opt['train_batch_size'] = 32 #
# Learning Rate
opt["start_learning_rate"] = 0.0001 # Training Epoch, use the as Real-ESRGAN: 0.0001 - 0.0002 is ok, based on your need
opt['decay_iteration'] = 100000 # Fixed decay gap
opt['double_milestones'] = [] # Just put this empty
# Perceptual loss
opt["danbooru_perceptual_loss_weight"] = 0.5 # ResNet50 Danbooru Perceptual loss weight scale
opt["vgg_perceptual_loss_weight"] = 0.5 # VGG PhotoRealistic Perceptual loss weight scale
opt['train_perceptual_vgg_type'] = 'vgg19' # VGG16/19 (Just use 19 by default)
opt['train_perceptual_layer_weights'] = {'conv1_2': 0.1, 'conv2_2': 0.1, 'conv3_4': 1, 'conv4_4': 1, 'conv5_4': 1} # Middle-Layer weight for VGG
opt['Danbooru_layer_weights'] = {"0": 0.1, "4_2_conv3": 20, "5_3_conv3": 25, "6_5_conv3": 1, "7_2_conv3": 1} # Middle-Layer weight for ResNet
# GAN loss
opt["discriminator_type"] = "PatchDiscriminator" # "PatchDiscriminator" || "UNetDiscriminator"
opt["gan_loss_weight"] = 0.2 #
elif opt['architecture'] == "CUNET":
# Setting for CUNET Training
opt['train_iterations'] = 500000 # Training Iterations (700K for large resolution large dataset overlap training)
opt['train_batch_size'] = 16
opt["start_learning_rate"] = 0.0002 # Training Epoch, use the as Real-ESRGAN: 0.0001 - 0.0002 is ok, based on your need
opt['decay_iteration'] = 100000 # Decay iteration
opt['double_milestones'] = [] # Iteration based time you double your learning rate
elif opt['architecture'] == "CUGAN":
# Setting for ESRGAN Training
opt['ESR_blocks_num'] = 6 # How many RRDB blocks you need
opt['train_iterations'] = 200000 # Training Iterations
opt['train_batch_size'] = 16
opt["start_learning_rate"] = 0.0001 # Training Epoch, use the as Real-ESRGAN: 0.0001 - 0.0002 is ok, based on your need
opt["perceptual_loss_weight"] = 1.0
opt['train_perceptual_vgg_type'] = 'vgg19'
opt['train_perceptual_layer_weights'] = {'conv1_2': 0.1, 'conv2_2': 0.1, 'conv3_4': 1, 'conv4_4': 1, 'conv5_4': 1}
opt['Danbooru_layer_weights'] = {"0": 0.1, "4_2_conv3": 20, "5_3_conv3": 25, "6_5_conv3": 1, "7_2_conv3": 1} # Middle-Layer weight for ResNet
opt["gan_loss_weight"] = 0.2 # This one is very important, Don't neglect it. Based on the paper, it should be 0.1 scale
opt['decay_iteration'] = 100000 # Decay iteration
opt['double_milestones'] = [] # Iteration based time you double your learning rate
elif opt['architecture'] == "GRL": # L1 loss training version
# Setting for GRL Training
opt['model_size'] = "tiny2" # "tiny2" in default
opt['train_iterations'] = 300000 # Training Iterations
opt['train_batch_size'] = 32 # 4x: 32 (256x256); 2x: 4?
# Learning Rate
opt["start_learning_rate"] = 0.0002 # Training Epoch, use the as Real-ESRGAN: 0.0001 - 0.0002 is ok, based on your need
opt['decay_iteration'] = 100000 # Decay iteration
opt['double_milestones'] = [] # Iteration based time you double your learning rate (Just ignore this one)
elif opt['architecture'] == "GRLGAN": # L1 + Preceptual + Discriminator Loss version
# Setting for GRL Training
opt['model_size'] = "tiny2" # "small" || "tiny" || "tiny2" (Use tiny2 by default, No need to change)
# Setting for GRL-GAN Traning
opt['train_iterations'] = 300000 # Training Iterations
opt['train_batch_size'] = 32 # 4x: 32 batch size (for 256x256); 2x: 4
# Learning Rate
opt["start_learning_rate"] = 0.0001 # Training Epoch, use the as Real-ESRGAN: 0.0001 - 0.0002 is ok, based on your need
opt['decay_iteration'] = 100000 # Fixed decay gap
opt['double_milestones'] = [] # Just put this empty
# Perceptual loss
opt["danbooru_perceptual_loss_weight"] = 0.5 # ResNet50 Danbooru Perceptual loss weight scale
opt["vgg_perceptual_loss_weight"] = 0.5 # VGG PhotoRealistic Perceptual loss weight scale
opt['train_perceptual_vgg_type'] = 'vgg19' # VGG16/19 (Just use 19 by default)
opt['train_perceptual_layer_weights'] = {'conv1_2': 0.1, 'conv2_2': 0.1, 'conv3_4': 1, 'conv4_4': 1, 'conv5_4': 1} # Middle-Layer weight for VGG
opt['Danbooru_layer_weights'] = {"0": 0.1, "4_2_conv3": 20, "5_3_conv3": 25, "6_5_conv3": 1, "7_2_conv3": 1} # Middle-Layer weight for ResNet
# GAN loss
opt["discriminator_type"] = "PatchDiscriminator" # "PatchDiscriminator" || "UNetDiscriminator"
opt["gan_loss_weight"] = 0.2 #
else:
raise NotImplementedError("Please check you architecture option setting!")
# Basic setting for degradation
opt["degradation_batch_size"] = 128 # Degradation batch size
opt["augment_prob"] = 0.5 # Probability of augmenting (Flip, Rotate) the HR and LR dataset in dataset loading part
if opt['architecture'] in ["ESRNET", "ESRGAN", "GRL", "GRLGAN", "CUNET", "CUGAN"]:
# Parallel Process
opt['parallel_num'] = 8 # Multi-Processing num; Recommend 6
# Blur kernel1
opt['kernel_range'] = [3, 11]
opt['kernel_list'] = ['iso', 'aniso', 'generalized_iso', 'generalized_aniso', 'plateau_iso', 'plateau_aniso']
opt['kernel_prob'] = [0.45, 0.25, 0.12, 0.03, 0.12, 0.03]
opt['sinc_prob'] = 0.1
opt['blur_sigma'] = [0.2, 3]
opt['betag_range'] = [0.5, 4]
opt['betap_range'] = [1, 2]
# Blur kernel2
opt['kernel_list2'] = ['iso', 'aniso', 'generalized_iso', 'generalized_aniso', 'plateau_iso', 'plateau_aniso']
opt['kernel_prob2'] = [0.45, 0.25, 0.12, 0.03, 0.12, 0.03]
opt['sinc_prob2'] = 0.1
opt['blur_sigma2'] = [0.2, 1.5]
opt['betag_range2'] = [0.5, 4]
opt['betap_range2'] = [1, 2]
# The first degradation process
opt['resize_prob'] = [0.2, 0.7, 0.1]
opt['resize_range'] = [0.1, 1.2] # Was [0.15, 1.5] in Real-ESRGAN
opt['gaussian_noise_prob'] = 0.5
opt['noise_range'] = [1, 30]
opt['poisson_scale_range'] = [0.05, 3]
opt['gray_noise_prob'] = 0.4
opt['jpeg_range'] = [30, 95]
# The second degradation process
opt['second_blur_prob'] = 0.8
opt['resize_prob2'] = [0.2, 0.7, 0.1] # [up, down, keep] Resize Probability
opt['resize_range2'] = [0.15, 1.2]
opt['gaussian_noise_prob2'] = 0.5
opt['noise_range2'] = [1, 25]
opt['poisson_scale_range2'] = [0.05, 2.5]
opt['gray_noise_prob2'] = 0.4
# Other common settings
opt['resize_options'] = ['area', 'bilinear', 'bicubic'] # Should be supported by F.interpolate
# First image compression
opt['compression_codec1'] = ["jpeg", "webp", "heif", "avif"] # Compression codec: heif is the intra frame version of HEVC (H.265) and avif is the intra frame version of AV1
opt['compression_codec_prob1'] = [0.4, 0.6, 0.0, 0.0]
# Specific Setting
opt["jpeg_quality_range1"] = [20, 95]
opt["webp_quality_range1"] = [20, 95]
opt["webp_encode_speed1"] = [0, 6]
opt["heif_quality_range1"] = [30, 100]
opt["heif_encode_speed1"] = [0, 6] # Useless now
opt["avif_quality_range1"] = [30, 100]
opt["avif_encode_speed1"] = [0, 6] # Useless now
######################################## Setting for Degradation with Intra-Prediction ########################################
opt['compression_codec2'] = ["jpeg", "webp", "avif", "mpeg2", "mpeg4", "h264", "h265"] # Compression codec: similar to VCISR but more intense degradation settings
opt['compression_codec_prob2'] = [0.06, 0.1, 0.1, 0.12, 0.12, 0.3, 0.2]
# Image compression setting
opt["jpeg_quality_range2"] = [20, 95]
opt["webp_quality_range2"] = [20, 95]
opt["webp_encode_speed2"] = [0, 6]
opt["avif_quality_range2"] = [20, 95]
opt["avif_encode_speed2"] = [0, 6] # Useless now
# Video compression I-Frame setting
opt['h264_crf_range2'] = [23, 38]
opt['h264_preset_mode2'] = ["slow", "medium", "fast", "faster", "superfast"]
opt['h264_preset_prob2'] = [0.05, 0.35, 0.3, 0.2, 0.1]
opt['h265_crf_range2'] = [28, 42]
opt['h265_preset_mode2'] = ["slow", "medium", "fast", "faster", "superfast"]
opt['h265_preset_prob2'] = [0.05, 0.35, 0.3, 0.2, 0.1]
opt['mpeg2_quality2'] = [8, 31] # linear scale 2-31 (the lower the higher quality)
opt['mpeg2_preset_mode2'] = ["slow", "medium", "fast", "faster", "superfast"]
opt['mpeg2_preset_prob2'] = [0.05, 0.35, 0.3, 0.2, 0.1]
opt['mpeg4_quality2'] = [8, 31] # should be the same as mpeg2_quality2
opt['mpeg4_preset_mode2'] = ["slow", "medium", "fast", "faster", "superfast"]
opt['mpeg4_preset_prob2'] = [0.05, 0.35, 0.3, 0.2, 0.1]