File size: 2,577 Bytes
4730cdc |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 |
#!/usr/bin/env python
# -*- coding:utf-8 -*-
# Power by Zongsheng Yue 2023-03-10 11:15:36
import sys
from pathlib import Path
sys.path.append(str(Path(__file__).parents[1]))
import os
import torch
import random
import argparse
import numpy as np
from omegaconf import OmegaConf
from basicsr.data.realesrgan_dataset import RealESRGANDataset
from utils import util_image
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",
)
args = parser.parse_args()
img_list = sorted([x for x in Path(args.indir).glob('**/*.JPEG')])
print(f'Number of images in imagenet validation dataset: {len(img_list)}')
random.seed(10000)
random.shuffle(img_list)
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)
num_imgs = 3000
configs = OmegaConf.load('./configs/degradation_testing_realesrgan.yaml')
opts, opts_degradation = configs.opts, configs.degradation
opts['dir_paths'] = [args.indir, ]
opts['length'] = num_imgs
dataset = RealESRGANDataset(opts, mode='testing')
for ii in range(num_imgs):
data_dict1 = dataset.__getitem__(ii)
if (ii + 1) % 100 == 0:
print(f'Processing: {ii+1}/{num_imgs}')
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()
|