updates

VERSION

@@ -1 +1 @@
-0.2.1
+0.2.3.0

inference_realesrgan.py

@@ -2,25 +2,32 @@ import argparse
 import cv2
 import glob
 import os
+from basicsr.archs.rrdbnet_arch import RRDBNet
 
 from realesrgan import RealESRGANer
+from realesrgan.archs.srvgg_arch import SRVGGNetCompact
 
 
 def main():
+    """Inference demo for Real-ESRGAN.
+    """
     parser = argparse.ArgumentParser()
-    parser.add_argument('--input', type=str, default='inputs', help='Input image or folder')
+    parser.add_argument('-i', '--input', type=str, default='inputs', help='Input image or folder')
     parser.add_argument(
-        '--model_path',
+        '-n',
+        '--model_name',
         type=str,
-        default='RealESRGAN_x4plus.pth',
-        help='Path to the pre-trained model')
-    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')
+        default='RealESRGAN_x4plus',
+        help=('Model names: RealESRGAN_x4plus | RealESRNet_x4plus | RealESRGAN_x4plus_anime_6B | RealESRGAN_x2plus'
+              'RealESRGANv2-anime-xsx2 | RealESRGANv2-animevideo-xsx2-nousm | RealESRGANv2-animevideo-xsx2'
+              'RealESRGANv2-anime-xsx4 | RealESRGANv2-animevideo-xsx4-nousm | RealESRGANv2-animevideo-xsx4'))
+    parser.add_argument('-o', '--output', type=str, default='results', help='Output folder')
+    parser.add_argument('-s', '--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('-t', '--tile', type=int, default=0, 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',
@@ -34,14 +41,55 @@ def main():
         help='Image extension. Options: auto | jpg | png, auto means using the same extension as inputs')
     args = parser.parse_args()
 
+    # determine models according to model names
+    args.model_name = args.model_name.split('.')[0]
+    if args.model_name in ['RealESRGAN_x4plus', 'RealESRNet_x4plus']:  # x4 RRDBNet model
+        model = RRDBNet(num_in_ch=3, num_out_ch=3, num_feat=64, num_block=23, num_grow_ch=32, scale=4)
+        netscale = 4
+    elif args.model_name in ['RealESRGAN_x4plus_anime_6B']:  # x4 RRDBNet model with 6 blocks
+        model = RRDBNet(num_in_ch=3, num_out_ch=3, num_feat=64, num_block=6, num_grow_ch=32, scale=4)
+        netscale = 4
+    elif args.model_name in ['RealESRGAN_x2plus']:  # x2 RRDBNet model
+        model = RRDBNet(num_in_ch=3, num_out_ch=3, num_feat=64, num_block=23, num_grow_ch=32, scale=2)
+        netscale = 2
+    elif args.model_name in [
+            'RealESRGANv2-anime-xsx2', 'RealESRGANv2-animevideo-xsx2-nousm', 'RealESRGANv2-animevideo-xsx2'
+    ]:  # x2 VGG-style model (XS size)
+        model = SRVGGNetCompact(num_in_ch=3, num_out_ch=3, num_feat=64, num_conv=16, upscale=2, act_type='prelu')
+        netscale = 2
+    elif args.model_name in [
+            'RealESRGANv2-anime-xsx4', 'RealESRGANv2-animevideo-xsx4-nousm', 'RealESRGANv2-animevideo-xsx4'
+    ]:  # x4 VGG-style model (XS size)
+        model = SRVGGNetCompact(num_in_ch=3, num_out_ch=3, num_feat=64, num_conv=16, upscale=4, act_type='prelu')
+        netscale = 4
+
+    # determine model paths
+    model_path = os.path.join('experiments/pretrained_models', args.model_name + '.pth')
+    if not os.path.isfile(model_path):
+        model_path = os.path.join('realesrgan/weights', args.model_name + '.pth')
+    if not os.path.isfile(model_path):
+        raise ValueError(f'Model {args.model_name} does not exist.')
+
+    # restorer
     upsampler = RealESRGANer(
-        scale=args.netscale,
-        model_path=args.model_path,
+        scale=netscale,
+        model_path=model_path,
+        model=model,
         tile=args.tile,
         tile_pad=args.tile_pad,
         pre_pad=args.pre_pad,
         half=args.half)
+
+    if args.face_enhance:  # Use GFPGAN for face enhancement
+        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:
@@ -52,18 +100,19 @@ def main():
         print('Testing', idx, imgname)
 
         img = cv2.imread(path, cv2.IMREAD_UNCHANGED)
-        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.')
+        if len(img.shape) == 3 and img.shape[2] == 4:
+            img_mode = 'RGBA'
+        else:
+            img_mode = None
 
         try:
-            output, img_mode = upsampler.enhance(img, outscale=args.outscale)
-        except Exception as error:
+            if args.face_enhance:
+                _, _, output = face_enhancer.enhance(img, has_aligned=False, only_center_face=False, paste_back=True)
+            else:
+                output, _ = upsampler.enhance(img, outscale=args.outscale)
+        except RuntimeError as error:
             print('Error', error)
+            print('If you encounter CUDA out of memory, try to set --tile with a smaller number.')
         else:
             if args.ext == 'auto':
                 extension = extension[1:]

inference_realesrgan_video.py

@@ -0,0 +1,199 @@
+import argparse
+import glob
+import mimetypes
+import os
+import queue
+import shutil
+import torch
+from basicsr.archs.rrdbnet_arch import RRDBNet
+from basicsr.utils.logger import AvgTimer
+from tqdm import tqdm
+
+from realesrgan import IOConsumer, PrefetchReader, RealESRGANer
+from realesrgan.archs.srvgg_arch import SRVGGNetCompact
+
+
+def main():
+    """Inference demo for Real-ESRGAN.
+    It mainly for restoring anime videos.
+
+    """
+    parser = argparse.ArgumentParser()
+    parser.add_argument('-i', '--input', type=str, default='inputs', help='Input image or folder')
+    parser.add_argument(
+        '-n',
+        '--model_name',
+        type=str,
+        default='RealESRGAN_x4plus',
+        help=('Model names: RealESRGAN_x4plus | RealESRNet_x4plus | RealESRGAN_x4plus_anime_6B | RealESRGAN_x2plus'
+              'RealESRGANv2-anime-xsx2 | RealESRGANv2-animevideo-xsx2-nousm | RealESRGANv2-animevideo-xsx2'
+              'RealESRGANv2-anime-xsx4 | RealESRGANv2-animevideo-xsx4-nousm | RealESRGANv2-animevideo-xsx4'))
+    parser.add_argument('-o', '--output', type=str, default='results', help='Output folder')
+    parser.add_argument('-s', '--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 video')
+    parser.add_argument('-t', '--tile', type=int, default=0, 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('-v', '--video', action='store_true', help='Output a video using ffmpeg')
+    parser.add_argument('-a', '--audio', action='store_true', help='Keep audio')
+    parser.add_argument('--fps', type=float, default=None, help='FPS of the output video')
+    parser.add_argument('--consumer', type=int, default=4, help='Number of IO consumers')
+
+    parser.add_argument(
+        '--alpha_upsampler',
+        type=str,
+        default='realesrgan',
+        help='The upsampler for the alpha channels. Options: realesrgan | bicubic')
+    parser.add_argument(
+        '--ext',
+        type=str,
+        default='auto',
+        help='Image extension. Options: auto | jpg | png, auto means using the same extension as inputs')
+    args = parser.parse_args()
+
+    # ---------------------- determine models according to model names ---------------------- #
+    args.model_name = args.model_name.split('.')[0]
+    if args.model_name in ['RealESRGAN_x4plus', 'RealESRNet_x4plus']:  # x4 RRDBNet model
+        model = RRDBNet(num_in_ch=3, num_out_ch=3, num_feat=64, num_block=23, num_grow_ch=32, scale=4)
+        netscale = 4
+    elif args.model_name in ['RealESRGAN_x4plus_anime_6B']:  # x4 RRDBNet model with 6 blocks
+        model = RRDBNet(num_in_ch=3, num_out_ch=3, num_feat=64, num_block=6, num_grow_ch=32, scale=4)
+        netscale = 4
+    elif args.model_name in ['RealESRGAN_x2plus']:  # x2 RRDBNet model
+        model = RRDBNet(num_in_ch=3, num_out_ch=3, num_feat=64, num_block=23, num_grow_ch=32, scale=2)
+        netscale = 2
+    elif args.model_name in [
+            'RealESRGANv2-anime-xsx2', 'RealESRGANv2-animevideo-xsx2-nousm', 'RealESRGANv2-animevideo-xsx2'
+    ]:  # x2 VGG-style model (XS size)
+        model = SRVGGNetCompact(num_in_ch=3, num_out_ch=3, num_feat=64, num_conv=16, upscale=2, act_type='prelu')
+        netscale = 2
+    elif args.model_name in [
+            'RealESRGANv2-anime-xsx4', 'RealESRGANv2-animevideo-xsx4-nousm', 'RealESRGANv2-animevideo-xsx4'
+    ]:  # x4 VGG-style model (XS size)
+        model = SRVGGNetCompact(num_in_ch=3, num_out_ch=3, num_feat=64, num_conv=16, upscale=4, act_type='prelu')
+        netscale = 4
+
+    # ---------------------- determine model paths ---------------------- #
+    model_path = os.path.join('experiments/pretrained_models', args.model_name + '.pth')
+    if not os.path.isfile(model_path):
+        model_path = os.path.join('realesrgan/weights', args.model_name + '.pth')
+    if not os.path.isfile(model_path):
+        raise ValueError(f'Model {args.model_name} does not exist.')
+
+    # restorer
+    upsampler = RealESRGANer(
+        scale=netscale,
+        model_path=model_path,
+        model=model,
+        tile=args.tile,
+        tile_pad=args.tile_pad,
+        pre_pad=args.pre_pad,
+        half=args.half)
+
+    if args.face_enhance:  # Use GFPGAN for face enhancement
+        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)
+    # for saving restored frames
+    save_frame_folder = os.path.join(args.output, 'frames_tmpout')
+    os.makedirs(save_frame_folder, exist_ok=True)
+
+    if mimetypes.guess_type(args.input)[0].startswith('video'):  # is a video file
+        video_name = os.path.splitext(os.path.basename(args.input))[0]
+        frame_folder = os.path.join('tmp_frames', video_name)
+        os.makedirs(frame_folder, exist_ok=True)
+        # use ffmpeg to extract frames
+        os.system(f'ffmpeg -i {args.input} -qscale:v 1 -qmin 1 -qmax 1 -vsync 0  {frame_folder}/frame%08d.png')
+        # get image path list
+        paths = sorted(glob.glob(os.path.join(frame_folder, '*')))
+        if args.video:
+            if args.fps is None:
+                # get input video fps
+                import ffmpeg
+                probe = ffmpeg.probe(args.input)
+                video_streams = [stream for stream in probe['streams'] if stream['codec_type'] == 'video']
+                args.fps = eval(video_streams[0]['avg_frame_rate'])
+    elif mimetypes.guess_type(args.input)[0].startswith('image'):  # is an image file
+        paths = [args.input]
+        video_name = 'video'
+    else:
+        paths = sorted(glob.glob(os.path.join(args.input, '*')))
+        video_name = 'video'
+
+    timer = AvgTimer()
+    timer.start()
+    pbar = tqdm(total=len(paths), unit='frame', desc='inference')
+    # set up prefetch reader
+    reader = PrefetchReader(paths, num_prefetch_queue=4)
+    reader.start()
+
+    que = queue.Queue()
+    consumers = [IOConsumer(args, que, f'IO_{i}') for i in range(args.consumer)]
+    for consumer in consumers:
+        consumer.start()
+
+    for idx, (path, img) in enumerate(zip(paths, reader)):
+        imgname, extension = os.path.splitext(os.path.basename(path))
+        if len(img.shape) == 3 and img.shape[2] == 4:
+            img_mode = 'RGBA'
+        else:
+            img_mode = None
+
+        try:
+            if args.face_enhance:
+                _, _, output = face_enhancer.enhance(img, has_aligned=False, only_center_face=False, paste_back=True)
+            else:
+                output, _ = upsampler.enhance(img, outscale=args.outscale)
+        except RuntimeError as error:
+            print('Error', error)
+            print('If you encounter CUDA out of memory, try to set --tile with a smaller number.')
+
+        else:
+            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 = os.path.join(save_frame_folder, f'{imgname}_out.{extension}')
+
+            que.put({'output': output, 'save_path': save_path})
+
+        pbar.update(1)
+        torch.cuda.synchronize()
+        timer.record()
+        avg_fps = 1. / (timer.get_avg_time() + 1e-7)
+        pbar.set_description(f'idx {idx}, fps {avg_fps:.2f}')
+
+    for _ in range(args.consumer):
+        que.put('quit')
+    for consumer in consumers:
+        consumer.join()
+    pbar.close()
+
+    # merge frames to video
+    if args.video:
+        video_save_path = os.path.join(args.output, f'{video_name}_{args.suffix}.mp4')
+        if args.audio:
+            os.system(
+                f'ffmpeg -r {args.fps} -i {save_frame_folder}/frame%08d_out.{extension} -i {args.input}'
+                f' -map 0:v:0 -map 1:a:0 -c:a copy -c:v libx264 -r {args.fps} -pix_fmt yuv420p  {video_save_path}')
+        else:
+            os.system(f'ffmpeg -r {args.fps} -i {save_frame_folder}/frame%08d_out.{extension} '
+                      f'-c:v libx264 -r {args.fps} -pix_fmt yuv420p {video_save_path}')
+
+        # delete tmp file
+        shutil.rmtree(save_frame_folder)
+        if os.path.isdir(frame_folder):
+            shutil.rmtree(frame_folder)
+
+
+if __name__ == '__main__':
+    main()

options/finetune_realesrgan_x4plus.yml

@@ -0,0 +1,188 @@
+# general settings
+name: finetune_RealESRGANx4plus_400k
+model_type: RealESRGANModel
+scale: 4
+num_gpu: auto
+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
+#       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

options/finetune_realesrgan_x4plus_pairdata.yml

@@ -0,0 +1,150 @@
+# general settings
+name: finetune_RealESRGANx4plus_400k_pairdata
+model_type: RealESRGANModel
+scale: 4
+num_gpu: auto
+manual_seed: 0
+
+# USM the ground-truth
+l1_gt_usm: True
+percep_gt_usm: True
+gan_gt_usm: False
+
+high_order_degradation: False # do not use the high-order degradation generation process
+
+# dataset and data loader settings
+datasets:
+  train:
+    name: DIV2K
+    type: RealESRGANPairedDataset
+    dataroot_gt: datasets/DF2K
+    dataroot_lq: datasets/DF2K
+    meta_info: datasets/DF2K/meta_info/meta_info_DIV2K_sub_pair.txt
+    io_backend:
+      type: disk
+
+    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
+#       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

options/setup.cfg

@@ -0,0 +1,33 @@
+[flake8]
+ignore =
+    # line break before binary operator (W503)
+    W503,
+    # line break after binary operator (W504)
+    W504,
+max-line-length=120
+
+[yapf]
+based_on_style = pep8
+column_limit = 120
+blank_line_before_nested_class_or_def = true
+split_before_expression_after_opening_paren = true
+
+[isort]
+line_length = 120
+multi_line_output = 0
+known_standard_library = pkg_resources,setuptools
+known_first_party = realesrgan
+known_third_party = PIL,basicsr,cv2,numpy,pytest,torch,torchvision,tqdm,yaml
+no_lines_before = STDLIB,LOCALFOLDER
+default_section = THIRDPARTY
+
+[codespell]
+skip = .git,./docs/build
+count =
+quiet-level = 3
+
+[aliases]
+test=pytest
+
+[tool:pytest]
+addopts=tests/

options/train_realesrgan_x2plus.yml

@@ -0,0 +1,186 @@
+# general settings
+name: train_RealESRGANx2plus_400k_B12G4
+model_type: RealESRGANModel
+scale: 2
+num_gpu: auto  # auto: can infer from your visible devices automatically. official: 4 GPUs
+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
+#       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

options/train_realesrgan_x4plus.yml

@@ -1,8 +1,8 @@
 # general settings
-name: train_RealESRGANx4plus_400k_B12G4_fromRealESRNet
+name: train_RealESRGANx4plus_400k_B12G4
 model_type: RealESRGANModel
 scale: 4
-num_gpu: 4
+num_gpu: auto  # auto: can infer from your visible devices automatically. official: 4 GPUs
 manual_seed: 0
 
 # ----------------- options for synthesizing training data in RealESRGANModel ----------------- #
@@ -39,7 +39,7 @@ datasets:
     name: DF2K+OST
     type: RealESRGANDataset
     dataroot_gt: datasets/DF2K
-    meta_info: realesrgan/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
 
@@ -90,7 +90,6 @@ network_g:
   num_block: 23
   num_grow_ch: 32
 
-
 network_d:
   type: UNetDiscriminatorSN
   num_in_ch: 3
@@ -100,7 +99,7 @@ network_d:
 # path
 path:
   # use the pre-trained Real-ESRNet model
-  pretrain_network_g: experiments/train_RealESRNetx4plus_1000k_B12G4_fromESRGAN/models/net_g_1000000.pth
+  pretrain_network_g: experiments/pretrained_models/RealESRNet_x4plus.pth
   param_key_g: params_ema
   strict_load_g: true
   resume_state: ~
@@ -166,7 +165,7 @@ train:
 #   save_img: True
 
 #   metrics:
-#     psnr: # metric name, can be arbitrary
+#     psnr: # metric name
 #       type: calculate_psnr
 #       crop_border: 4
 #       test_y_channel: false

options/train_realesrnet_x2plus.yml

@@ -0,0 +1,145 @@
+# general settings
+name: train_RealESRNetx2plus_1000k_B12G4
+model_type: RealESRNetModel
+scale: 2
+num_gpu: auto  # auto: can infer from your visible devices automatically. official: 4 GPUs
+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
+#       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

options/train_realesrnet_x4plus.yml

@@ -1,8 +1,8 @@
 # general settings
-name: train_RealESRNetx4plus_1000k_B12G4_fromESRGAN
+name: train_RealESRNetx4plus_1000k_B12G4
 model_type: RealESRNetModel
 scale: 4
-num_gpu: 4
+num_gpu: auto  # auto: can infer from your visible devices automatically. official: 4 GPUs
 manual_seed: 0
 
 # ----------------- options for synthesizing training data in RealESRNetModel ----------------- #
@@ -36,7 +36,7 @@ datasets:
     name: DF2K+OST
     type: RealESRGANDataset
     dataroot_gt: datasets/DF2K
-    meta_info: realesrgan/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
 
@@ -124,7 +124,7 @@ train:
 #   save_img: True
 
 #   metrics:
-#     psnr: # metric name, can be arbitrary
+#     psnr: # metric name
 #       type: calculate_psnr
 #       crop_border: 4
 #       test_y_channel: false

realesrgan/__init__.py

@@ -3,4 +3,4 @@ from .archs import *
 from .data import *
 from .models import *
 from .utils import *
-#from .version import __gitsha__, __version__
+from .version import *

realesrgan/archs/discriminator_arch.py

@@ -6,15 +6,23 @@ from torch.nn.utils import spectral_norm
 
 @ARCH_REGISTRY.register()
 class UNetDiscriminatorSN(nn.Module):
-    """Defines a U-Net discriminator with spectral normalization (SN)"""
+    """Defines a U-Net discriminator with spectral normalization (SN)
+
+    It is used in Real-ESRGAN: Training Real-World Blind Super-Resolution with Pure Synthetic Data.
+
+    Arg:
+        num_in_ch (int): Channel number of inputs. Default: 3.
+        num_feat (int): Channel number of base intermediate features. Default: 64.
+        skip_connection (bool): Whether to use skip connections between U-Net. Default: True.
+    """
 
     def __init__(self, num_in_ch, num_feat=64, skip_connection=True):
         super(UNetDiscriminatorSN, self).__init__()
         self.skip_connection = skip_connection
         norm = spectral_norm
-
+        # the first convolution
         self.conv0 = nn.Conv2d(num_in_ch, num_feat, kernel_size=3, stride=1, padding=1)
-
+        # downsample
         self.conv1 = norm(nn.Conv2d(num_feat, num_feat * 2, 4, 2, 1, bias=False))
         self.conv2 = norm(nn.Conv2d(num_feat * 2, num_feat * 4, 4, 2, 1, bias=False))
         self.conv3 = norm(nn.Conv2d(num_feat * 4, num_feat * 8, 4, 2, 1, bias=False))
@@ -22,14 +30,13 @@ class UNetDiscriminatorSN(nn.Module):
         self.conv4 = norm(nn.Conv2d(num_feat * 8, num_feat * 4, 3, 1, 1, bias=False))
         self.conv5 = norm(nn.Conv2d(num_feat * 4, num_feat * 2, 3, 1, 1, bias=False))
         self.conv6 = norm(nn.Conv2d(num_feat * 2, num_feat, 3, 1, 1, bias=False))
-
-        # extra
+        # extra convolutions
         self.conv7 = norm(nn.Conv2d(num_feat, num_feat, 3, 1, 1, bias=False))
         self.conv8 = norm(nn.Conv2d(num_feat, num_feat, 3, 1, 1, bias=False))
-
         self.conv9 = nn.Conv2d(num_feat, 1, 3, 1, 1)
 
     def forward(self, x):
+        # downsample
         x0 = F.leaky_relu(self.conv0(x), negative_slope=0.2, inplace=True)
         x1 = F.leaky_relu(self.conv1(x0), negative_slope=0.2, inplace=True)
         x2 = F.leaky_relu(self.conv2(x1), negative_slope=0.2, inplace=True)
@@ -52,7 +59,7 @@ class UNetDiscriminatorSN(nn.Module):
         if self.skip_connection:
             x6 = x6 + x0
 
-        # extra
+        # extra convolutions
         out = F.leaky_relu(self.conv7(x6), negative_slope=0.2, inplace=True)
         out = F.leaky_relu(self.conv8(out), negative_slope=0.2, inplace=True)
         out = self.conv9(out)

realesrgan/archs/srvgg_arch.py

@@ -0,0 +1,69 @@
+from basicsr.utils.registry import ARCH_REGISTRY
+from torch import nn as nn
+from torch.nn import functional as F
+
+
+@ARCH_REGISTRY.register()
+class SRVGGNetCompact(nn.Module):
+    """A compact VGG-style network structure for super-resolution.
+
+    It is a compact network structure, which performs upsampling in the last layer and no convolution is
+    conducted on the HR feature space.
+
+    Args:
+        num_in_ch (int): Channel number of inputs. Default: 3.
+        num_out_ch (int): Channel number of outputs. Default: 3.
+        num_feat (int): Channel number of intermediate features. Default: 64.
+        num_conv (int): Number of convolution layers in the body network. Default: 16.
+        upscale (int): Upsampling factor. Default: 4.
+        act_type (str): Activation type, options: 'relu', 'prelu', 'leakyrelu'. Default: prelu.
+    """
+
+    def __init__(self, num_in_ch=3, num_out_ch=3, num_feat=64, num_conv=16, upscale=4, act_type='prelu'):
+        super(SRVGGNetCompact, self).__init__()
+        self.num_in_ch = num_in_ch
+        self.num_out_ch = num_out_ch
+        self.num_feat = num_feat
+        self.num_conv = num_conv
+        self.upscale = upscale
+        self.act_type = act_type
+
+        self.body = nn.ModuleList()
+        # the first conv
+        self.body.append(nn.Conv2d(num_in_ch, num_feat, 3, 1, 1))
+        # the first activation
+        if act_type == 'relu':
+            activation = nn.ReLU(inplace=True)
+        elif act_type == 'prelu':
+            activation = nn.PReLU(num_parameters=num_feat)
+        elif act_type == 'leakyrelu':
+            activation = nn.LeakyReLU(negative_slope=0.1, inplace=True)
+        self.body.append(activation)
+
+        # the body structure
+        for _ in range(num_conv):
+            self.body.append(nn.Conv2d(num_feat, num_feat, 3, 1, 1))
+            # activation
+            if act_type == 'relu':
+                activation = nn.ReLU(inplace=True)
+            elif act_type == 'prelu':
+                activation = nn.PReLU(num_parameters=num_feat)
+            elif act_type == 'leakyrelu':
+                activation = nn.LeakyReLU(negative_slope=0.1, inplace=True)
+            self.body.append(activation)
+
+        # the last conv
+        self.body.append(nn.Conv2d(num_feat, num_out_ch * upscale * upscale, 3, 1, 1))
+        # upsample
+        self.upsampler = nn.PixelShuffle(upscale)
+
+    def forward(self, x):
+        out = x
+        for i in range(0, len(self.body)):
+            out = self.body[i](out)
+
+        out = self.upsampler(out)
+        # add the nearest upsampled image, so that the network learns the residual
+        base = F.interpolate(x, scale_factor=self.upscale, mode='nearest')
+        out += base
+        return out

realesrgan/data/realesrgan_dataset.py

@@ -15,18 +15,31 @@ from torch.utils import data as data
 
 @DATASET_REGISTRY.register()
 class RealESRGANDataset(data.Dataset):
-    """
-    Dataset used for Real-ESRGAN model.
+    """Dataset used for Real-ESRGAN model:
+    Real-ESRGAN: Training Real-World Blind Super-Resolution with Pure Synthetic Data.
+
+    It loads gt (Ground-Truth) images, and augments them.
+    It also generates blur kernels and sinc kernels for generating low-quality images.
+    Note that the low-quality images are processed in tensors on GPUS for faster processing.
+
+    Args:
+        opt (dict): Config for train datasets. It contains the following keys:
+            dataroot_gt (str): Data root path for gt.
+            meta_info (str): Path for meta information file.
+            io_backend (dict): IO backend type and other kwarg.
+            use_hflip (bool): Use horizontal flips.
+            use_rot (bool): Use rotation (use vertical flip and transposing h and w for implementation).
+            Please see more options in the codes.
     """
 
     def __init__(self, opt):
         super(RealESRGANDataset, self).__init__()
         self.opt = opt
-        # file client (io backend)
         self.file_client = None
         self.io_backend_opt = opt['io_backend']
         self.gt_folder = opt['dataroot_gt']
 
+        # file client (lmdb io backend)
         if self.io_backend_opt['type'] == 'lmdb':
             self.io_backend_opt['db_paths'] = [self.gt_folder]
             self.io_backend_opt['client_keys'] = ['gt']
@@ -35,18 +48,20 @@ class RealESRGANDataset(data.Dataset):
             with open(osp.join(self.gt_folder, 'meta_info.txt')) as fin:
                 self.paths = [line.split('.')[0] for line in fin]
         else:
+            # disk backend with meta_info
+            # Each line in the meta_info describes the relative path to an image
             with open(self.opt['meta_info']) as fin:
-                paths = [line.strip() for line in fin]
+                paths = [line.strip().split(' ')[0] for line in fin]
                 self.paths = [os.path.join(self.gt_folder, v) for v in paths]
 
         # blur settings for the first degradation
         self.blur_kernel_size = opt['blur_kernel_size']
         self.kernel_list = opt['kernel_list']
-        self.kernel_prob = opt['kernel_prob']
+        self.kernel_prob = opt['kernel_prob']  # a list for each kernel probability
         self.blur_sigma = opt['blur_sigma']
-        self.betag_range = opt['betag_range']
-        self.betap_range = opt['betap_range']
-        self.sinc_prob = opt['sinc_prob']
+        self.betag_range = opt['betag_range']  # betag used in generalized Gaussian blur kernels
+        self.betap_range = opt['betap_range']  # betap used in plateau blur kernels
+        self.sinc_prob = opt['sinc_prob']  # the probability for sinc filters
 
         # blur settings for the second degradation
         self.blur_kernel_size2 = opt['blur_kernel_size2']
@@ -61,6 +76,7 @@ class RealESRGANDataset(data.Dataset):
         self.final_sinc_prob = opt['final_sinc_prob']
 
         self.kernel_range = [2 * v + 1 for v in range(3, 11)]  # kernel size ranges from 7 to 21
+        # TODO: kernel range is now hard-coded, should be in the configure file
         self.pulse_tensor = torch.zeros(21, 21).float()  # convolving with pulse tensor brings no blurry effect
         self.pulse_tensor[10, 10] = 1
 
@@ -76,7 +92,7 @@ class RealESRGANDataset(data.Dataset):
         while retry > 0:
             try:
                 img_bytes = self.file_client.get(gt_path, 'gt')
-            except Exception as e:
+            except (IOError, OSError) as e:
                 logger = get_root_logger()
                 logger.warn(f'File client error: {e}, remaining retry times: {retry - 1}')
                 # change another file to read
@@ -89,10 +105,11 @@ class RealESRGANDataset(data.Dataset):
                 retry -= 1
         img_gt = imfrombytes(img_bytes, float32=True)
 
-        # -------------------- augmentation for training: flip, rotation -------------------- #
+        # -------------------- Do augmentation for training: flip, rotation -------------------- #
         img_gt = augment(img_gt, self.opt['use_hflip'], self.opt['use_rot'])
 
-        # crop or pad to 400: 400 is hard-coded. You may change it accordingly
+        # crop or pad to 400
+        # TODO: 400 is hard-coded. You may change it accordingly
         h, w = img_gt.shape[0:2]
         crop_pad_size = 400
         # pad
@@ -154,7 +171,7 @@ class RealESRGANDataset(data.Dataset):
         pad_size = (21 - kernel_size) // 2
         kernel2 = np.pad(kernel2, ((pad_size, pad_size), (pad_size, pad_size)))
 
-        # ------------------------------------- sinc kernel ------------------------------------- #
+        # ------------------------------------- the final sinc kernel ------------------------------------- #
         if np.random.uniform() < self.opt['final_sinc_prob']:
             kernel_size = random.choice(self.kernel_range)
             omega_c = np.random.uniform(np.pi / 3, np.pi)

realesrgan/data/realesrgan_paired_dataset.py

@@ -0,0 +1,108 @@
+import os
+from basicsr.data.data_util import paired_paths_from_folder, paired_paths_from_lmdb
+from basicsr.data.transforms import augment, paired_random_crop
+from basicsr.utils import FileClient, imfrombytes, img2tensor
+from basicsr.utils.registry import DATASET_REGISTRY
+from torch.utils import data as data
+from torchvision.transforms.functional import normalize
+
+
+@DATASET_REGISTRY.register()
+class RealESRGANPairedDataset(data.Dataset):
+    """Paired image dataset for image restoration.
+
+    Read LQ (Low Quality, e.g. LR (Low Resolution), blurry, noisy, etc) and GT image pairs.
+
+    There are three modes:
+    1. 'lmdb': Use lmdb files.
+        If opt['io_backend'] == lmdb.
+    2. 'meta_info': Use meta information file to generate paths.
+        If opt['io_backend'] != lmdb and opt['meta_info'] is not None.
+    3. 'folder': Scan folders to generate paths.
+        The rest.
+
+    Args:
+        opt (dict): Config for train datasets. It contains the following keys:
+            dataroot_gt (str): Data root path for gt.
+            dataroot_lq (str): Data root path for lq.
+            meta_info (str): Path for meta information file.
+            io_backend (dict): IO backend type and other kwarg.
+            filename_tmpl (str): Template for each filename. Note that the template excludes the file extension.
+                Default: '{}'.
+            gt_size (int): Cropped patched size for gt patches.
+            use_hflip (bool): Use horizontal flips.
+            use_rot (bool): Use rotation (use vertical flip and transposing h
+                and w for implementation).
+
+            scale (bool): Scale, which will be added automatically.
+            phase (str): 'train' or 'val'.
+    """
+
+    def __init__(self, opt):
+        super(RealESRGANPairedDataset, self).__init__()
+        self.opt = opt
+        self.file_client = None
+        self.io_backend_opt = opt['io_backend']
+        # mean and std for normalizing the input images
+        self.mean = opt['mean'] if 'mean' in opt else None
+        self.std = opt['std'] if 'std' in opt else None
+
+        self.gt_folder, self.lq_folder = opt['dataroot_gt'], opt['dataroot_lq']
+        self.filename_tmpl = opt['filename_tmpl'] if 'filename_tmpl' in opt else '{}'
+
+        # file client (lmdb io backend)
+        if self.io_backend_opt['type'] == 'lmdb':
+            self.io_backend_opt['db_paths'] = [self.lq_folder, self.gt_folder]
+            self.io_backend_opt['client_keys'] = ['lq', 'gt']
+            self.paths = paired_paths_from_lmdb([self.lq_folder, self.gt_folder], ['lq', 'gt'])
+        elif 'meta_info' in self.opt and self.opt['meta_info'] is not None:
+            # disk backend with meta_info
+            # Each line in the meta_info describes the relative path to an image
+            with open(self.opt['meta_info']) as fin:
+                paths = [line.strip() for line in fin]
+            self.paths = []
+            for path in paths:
+                gt_path, lq_path = path.split(', ')
+                gt_path = os.path.join(self.gt_folder, gt_path)
+                lq_path = os.path.join(self.lq_folder, lq_path)
+                self.paths.append(dict([('gt_path', gt_path), ('lq_path', lq_path)]))
+        else:
+            # disk backend
+            # it will scan the whole folder to get meta info
+            # it will be time-consuming for folders with too many files. It is recommended using an extra meta txt file
+            self.paths = paired_paths_from_folder([self.lq_folder, self.gt_folder], ['lq', 'gt'], self.filename_tmpl)
+
+    def __getitem__(self, index):
+        if self.file_client is None:
+            self.file_client = FileClient(self.io_backend_opt.pop('type'), **self.io_backend_opt)
+
+        scale = self.opt['scale']
+
+        # Load gt and lq images. Dimension order: HWC; channel order: BGR;
+        # image range: [0, 1], float32.
+        gt_path = self.paths[index]['gt_path']
+        img_bytes = self.file_client.get(gt_path, 'gt')
+        img_gt = imfrombytes(img_bytes, float32=True)
+        lq_path = self.paths[index]['lq_path']
+        img_bytes = self.file_client.get(lq_path, 'lq')
+        img_lq = imfrombytes(img_bytes, float32=True)
+
+        # augmentation for training
+        if self.opt['phase'] == 'train':
+            gt_size = self.opt['gt_size']
+            # random crop
+            img_gt, img_lq = paired_random_crop(img_gt, img_lq, gt_size, scale, gt_path)
+            # flip, rotation
+            img_gt, img_lq = augment([img_gt, img_lq], self.opt['use_hflip'], self.opt['use_rot'])
+
+        # BGR to RGB, HWC to CHW, numpy to tensor
+        img_gt, img_lq = img2tensor([img_gt, img_lq], bgr2rgb=True, float32=True)
+        # normalize
+        if self.mean is not None or self.std is not None:
+            normalize(img_lq, self.mean, self.std, inplace=True)
+            normalize(img_gt, self.mean, self.std, inplace=True)
+
+        return {'lq': img_lq, 'gt': img_gt, 'lq_path': lq_path, 'gt_path': gt_path}
+
+    def __len__(self):
+        return len(self.paths)

realesrgan/models/realesrgan_model.py

@@ -13,35 +13,45 @@ from torch.nn import functional as F
 
 @MODEL_REGISTRY.register()
 class RealESRGANModel(SRGANModel):
-    """RealESRGAN Model"""
+    """RealESRGAN Model for Real-ESRGAN: Training Real-World Blind Super-Resolution with Pure Synthetic Data.
+
+    It mainly performs:
+    1. randomly synthesize LQ images in GPU tensors
+    2. optimize the networks with GAN training.
+    """
 
     def __init__(self, opt):
         super(RealESRGANModel, self).__init__(opt)
-        self.jpeger = DiffJPEG(differentiable=False).cuda()
-        self.usm_sharpener = USMSharp().cuda()
-        self.queue_size = opt['queue_size']
+        self.jpeger = DiffJPEG(differentiable=False).cuda()  # simulate JPEG compression artifacts
+        self.usm_sharpener = USMSharp().cuda()  # do usm sharpening
+        self.queue_size = opt.get('queue_size', 180)
 
     @torch.no_grad()
     def _dequeue_and_enqueue(self):
-        # training pair pool
+        """It is the training pair pool for increasing the diversity in a batch.
+
+        Batch processing limits the diversity of synthetic degradations in a batch. For example, samples in a
+        batch could not have different resize scaling factors. Therefore, we employ this training pair pool
+        to increase the degradation diversity in a batch.
+        """
         # initialize
         b, c, h, w = self.lq.size()
         if not hasattr(self, 'queue_lr'):
-            assert self.queue_size % b == 0, 'queue size should be divisible by batch size'
+            assert self.queue_size % b == 0, f'queue size {self.queue_size} should be divisible by batch size {b}'
             self.queue_lr = torch.zeros(self.queue_size, c, h, w).cuda()
             _, c, h, w = self.gt.size()
             self.queue_gt = torch.zeros(self.queue_size, c, h, w).cuda()
             self.queue_ptr = 0
-        if self.queue_ptr == self.queue_size:  # full
+        if self.queue_ptr == self.queue_size:  # the pool is full
             # do dequeue and enqueue
             # shuffle
             idx = torch.randperm(self.queue_size)
             self.queue_lr = self.queue_lr[idx]
             self.queue_gt = self.queue_gt[idx]
-            # get
+            # get first b samples
             lq_dequeue = self.queue_lr[0:b, :, :, :].clone()
             gt_dequeue = self.queue_gt[0:b, :, :, :].clone()
-            # update
+            # update the queue
             self.queue_lr[0:b, :, :, :] = self.lq.clone()
             self.queue_gt[0:b, :, :, :] = self.gt.clone()
 
@@ -55,7 +65,9 @@ class RealESRGANModel(SRGANModel):
 
     @torch.no_grad()
     def feed_data(self, data):
-        if self.is_train:
+        """Accept data from dataloader, and then add two-order degradations to obtain LQ images.
+        """
+        if self.is_train and self.opt.get('high_order_degradation', True):
             # training data synthesis
             self.gt = data['gt'].to(self.device)
             self.gt_usm = self.usm_sharpener(self.gt)
@@ -79,7 +91,7 @@ class RealESRGANModel(SRGANModel):
                 scale = 1
             mode = random.choice(['area', 'bilinear', 'bicubic'])
             out = F.interpolate(out, scale_factor=scale, mode=mode)
-            # noise
+            # add noise
             gray_noise_prob = self.opt['gray_noise_prob']
             if np.random.uniform() < self.opt['gaussian_noise_prob']:
                 out = random_add_gaussian_noise_pt(
@@ -93,7 +105,7 @@ class RealESRGANModel(SRGANModel):
                     rounds=False)
             # JPEG compression
             jpeg_p = out.new_zeros(out.size(0)).uniform_(*self.opt['jpeg_range'])
-            out = torch.clamp(out, 0, 1)
+            out = torch.clamp(out, 0, 1)  # clamp to [0, 1], otherwise JPEGer will result in unpleasant artifacts
             out = self.jpeger(out, quality=jpeg_p)
 
             # ----------------------- The second degradation process ----------------------- #
@@ -111,7 +123,7 @@ class RealESRGANModel(SRGANModel):
             mode = random.choice(['area', 'bilinear', 'bicubic'])
             out = F.interpolate(
                 out, size=(int(ori_h / self.opt['scale'] * scale), int(ori_w / self.opt['scale'] * scale)), mode=mode)
-            # noise
+            # add noise
             gray_noise_prob = self.opt['gray_noise_prob2']
             if np.random.uniform() < self.opt['gaussian_noise_prob2']:
                 out = random_add_gaussian_noise_pt(
@@ -162,10 +174,13 @@ class RealESRGANModel(SRGANModel):
             self._dequeue_and_enqueue()
             # sharpen self.gt again, as we have changed the self.gt with self._dequeue_and_enqueue
             self.gt_usm = self.usm_sharpener(self.gt)
+            self.lq = self.lq.contiguous()  # for the warning: grad and param do not obey the gradient layout contract
         else:
+            # for paired training or validation
             self.lq = data['lq'].to(self.device)
             if 'gt' in data:
                 self.gt = data['gt'].to(self.device)
+                self.gt_usm = self.usm_sharpener(self.gt)
 
     def nondist_validation(self, dataloader, current_iter, tb_logger, save_img):
         # do not use the synthetic process during validation
@@ -174,6 +189,7 @@ class RealESRGANModel(SRGANModel):
         self.is_train = True
 
     def optimize_parameters(self, current_iter):
+        # usm sharpening
         l1_gt = self.gt_usm
         percep_gt = self.gt_usm
         gan_gt = self.gt_usm

realesrgan/models/realesrnet_model.py

@@ -12,35 +12,46 @@ from torch.nn import functional as F
 
 @MODEL_REGISTRY.register()
 class RealESRNetModel(SRModel):
-    """RealESRNet Model"""
+    """RealESRNet Model for Real-ESRGAN: Training Real-World Blind Super-Resolution with Pure Synthetic Data.
+
+    It is trained without GAN losses.
+    It mainly performs:
+    1. randomly synthesize LQ images in GPU tensors
+    2. optimize the networks with GAN training.
+    """
 
     def __init__(self, opt):
         super(RealESRNetModel, self).__init__(opt)
-        self.jpeger = DiffJPEG(differentiable=False).cuda()
-        self.usm_sharpener = USMSharp().cuda()
-        self.queue_size = opt['queue_size']
+        self.jpeger = DiffJPEG(differentiable=False).cuda()  # simulate JPEG compression artifacts
+        self.usm_sharpener = USMSharp().cuda()  # do usm sharpening
+        self.queue_size = opt.get('queue_size', 180)
 
     @torch.no_grad()
     def _dequeue_and_enqueue(self):
-        # training pair pool
+        """It is the training pair pool for increasing the diversity in a batch.
+
+        Batch processing limits the diversity of synthetic degradations in a batch. For example, samples in a
+        batch could not have different resize scaling factors. Therefore, we employ this training pair pool
+        to increase the degradation diversity in a batch.
+        """
         # initialize
         b, c, h, w = self.lq.size()
         if not hasattr(self, 'queue_lr'):
-            assert self.queue_size % b == 0, 'queue size should be divisible by batch size'
+            assert self.queue_size % b == 0, f'queue size {self.queue_size} should be divisible by batch size {b}'
             self.queue_lr = torch.zeros(self.queue_size, c, h, w).cuda()
             _, c, h, w = self.gt.size()
             self.queue_gt = torch.zeros(self.queue_size, c, h, w).cuda()
             self.queue_ptr = 0
-        if self.queue_ptr == self.queue_size:  # full
+        if self.queue_ptr == self.queue_size:  # the pool is full
             # do dequeue and enqueue
             # shuffle
             idx = torch.randperm(self.queue_size)
             self.queue_lr = self.queue_lr[idx]
             self.queue_gt = self.queue_gt[idx]
-            # get
+            # get first b samples
             lq_dequeue = self.queue_lr[0:b, :, :, :].clone()
             gt_dequeue = self.queue_gt[0:b, :, :, :].clone()
-            # update
+            # update the queue
             self.queue_lr[0:b, :, :, :] = self.lq.clone()
             self.queue_gt[0:b, :, :, :] = self.gt.clone()
 
@@ -54,10 +65,12 @@ class RealESRNetModel(SRModel):
 
     @torch.no_grad()
     def feed_data(self, data):
-        if self.is_train:
+        """Accept data from dataloader, and then add two-order degradations to obtain LQ images.
+        """
+        if self.is_train and self.opt.get('high_order_degradation', True):
             # training data synthesis
             self.gt = data['gt'].to(self.device)
-            # USM the GT images
+            # USM sharpen the GT images
             if self.opt['gt_usm'] is True:
                 self.gt = self.usm_sharpener(self.gt)
 
@@ -80,7 +93,7 @@ class RealESRNetModel(SRModel):
                 scale = 1
             mode = random.choice(['area', 'bilinear', 'bicubic'])
             out = F.interpolate(out, scale_factor=scale, mode=mode)
-            # noise
+            # add noise
             gray_noise_prob = self.opt['gray_noise_prob']
             if np.random.uniform() < self.opt['gaussian_noise_prob']:
                 out = random_add_gaussian_noise_pt(
@@ -94,7 +107,7 @@ class RealESRNetModel(SRModel):
                     rounds=False)
             # JPEG compression
             jpeg_p = out.new_zeros(out.size(0)).uniform_(*self.opt['jpeg_range'])
-            out = torch.clamp(out, 0, 1)
+            out = torch.clamp(out, 0, 1)  # clamp to [0, 1], otherwise JPEGer will result in unpleasant artifacts
             out = self.jpeger(out, quality=jpeg_p)
 
             # ----------------------- The second degradation process ----------------------- #
@@ -112,7 +125,7 @@ class RealESRNetModel(SRModel):
             mode = random.choice(['area', 'bilinear', 'bicubic'])
             out = F.interpolate(
                 out, size=(int(ori_h / self.opt['scale'] * scale), int(ori_w / self.opt['scale'] * scale)), mode=mode)
-            # noise
+            # add noise
             gray_noise_prob = self.opt['gray_noise_prob2']
             if np.random.uniform() < self.opt['gaussian_noise_prob2']:
                 out = random_add_gaussian_noise_pt(
@@ -160,10 +173,13 @@ class RealESRNetModel(SRModel):
 
             # training pair pool
             self._dequeue_and_enqueue()
+            self.lq = self.lq.contiguous()  # for the warning: grad and param do not obey the gradient layout contract
         else:
+            # for paired training or validation
             self.lq = data['lq'].to(self.device)
             if 'gt' in data:
                 self.gt = data['gt'].to(self.device)
+                self.gt_usm = self.usm_sharpener(self.gt)
 
     def nondist_validation(self, dataloader, current_iter, tb_logger, save_img):
         # do not use the synthetic process during validation

realesrgan/utils.py

@@ -2,18 +2,31 @@ import cv2
 import math
 import numpy as np
 import os
+import queue
+import threading
 import torch
-from basicsr.archs.rrdbnet_arch import RRDBNet
-from torch.hub import download_url_to_file, get_dir
+from basicsr.utils.download_util import load_file_from_url
 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():
+    """A helper class for upsampling images with RealESRGAN.
 
-    def __init__(self, scale, model_path, tile=0, tile_pad=10, pre_pad=10, half=False):
+    Args:
+        scale (int): Upsampling scale factor used in the networks. It is usually 2 or 4.
+        model_path (str): The path to the pretrained model. It can be urls (will first download it automatically).
+        model (nn.Module): The defined network. Default: None.
+        tile (int): As too large images result in the out of GPU memory issue, so this tile option will first crop
+            input images into tiles, and then process each of them. Finally, they will be merged into one image.
+            0 denotes for do not use tile. Default: 0.
+        tile_pad (int): The pad size for each tile, to remove border artifacts. Default: 10.
+        pre_pad (int): Pad the input images to avoid border artifacts. Default: 10.
+        half (float): Whether to use half precision during inference. Default: False.
+    """
+
+    def __init__(self, scale, model_path, model=None, tile=0, tile_pad=10, pre_pad=10, half=False):
         self.scale = scale
         self.tile_size = tile
         self.tile_pad = tile_pad
@@ -23,12 +36,12 @@ class RealESRGANer():
 
         # 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 the model_path starts with https, it will first download models to the folder: realesrgan/weights
         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)
+                url=model_path, model_dir=os.path.join(ROOT_DIR, 'realesrgan/weights'), progress=True, file_name=None)
+        loadnet = torch.load(model_path, map_location=torch.device('cpu'))
+        # prefer to use params_ema
         if 'params_ema' in loadnet:
             keyname = 'params_ema'
         else:
@@ -40,6 +53,8 @@ class RealESRGANer():
             self.model = self.model.half()
 
     def pre_process(self, img):
+        """Pre-process, such as pre-pad and mod pad, so that the images can be divisible
+        """
         img = torch.from_numpy(np.transpose(img, (2, 0, 1))).float()
         self.img = img.unsqueeze(0).to(self.device)
         if self.half:
@@ -48,7 +63,7 @@ class RealESRGANer():
         # pre_pad
         if self.pre_pad != 0:
             self.img = F.pad(self.img, (0, self.pre_pad, 0, self.pre_pad), 'reflect')
-        # mod pad
+        # mod pad for divisible borders
         if self.scale == 2:
             self.mod_scale = 2
         elif self.scale == 1:
@@ -63,10 +78,14 @@ class RealESRGANer():
             self.img = F.pad(self.img, (0, self.mod_pad_w, 0, self.mod_pad_h), 'reflect')
 
     def process(self):
+        # model inference
         self.output = self.model(self.img)
 
     def tile_process(self):
-        """Modified from: https://github.com/ata4/esrgan-launcher
+        """It will first crop input images to tiles, and then process each tile.
+        Finally, all the processed tiles are merged into one images.
+
+        Modified from: https://github.com/ata4/esrgan-launcher
         """
         batch, channel, height, width = self.img.shape
         output_height = height * self.scale
@@ -106,7 +125,7 @@ class RealESRGANer():
                 try:
                     with torch.no_grad():
                         output_tile = self.model(input_tile)
-                except Exception as error:
+                except RuntimeError as error:
                     print('Error', error)
                 print(f'\tTile {tile_idx}/{tiles_x * tiles_y}')
 
@@ -143,7 +162,7 @@ class RealESRGANer():
         h_input, w_input = img.shape[0:2]
         # img: numpy
         img = img.astype(np.float32)
-        if np.max(img) > 255:  # 16-bit image
+        if np.max(img) > 256:  # 16-bit image
             max_range = 65535
             print('\tInput is a 16-bit image')
         else:
@@ -187,7 +206,7 @@ class RealESRGANer():
                 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:
+            else:  # use the cv2 resize for alpha channel
                 h, w = alpha.shape[0:2]
                 output_alpha = cv2.resize(alpha, (w * self.scale, h * self.scale), interpolation=cv2.INTER_LINEAR)
 
@@ -211,21 +230,51 @@ class RealESRGANer():
         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
+class PrefetchReader(threading.Thread):
+    """Prefetch images.
+
+    Args:
+        img_list (list[str]): A image list of image paths to be read.
+        num_prefetch_queue (int): Number of prefetch queue.
     """
-    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
+
+    def __init__(self, img_list, num_prefetch_queue):
+        super().__init__()
+        self.que = queue.Queue(num_prefetch_queue)
+        self.img_list = img_list
+
+    def run(self):
+        for img_path in self.img_list:
+            img = cv2.imread(img_path, cv2.IMREAD_UNCHANGED)
+            self.que.put(img)
+
+        self.que.put(None)
+
+    def __next__(self):
+        next_item = self.que.get()
+        if next_item is None:
+            raise StopIteration
+        return next_item
+
+    def __iter__(self):
+        return self
+
+
+class IOConsumer(threading.Thread):
+
+    def __init__(self, opt, que, qid):
+        super().__init__()
+        self._queue = que
+        self.qid = qid
+        self.opt = opt
+
+    def run(self):
+        while True:
+            msg = self._queue.get()
+            if isinstance(msg, str) and msg == 'quit':
+                break
+
+            output = msg['output']
+            save_path = msg['save_path']
+            cv2.imwrite(save_path, output)
+        print(f'IO worker {self.qid} is done.')

scripts/extract_subimages.py

@@ -0,0 +1,135 @@
+import argparse
+import cv2
+import numpy as np
+import os
+import sys
+from basicsr.utils import scandir
+from multiprocessing import Pool
+from os import path as osp
+from tqdm import tqdm
+
+
+def main(args):
+    """A multi-thread tool to crop large images to sub-images for faster IO.
+
+    opt (dict): Configuration dict. It contains:
+        n_thread (int): Thread number.
+        compression_level (int):  CV_IMWRITE_PNG_COMPRESSION from 0 to 9. A higher value means a smaller size
+            and longer compression time. Use 0 for faster CPU decompression. Default: 3, same in cv2.
+        input_folder (str): Path to the input folder.
+        save_folder (str): Path to save folder.
+        crop_size (int): Crop size.
+        step (int): Step for overlapped sliding window.
+        thresh_size (int): Threshold size. Patches whose size is lower than thresh_size will be dropped.
+
+    Usage:
+        For each folder, run this script.
+        Typically, there are GT folder and LQ folder to be processed for DIV2K dataset.
+        After process, each sub_folder should have the same number of subimages.
+        Remember to modify opt configurations according to your settings.
+    """
+
+    opt = {}
+    opt['n_thread'] = args.n_thread
+    opt['compression_level'] = args.compression_level
+    opt['input_folder'] = args.input
+    opt['save_folder'] = args.output
+    opt['crop_size'] = args.crop_size
+    opt['step'] = args.step
+    opt['thresh_size'] = args.thresh_size
+    extract_subimages(opt)
+
+
+def extract_subimages(opt):
+    """Crop images to subimages.
+
+    Args:
+        opt (dict): Configuration dict. It contains:
+            input_folder (str): Path to the input folder.
+            save_folder (str): Path to save folder.
+            n_thread (int): Thread number.
+    """
+    input_folder = opt['input_folder']
+    save_folder = opt['save_folder']
+    if not osp.exists(save_folder):
+        os.makedirs(save_folder)
+        print(f'mkdir {save_folder} ...')
+    else:
+        print(f'Folder {save_folder} already exists. Exit.')
+        sys.exit(1)
+
+    # scan all images
+    img_list = list(scandir(input_folder, full_path=True))
+
+    pbar = tqdm(total=len(img_list), unit='image', desc='Extract')
+    pool = Pool(opt['n_thread'])
+    for path in img_list:
+        pool.apply_async(worker, args=(path, opt), callback=lambda arg: pbar.update(1))
+    pool.close()
+    pool.join()
+    pbar.close()
+    print('All processes done.')
+
+
+def worker(path, opt):
+    """Worker for each process.
+
+    Args:
+        path (str): Image path.
+        opt (dict): Configuration dict. It contains:
+            crop_size (int): Crop size.
+            step (int): Step for overlapped sliding window.
+            thresh_size (int): Threshold size. Patches whose size is lower than thresh_size will be dropped.
+            save_folder (str): Path to save folder.
+            compression_level (int): for cv2.IMWRITE_PNG_COMPRESSION.
+
+    Returns:
+        process_info (str): Process information displayed in progress bar.
+    """
+    crop_size = opt['crop_size']
+    step = opt['step']
+    thresh_size = opt['thresh_size']
+    img_name, extension = osp.splitext(osp.basename(path))
+
+    # remove the x2, x3, x4 and x8 in the filename for DIV2K
+    img_name = img_name.replace('x2', '').replace('x3', '').replace('x4', '').replace('x8', '')
+
+    img = cv2.imread(path, cv2.IMREAD_UNCHANGED)
+
+    h, w = img.shape[0:2]
+    h_space = np.arange(0, h - crop_size + 1, step)
+    if h - (h_space[-1] + crop_size) > thresh_size:
+        h_space = np.append(h_space, h - crop_size)
+    w_space = np.arange(0, w - crop_size + 1, step)
+    if w - (w_space[-1] + crop_size) > thresh_size:
+        w_space = np.append(w_space, w - crop_size)
+
+    index = 0
+    for x in h_space:
+        for y in w_space:
+            index += 1
+            cropped_img = img[x:x + crop_size, y:y + crop_size, ...]
+            cropped_img = np.ascontiguousarray(cropped_img)
+            cv2.imwrite(
+                osp.join(opt['save_folder'], f'{img_name}_s{index:03d}{extension}'), cropped_img,
+                [cv2.IMWRITE_PNG_COMPRESSION, opt['compression_level']])
+    process_info = f'Processing {img_name} ...'
+    return process_info
+
+
+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_HR_sub', help='Output folder')
+    parser.add_argument('--crop_size', type=int, default=480, help='Crop size')
+    parser.add_argument('--step', type=int, default=240, help='Step for overlapped sliding window')
+    parser.add_argument(
+        '--thresh_size',
+        type=int,
+        default=0,
+        help='Threshold size. Patches whose size is lower than thresh_size will be dropped.')
+    parser.add_argument('--n_thread', type=int, default=20, help='Thread number.')
+    parser.add_argument('--compression_level', type=int, default=3, help='Compression level')
+    args = parser.parse_args()
+
+    main(args)

scripts/generate_meta_info.py

@@ -0,0 +1,58 @@
+import argparse
+import cv2
+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:
+            status = True
+            if args.check:
+                # read the image once for check, as some images may have errors
+                try:
+                    img = cv2.imread(img_path)
+                except (IOError, OSError) as error:
+                    print(f'Read {img_path} error: {error}')
+                    status = False
+                if img is None:
+                    status = False
+                    print(f'Img is None: {img_path}')
+            if status:
+                # get the relative path
+                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')
+    parser.add_argument('--check', action='store_true', help='Read image to check whether it is ok')
+    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)}.')
+    os.makedirs(os.path.dirname(args.meta_info), exist_ok=True)
+
+    main(args)

scripts/generate_meta_info_pairdata.py

@@ -0,0 +1,49 @@
+import argparse
+import glob
+import os
+
+
+def main(args):
+    txt_file = open(args.meta_info, 'w')
+    # sca images
+    img_paths_gt = sorted(glob.glob(os.path.join(args.input[0], '*')))
+    img_paths_lq = sorted(glob.glob(os.path.join(args.input[1], '*')))
+
+    assert len(img_paths_gt) == len(img_paths_lq), ('GT folder and LQ folder should have the same length, but got '
+                                                    f'{len(img_paths_gt)} and {len(img_paths_lq)}.')
+
+    for img_path_gt, img_path_lq in zip(img_paths_gt, img_paths_lq):
+        # get the relative paths
+        img_name_gt = os.path.relpath(img_path_gt, args.root[0])
+        img_name_lq = os.path.relpath(img_path_lq, args.root[1])
+        print(f'{img_name_gt}, {img_name_lq}')
+        txt_file.write(f'{img_name_gt}, {img_name_lq}\n')
+
+
+if __name__ == '__main__':
+    """This script is used to generate meta info (txt file) for paired images.
+    """
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        '--input',
+        nargs='+',
+        default=['datasets/DF2K/DIV2K_train_HR_sub', 'datasets/DF2K/DIV2K_train_LR_bicubic_X4_sub'],
+        help='Input folder, should be [gt_folder, lq_folder]')
+    parser.add_argument('--root', nargs='+', default=[None, None], help='Folder root, will use the ')
+    parser.add_argument(
+        '--meta_info',
+        type=str,
+        default='datasets/DF2K/meta_info/meta_info_DIV2K_sub_pair.txt',
+        help='txt path for meta info')
+    args = parser.parse_args()
+
+    assert len(args.input) == 2, 'Input folder should have two elements: gt folder and lq folder'
+    assert len(args.root) == 2, 'Root path should have two elements: root for gt folder and lq folder'
+    os.makedirs(os.path.dirname(args.meta_info), exist_ok=True)
+    for i in range(2):
+        if args.input[i].endswith('/'):
+            args.input[i] = args.input[i][:-1]
+        if args.root[i] is None:
+            args.root[i] = os.path.dirname(args.input[i])
+
+    main(args)

scripts/generate_multiscale_DF2K.py

@@ -0,0 +1,48 @@
+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__':
+    """Generate multi-scale versions for GT images with LANCZOS resampling.
+    It is now used for DF2K dataset (DIV2K + Flickr 2K)
+    """
+    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)

scripts/pytorch2onnx.py

@@ -1,17 +1,36 @@
+import argparse
 import torch
 import torch.onnx
 from basicsr.archs.rrdbnet_arch import RRDBNet
 
-# An instance of your model
-model = RRDBNet(num_in_ch=3, num_out_ch=3, num_feat=64, num_block=23, num_grow_ch=32)
-model.load_state_dict(torch.load('experiments/pretrained_models/RealESRGAN_x4plus.pth')['params_ema'])
-# set the train mode to false since we will only run the forward pass.
-model.train(False)
-model.cpu().eval()
 
-# An example input you would normally provide to your model's forward() method
-x = torch.rand(1, 3, 64, 64)
+def main(args):
+    # An instance of the model
+    model = RRDBNet(num_in_ch=3, num_out_ch=3, num_feat=64, num_block=23, num_grow_ch=32, scale=4)
+    if args.params:
+        keyname = 'params'
+    else:
+        keyname = 'params_ema'
+    model.load_state_dict(torch.load(args.input)[keyname])
+    # set the train mode to false since we will only run the forward pass.
+    model.train(False)
+    model.cpu().eval()
 
-# Export the model
-with torch.no_grad():
-    torch_out = torch.onnx._export(model, x, 'realesrgan-x4.onnx', opset_version=11, export_params=True)
+    # An example input
+    x = torch.rand(1, 3, 64, 64)
+    # Export the model
+    with torch.no_grad():
+        torch_out = torch.onnx._export(model, x, args.output, opset_version=11, export_params=True)
+    print(torch_out.shape)
+
+
+if __name__ == '__main__':
+    """Convert pytorch model to onnx models"""
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        '--input', type=str, default='experiments/pretrained_models/RealESRGAN_x4plus.pth', help='Input model path')
+    parser.add_argument('--output', type=str, default='realesrgan-x4.onnx', help='Output onnx path')
+    parser.add_argument('--params', action='store_false', help='Use params instead of params_ema')
+    args = parser.parse_args()
+
+    main(args)

setup.py

@@ -6,7 +6,7 @@ import os
 import subprocess
 import time
 
-version_file = 'version.py'
+version_file = 'realesrgan/version.py'
 
 
 def readme():
@@ -43,12 +43,6 @@ def get_git_hash():
 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'
 

tests/data/gt.lmdb/meta_info.txt

@@ -0,0 +1,2 @@
+baboon.png (480,500,3) 1
+comic.png (360,240,3) 1

tests/data/lq.lmdb/meta_info.txt

@@ -0,0 +1,2 @@
+baboon.png (120,125,3) 1
+comic.png (80,60,3) 1

tests/data/meta_info_gt.txt

@@ -0,0 +1,2 @@
+baboon.png
+comic.png

tests/data/meta_info_pair.txt

@@ -0,0 +1,2 @@
+gt/baboon.png, lq/baboon.png
+gt/comic.png, lq/comic.png

tests/data/test_realesrgan_dataset.yml

@@ -0,0 +1,28 @@
+name: Demo
+type: RealESRGANDataset
+dataroot_gt: tests/data/gt
+meta_info: tests/data/meta_info_gt.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: 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: 1
+blur_sigma2: [0.2, 1.5]
+betag_range2: [0.5, 4]
+betap_range2: [1, 2]
+
+final_sinc_prob: 1
+
+gt_size: 128
+use_hflip: True
+use_rot: False

tests/data/test_realesrgan_model.yml

@@ -0,0 +1,115 @@
+scale: 4
+num_gpu: 1
+manual_seed: 0
+is_train: True
+dist: False
+
+# ----------------- options for synthesizing training data ----------------- #
+# 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: 1
+noise_range: [1, 30]
+poisson_scale_range: [0.05, 3]
+gray_noise_prob: 1
+jpeg_range: [30, 95]
+
+# the second degradation process
+second_blur_prob: 1
+resize_prob2: [0.3, 0.4, 0.3]  # up, down, keep
+resize_range2: [0.3, 1.2]
+gaussian_noise_prob2: 1
+noise_range2: [1, 25]
+poisson_scale_range2: [0.05, 2.5]
+gray_noise_prob2: 1
+jpeg_range2: [30, 95]
+
+gt_size: 32
+queue_size: 1
+
+# network structures
+network_g:
+  type: RRDBNet
+  num_in_ch: 3
+  num_out_ch: 3
+  num_feat: 4
+  num_block: 1
+  num_grow_ch: 2
+
+network_d:
+  type: UNetDiscriminatorSN
+  num_in_ch: 3
+  num_feat: 2
+  skip_connection: True
+
+# path
+path:
+  pretrain_network_g: ~
+  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
+
+
+# validation settings
+val:
+  val_freq: !!float 5e3
+  save_img: False

tests/data/test_realesrgan_paired_dataset.yml

@@ -0,0 +1,13 @@
+name: Demo
+type: RealESRGANPairedDataset
+scale: 4
+dataroot_gt: tests/data
+dataroot_lq: tests/data
+meta_info: tests/data/meta_info_pair.txt
+io_backend:
+  type: disk
+
+phase: train
+gt_size: 128
+use_hflip: True
+use_rot: False

tests/data/test_realesrnet_model.yml

@@ -0,0 +1,75 @@
+scale: 4
+num_gpu: 1
+manual_seed: 0
+is_train: True
+dist: False
+
+# ----------------- options for synthesizing training data ----------------- #
+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: 1
+noise_range: [1, 30]
+poisson_scale_range: [0.05, 3]
+gray_noise_prob: 1
+jpeg_range: [30, 95]
+
+# the second degradation process
+second_blur_prob: 1
+resize_prob2: [0.3, 0.4, 0.3]  # up, down, keep
+resize_range2: [0.3, 1.2]
+gaussian_noise_prob2: 1
+noise_range2: [1, 25]
+poisson_scale_range2: [0.05, 2.5]
+gray_noise_prob2: 1
+jpeg_range2: [30, 95]
+
+gt_size: 32
+queue_size: 1
+
+# network structures
+network_g:
+  type: RRDBNet
+  num_in_ch: 3
+  num_out_ch: 3
+  num_feat: 4
+  num_block: 1
+  num_grow_ch: 2
+
+# path
+path:
+  pretrain_network_g: ~
+  param_key_g: params_ema
+  strict_load_g: true
+  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
+
+
+# validation settings
+val:
+  val_freq: !!float 5e3
+  save_img: False

tests/test_dataset.py

@@ -0,0 +1,151 @@
+import pytest
+import yaml
+
+from realesrgan.data.realesrgan_dataset import RealESRGANDataset
+from realesrgan.data.realesrgan_paired_dataset import RealESRGANPairedDataset
+
+
+def test_realesrgan_dataset():
+
+    with open('tests/data/test_realesrgan_dataset.yml', mode='r') as f:
+        opt = yaml.load(f, Loader=yaml.FullLoader)
+
+    dataset = RealESRGANDataset(opt)
+    assert dataset.io_backend_opt['type'] == 'disk'  # io backend
+    assert len(dataset) == 2  # whether to read correct meta info
+    assert dataset.kernel_list == [
+        'iso', 'aniso', 'generalized_iso', 'generalized_aniso', 'plateau_iso', 'plateau_aniso'
+    ]  # correct initialization the degradation configurations
+    assert dataset.betag_range2 == [0.5, 4]
+
+    # test __getitem__
+    result = dataset.__getitem__(0)
+    # check returned keys
+    expected_keys = ['gt', 'kernel1', 'kernel2', 'sinc_kernel', 'gt_path']
+    assert set(expected_keys).issubset(set(result.keys()))
+    # check shape and contents
+    assert result['gt'].shape == (3, 400, 400)
+    assert result['kernel1'].shape == (21, 21)
+    assert result['kernel2'].shape == (21, 21)
+    assert result['sinc_kernel'].shape == (21, 21)
+    assert result['gt_path'] == 'tests/data/gt/baboon.png'
+
+    # ------------------ test lmdb backend -------------------- #
+    opt['dataroot_gt'] = 'tests/data/gt.lmdb'
+    opt['io_backend']['type'] = 'lmdb'
+
+    dataset = RealESRGANDataset(opt)
+    assert dataset.io_backend_opt['type'] == 'lmdb'  # io backend
+    assert len(dataset.paths) == 2  # whether to read correct meta info
+    assert dataset.kernel_list == [
+        'iso', 'aniso', 'generalized_iso', 'generalized_aniso', 'plateau_iso', 'plateau_aniso'
+    ]  # correct initialization the degradation configurations
+    assert dataset.betag_range2 == [0.5, 4]
+
+    # test __getitem__
+    result = dataset.__getitem__(1)
+    # check returned keys
+    expected_keys = ['gt', 'kernel1', 'kernel2', 'sinc_kernel', 'gt_path']
+    assert set(expected_keys).issubset(set(result.keys()))
+    # check shape and contents
+    assert result['gt'].shape == (3, 400, 400)
+    assert result['kernel1'].shape == (21, 21)
+    assert result['kernel2'].shape == (21, 21)
+    assert result['sinc_kernel'].shape == (21, 21)
+    assert result['gt_path'] == 'comic'
+
+    # ------------------ test with sinc_prob = 0 -------------------- #
+    opt['dataroot_gt'] = 'tests/data/gt.lmdb'
+    opt['io_backend']['type'] = 'lmdb'
+    opt['sinc_prob'] = 0
+    opt['sinc_prob2'] = 0
+    opt['final_sinc_prob'] = 0
+    dataset = RealESRGANDataset(opt)
+    result = dataset.__getitem__(0)
+    # check returned keys
+    expected_keys = ['gt', 'kernel1', 'kernel2', 'sinc_kernel', 'gt_path']
+    assert set(expected_keys).issubset(set(result.keys()))
+    # check shape and contents
+    assert result['gt'].shape == (3, 400, 400)
+    assert result['kernel1'].shape == (21, 21)
+    assert result['kernel2'].shape == (21, 21)
+    assert result['sinc_kernel'].shape == (21, 21)
+    assert result['gt_path'] == 'baboon'
+
+    # ------------------ lmdb backend should have paths ends with lmdb -------------------- #
+    with pytest.raises(ValueError):
+        opt['dataroot_gt'] = 'tests/data/gt'
+        opt['io_backend']['type'] = 'lmdb'
+        dataset = RealESRGANDataset(opt)
+
+
+def test_realesrgan_paired_dataset():
+
+    with open('tests/data/test_realesrgan_paired_dataset.yml', mode='r') as f:
+        opt = yaml.load(f, Loader=yaml.FullLoader)
+
+    dataset = RealESRGANPairedDataset(opt)
+    assert dataset.io_backend_opt['type'] == 'disk'  # io backend
+    assert len(dataset) == 2  # whether to read correct meta info
+
+    # test __getitem__
+    result = dataset.__getitem__(0)
+    # check returned keys
+    expected_keys = ['gt', 'lq', 'gt_path', 'lq_path']
+    assert set(expected_keys).issubset(set(result.keys()))
+    # check shape and contents
+    assert result['gt'].shape == (3, 128, 128)
+    assert result['lq'].shape == (3, 32, 32)
+    assert result['gt_path'] == 'tests/data/gt/baboon.png'
+    assert result['lq_path'] == 'tests/data/lq/baboon.png'
+
+    # ------------------ test lmdb backend -------------------- #
+    opt['dataroot_gt'] = 'tests/data/gt.lmdb'
+    opt['dataroot_lq'] = 'tests/data/lq.lmdb'
+    opt['io_backend']['type'] = 'lmdb'
+
+    dataset = RealESRGANPairedDataset(opt)
+    assert dataset.io_backend_opt['type'] == 'lmdb'  # io backend
+    assert len(dataset) == 2  # whether to read correct meta info
+
+    # test __getitem__
+    result = dataset.__getitem__(1)
+    # check returned keys
+    expected_keys = ['gt', 'lq', 'gt_path', 'lq_path']
+    assert set(expected_keys).issubset(set(result.keys()))
+    # check shape and contents
+    assert result['gt'].shape == (3, 128, 128)
+    assert result['lq'].shape == (3, 32, 32)
+    assert result['gt_path'] == 'comic'
+    assert result['lq_path'] == 'comic'
+
+    # ------------------ test paired_paths_from_folder -------------------- #
+    opt['dataroot_gt'] = 'tests/data/gt'
+    opt['dataroot_lq'] = 'tests/data/lq'
+    opt['io_backend'] = dict(type='disk')
+    opt['meta_info'] = None
+
+    dataset = RealESRGANPairedDataset(opt)
+    assert dataset.io_backend_opt['type'] == 'disk'  # io backend
+    assert len(dataset) == 2  # whether to read correct meta info
+
+    # test __getitem__
+    result = dataset.__getitem__(0)
+    # check returned keys
+    expected_keys = ['gt', 'lq', 'gt_path', 'lq_path']
+    assert set(expected_keys).issubset(set(result.keys()))
+    # check shape and contents
+    assert result['gt'].shape == (3, 128, 128)
+    assert result['lq'].shape == (3, 32, 32)
+
+    # ------------------ test normalization -------------------- #
+    dataset.mean = [0.5, 0.5, 0.5]
+    dataset.std = [0.5, 0.5, 0.5]
+    # test __getitem__
+    result = dataset.__getitem__(0)
+    # check returned keys
+    expected_keys = ['gt', 'lq', 'gt_path', 'lq_path']
+    assert set(expected_keys).issubset(set(result.keys()))
+    # check shape and contents
+    assert result['gt'].shape == (3, 128, 128)
+    assert result['lq'].shape == (3, 32, 32)

tests/test_discriminator_arch.py

@@ -0,0 +1,19 @@
+import torch
+
+from realesrgan.archs.discriminator_arch import UNetDiscriminatorSN
+
+
+def test_unetdiscriminatorsn():
+    """Test arch: UNetDiscriminatorSN."""
+
+    # model init and forward (cpu)
+    net = UNetDiscriminatorSN(num_in_ch=3, num_feat=4, skip_connection=True)
+    img = torch.rand((1, 3, 32, 32), dtype=torch.float32)
+    output = net(img)
+    assert output.shape == (1, 1, 32, 32)
+
+    # model init and forward (gpu)
+    if torch.cuda.is_available():
+        net.cuda()
+        output = net(img.cuda())
+        assert output.shape == (1, 1, 32, 32)

tests/test_model.py

@@ -0,0 +1,126 @@
+import torch
+import yaml
+from basicsr.archs.rrdbnet_arch import RRDBNet
+from basicsr.data.paired_image_dataset import PairedImageDataset
+from basicsr.losses.losses import GANLoss, L1Loss, PerceptualLoss
+
+from realesrgan.archs.discriminator_arch import UNetDiscriminatorSN
+from realesrgan.models.realesrgan_model import RealESRGANModel
+from realesrgan.models.realesrnet_model import RealESRNetModel
+
+
+def test_realesrnet_model():
+    with open('tests/data/test_realesrnet_model.yml', mode='r') as f:
+        opt = yaml.load(f, Loader=yaml.FullLoader)
+
+    # build model
+    model = RealESRNetModel(opt)
+    # test attributes
+    assert model.__class__.__name__ == 'RealESRNetModel'
+    assert isinstance(model.net_g, RRDBNet)
+    assert isinstance(model.cri_pix, L1Loss)
+    assert isinstance(model.optimizers[0], torch.optim.Adam)
+
+    # prepare data
+    gt = torch.rand((1, 3, 32, 32), dtype=torch.float32)
+    kernel1 = torch.rand((1, 5, 5), dtype=torch.float32)
+    kernel2 = torch.rand((1, 5, 5), dtype=torch.float32)
+    sinc_kernel = torch.rand((1, 5, 5), dtype=torch.float32)
+    data = dict(gt=gt, kernel1=kernel1, kernel2=kernel2, sinc_kernel=sinc_kernel)
+    model.feed_data(data)
+    # check dequeue
+    model.feed_data(data)
+    # check data shape
+    assert model.lq.shape == (1, 3, 8, 8)
+    assert model.gt.shape == (1, 3, 32, 32)
+
+    # change probability to test if-else
+    model.opt['gaussian_noise_prob'] = 0
+    model.opt['gray_noise_prob'] = 0
+    model.opt['second_blur_prob'] = 0
+    model.opt['gaussian_noise_prob2'] = 0
+    model.opt['gray_noise_prob2'] = 0
+    model.feed_data(data)
+    # check data shape
+    assert model.lq.shape == (1, 3, 8, 8)
+    assert model.gt.shape == (1, 3, 32, 32)
+
+    # ----------------- test nondist_validation -------------------- #
+    # construct dataloader
+    dataset_opt = dict(
+        name='Demo',
+        dataroot_gt='tests/data/gt',
+        dataroot_lq='tests/data/lq',
+        io_backend=dict(type='disk'),
+        scale=4,
+        phase='val')
+    dataset = PairedImageDataset(dataset_opt)
+    dataloader = torch.utils.data.DataLoader(dataset=dataset, batch_size=1, shuffle=False, num_workers=0)
+    assert model.is_train is True
+    model.nondist_validation(dataloader, 1, None, False)
+    assert model.is_train is True
+
+
+def test_realesrgan_model():
+    with open('tests/data/test_realesrgan_model.yml', mode='r') as f:
+        opt = yaml.load(f, Loader=yaml.FullLoader)
+
+    # build model
+    model = RealESRGANModel(opt)
+    # test attributes
+    assert model.__class__.__name__ == 'RealESRGANModel'
+    assert isinstance(model.net_g, RRDBNet)  # generator
+    assert isinstance(model.net_d, UNetDiscriminatorSN)  # discriminator
+    assert isinstance(model.cri_pix, L1Loss)
+    assert isinstance(model.cri_perceptual, PerceptualLoss)
+    assert isinstance(model.cri_gan, GANLoss)
+    assert isinstance(model.optimizers[0], torch.optim.Adam)
+    assert isinstance(model.optimizers[1], torch.optim.Adam)
+
+    # prepare data
+    gt = torch.rand((1, 3, 32, 32), dtype=torch.float32)
+    kernel1 = torch.rand((1, 5, 5), dtype=torch.float32)
+    kernel2 = torch.rand((1, 5, 5), dtype=torch.float32)
+    sinc_kernel = torch.rand((1, 5, 5), dtype=torch.float32)
+    data = dict(gt=gt, kernel1=kernel1, kernel2=kernel2, sinc_kernel=sinc_kernel)
+    model.feed_data(data)
+    # check dequeue
+    model.feed_data(data)
+    # check data shape
+    assert model.lq.shape == (1, 3, 8, 8)
+    assert model.gt.shape == (1, 3, 32, 32)
+
+    # change probability to test if-else
+    model.opt['gaussian_noise_prob'] = 0
+    model.opt['gray_noise_prob'] = 0
+    model.opt['second_blur_prob'] = 0
+    model.opt['gaussian_noise_prob2'] = 0
+    model.opt['gray_noise_prob2'] = 0
+    model.feed_data(data)
+    # check data shape
+    assert model.lq.shape == (1, 3, 8, 8)
+    assert model.gt.shape == (1, 3, 32, 32)
+
+    # ----------------- test nondist_validation -------------------- #
+    # construct dataloader
+    dataset_opt = dict(
+        name='Demo',
+        dataroot_gt='tests/data/gt',
+        dataroot_lq='tests/data/lq',
+        io_backend=dict(type='disk'),
+        scale=4,
+        phase='val')
+    dataset = PairedImageDataset(dataset_opt)
+    dataloader = torch.utils.data.DataLoader(dataset=dataset, batch_size=1, shuffle=False, num_workers=0)
+    assert model.is_train is True
+    model.nondist_validation(dataloader, 1, None, False)
+    assert model.is_train is True
+
+    # ----------------- test optimize_parameters -------------------- #
+    model.feed_data(data)
+    model.optimize_parameters(1)
+    assert model.output.shape == (1, 3, 32, 32)
+    assert isinstance(model.log_dict, dict)
+    # check returned keys
+    expected_keys = ['l_g_pix', 'l_g_percep', 'l_g_gan', 'l_d_real', 'out_d_real', 'l_d_fake', 'out_d_fake']
+    assert set(expected_keys).issubset(set(model.log_dict.keys()))

tests/test_utils.py

@@ -0,0 +1,87 @@
+import numpy as np
+from basicsr.archs.rrdbnet_arch import RRDBNet
+
+from realesrgan.utils import RealESRGANer
+
+
+def test_realesrganer():
+    # initialize with default model
+    restorer = RealESRGANer(
+        scale=4,
+        model_path='experiments/pretrained_models/RealESRGAN_x4plus.pth',
+        model=None,
+        tile=10,
+        tile_pad=10,
+        pre_pad=2,
+        half=False)
+    assert isinstance(restorer.model, RRDBNet)
+    assert restorer.half is False
+    # initialize with user-defined model
+    model = RRDBNet(num_in_ch=3, num_out_ch=3, num_feat=64, num_block=6, num_grow_ch=32, scale=4)
+    restorer = RealESRGANer(
+        scale=4,
+        model_path='experiments/pretrained_models/RealESRGAN_x4plus_anime_6B.pth',
+        model=model,
+        tile=10,
+        tile_pad=10,
+        pre_pad=2,
+        half=True)
+    # test attribute
+    assert isinstance(restorer.model, RRDBNet)
+    assert restorer.half is True
+
+    # ------------------ test pre_process ---------------- #
+    img = np.random.random((12, 12, 3)).astype(np.float32)
+    restorer.pre_process(img)
+    assert restorer.img.shape == (1, 3, 14, 14)
+    # with modcrop
+    restorer.scale = 1
+    restorer.pre_process(img)
+    assert restorer.img.shape == (1, 3, 16, 16)
+
+    # ------------------ test process ---------------- #
+    restorer.process()
+    assert restorer.output.shape == (1, 3, 64, 64)
+
+    # ------------------ test post_process ---------------- #
+    restorer.mod_scale = 4
+    output = restorer.post_process()
+    assert output.shape == (1, 3, 60, 60)
+
+    # ------------------ test tile_process ---------------- #
+    restorer.scale = 4
+    img = np.random.random((12, 12, 3)).astype(np.float32)
+    restorer.pre_process(img)
+    restorer.tile_process()
+    assert restorer.output.shape == (1, 3, 64, 64)
+
+    # ------------------ test enhance ---------------- #
+    img = np.random.random((12, 12, 3)).astype(np.float32)
+    result = restorer.enhance(img, outscale=2)
+    assert result[0].shape == (24, 24, 3)
+    assert result[1] == 'RGB'
+
+    # ------------------ test enhance with 16-bit image---------------- #
+    img = np.random.random((4, 4, 3)).astype(np.uint16) + 512
+    result = restorer.enhance(img, outscale=2)
+    assert result[0].shape == (8, 8, 3)
+    assert result[1] == 'RGB'
+
+    # ------------------ test enhance with gray image---------------- #
+    img = np.random.random((4, 4)).astype(np.float32)
+    result = restorer.enhance(img, outscale=2)
+    assert result[0].shape == (8, 8)
+    assert result[1] == 'L'
+
+    # ------------------ test enhance with RGBA---------------- #
+    img = np.random.random((4, 4, 4)).astype(np.float32)
+    result = restorer.enhance(img, outscale=2)
+    assert result[0].shape == (8, 8, 4)
+    assert result[1] == 'RGBA'
+
+    # ------------------ test enhance with RGBA, alpha_upsampler---------------- #
+    restorer.tile_size = 0
+    img = np.random.random((4, 4, 4)).astype(np.float32)
+    result = restorer.enhance(img, outscale=2, alpha_upsampler=None)
+    assert result[0].shape == (8, 8, 4)
+    assert result[1] == 'RGBA'