first commit

scripts/prepare_sr_testing_syn.py

@@ -0,0 +1,107 @@
+#!/usr/bin/env python
+# -*- coding:utf-8 -*-
+# Power by Zongsheng Yue 2024-04-07 20:57:36
+
+import os
+import torch
+import random
+import argparse
+import numpy as np
+from omegaconf import OmegaConf
+
+import sys
+from pathlib import Path
+sys.path.append(str(Path(__file__).parents[1]))
+
+from basicsr.data.realesrgan_dataset import RealESRGANDataset
+from utils import util_image
+from utils import util_common
+
+def main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+            "-i",
+            "--indir",
+            type=str,
+            default="/mnt/lustre/share/zhangwenwei/data/imagenet/val",
+            help="Folder to save the checkpoints and training log",
+            )
+    parser.add_argument(
+            "-o",
+            "--outdir",
+            type=str,
+            default="./ImageNet-Test",
+            help="Folder to save the checkpoints and training log",
+            )
+    parser.add_argument(
+            "-r",
+            "--resolution",
+            type=int,
+            default=1024,
+            help="Image resolution of the ground truth",
+            )
+    parser.add_argument(
+            "--num_imgs",
+            type=int,
+            default=-1,
+            help="Number of images.",
+            )
+    args = parser.parse_args()
+
+    if Path(args.indir).is_dir():
+        img_list = sorted([x for x in Path(args.indir).glob('*.[JjPp][PpNn]*[Gg]')])
+    elif args.indir.endswith('txt'):
+        img_list = util_common.readline_txt(args.indir)
+    else:
+        raise ValueError('Please input valid args.indir!')
+    print(f'Number of images in the input folder: {len(img_list)}')
+
+    random.seed(10000)
+    random.shuffle(img_list)
+
+    num_imgs = args.num_imgs
+    if num_imgs > 0:
+        assert num_imgs <= len(img_list)
+        img_list = random.sample(img_list, k=num_imgs)
+
+    gt_dir = Path(args.outdir) / 'gt'
+    if not gt_dir.exists():
+        gt_dir.mkdir(parents=True)
+    lq_dir = Path(args.outdir) / 'lq'
+    if not lq_dir.exists():
+        lq_dir.mkdir(parents=True)
+
+    # Loading configuration
+    configs = OmegaConf.load('./configs/degradation_testing_realesrgan.yaml')
+    opts, opts_degradation = configs.opts, configs.degradation
+    opts['gt_size'] = args.resolution
+    opts_degradation['gt_size'] = args.resolution
+
+    dataset = RealESRGANDataset(opts, mode='testing')
+    dataset.image_paths = img_list
+    dataset.text_paths = [None, ] * len(img_list)
+    dataset.moment_paths = [None, ] * len(img_list)
+    for ii in range(len(img_list)):
+        data_dict1 = dataset.__getitem__(ii)
+        if (ii + 1) % 100 == 0:
+            print(f'Processing: {ii+1}/{len(img_list)}')
+        prefix = 'realesrgan'
+        data_dict2 = dataset.degrade_fun(
+                opts_degradation,
+                im_gt=data_dict1['gt'].unsqueeze(0),
+                kernel1=data_dict1['kernel1'],
+                kernel2=data_dict1['kernel2'],
+                sinc_kernel=data_dict1['sinc_kernel'],
+                )
+        im_lq, im_gt = data_dict2['lq'], data_dict2['gt']
+        im_lq, im_gt = util_image.tensor2img([im_lq, im_gt], rgb2bgr=True, min_max=(0,1) ) # uint8
+
+        im_name = Path(data_dict1['gt_path']).stem
+        im_path_gt = gt_dir / f'{im_name}.png'
+        util_image.imwrite(im_gt, im_path_gt, chn='bgr', dtype_in='uint8')
+
+        im_path_lq = lq_dir / f'{im_name}.png'
+        util_image.imwrite(im_lq, im_path_lq, chn='bgr', dtype_in='uint8')
+
+if __name__ == "__main__":
+    main()