Real-ESRGAN

1
import argparse
2
import cv2
3
import numpy as np
4
import os
5
import sys
6
from basicsr.utils import scandir
7
from multiprocessing import Pool
8
from os import path as osp
9
from tqdm import tqdm
10

11

12
def main(args):
13
    """A multi-thread tool to crop large images to sub-images for faster IO.
14

15
    opt (dict): Configuration dict. It contains:
16
        n_thread (int): Thread number.
17
        compression_level (int):  CV_IMWRITE_PNG_COMPRESSION from 0 to 9. A higher value means a smaller size
18
            and longer compression time. Use 0 for faster CPU decompression. Default: 3, same in cv2.
19
        input_folder (str): Path to the input folder.
20
        save_folder (str): Path to save folder.
21
        crop_size (int): Crop size.
22
        step (int): Step for overlapped sliding window.
23
        thresh_size (int): Threshold size. Patches whose size is lower than thresh_size will be dropped.
24

25
    Usage:
26
        For each folder, run this script.
27
        Typically, there are GT folder and LQ folder to be processed for DIV2K dataset.
28
        After process, each sub_folder should have the same number of subimages.
29
        Remember to modify opt configurations according to your settings.
30
    """
31

32
    opt = {}
33
    opt['n_thread'] = args.n_thread
34
    opt['compression_level'] = args.compression_level
35
    opt['input_folder'] = args.input
36
    opt['save_folder'] = args.output
37
    opt['crop_size'] = args.crop_size
38
    opt['step'] = args.step
39
    opt['thresh_size'] = args.thresh_size
40
    extract_subimages(opt)
41

42

43
def extract_subimages(opt):
44
    """Crop images to subimages.
45

46
    Args:
47
        opt (dict): Configuration dict. It contains:
48
            input_folder (str): Path to the input folder.
49
            save_folder (str): Path to save folder.
50
            n_thread (int): Thread number.
51
    """
52
    input_folder = opt['input_folder']
53
    save_folder = opt['save_folder']
54
    if not osp.exists(save_folder):
55
        os.makedirs(save_folder)
56
        print(f'mkdir {save_folder} ...')
57
    else:
58
        print(f'Folder {save_folder} already exists. Exit.')
59
        sys.exit(1)
60

61
    # scan all images
62
    img_list = list(scandir(input_folder, full_path=True))
63

64
    pbar = tqdm(total=len(img_list), unit='image', desc='Extract')
65
    pool = Pool(opt['n_thread'])
66
    for path in img_list:
67
        pool.apply_async(worker, args=(path, opt), callback=lambda arg: pbar.update(1))
68
    pool.close()
69
    pool.join()
70
    pbar.close()
71
    print('All processes done.')
72

73

74
def worker(path, opt):
75
    """Worker for each process.
76

77
    Args:
78
        path (str): Image path.
79
        opt (dict): Configuration dict. It contains:
80
            crop_size (int): Crop size.
81
            step (int): Step for overlapped sliding window.
82
            thresh_size (int): Threshold size. Patches whose size is lower than thresh_size will be dropped.
83
            save_folder (str): Path to save folder.
84
            compression_level (int): for cv2.IMWRITE_PNG_COMPRESSION.
85

86
    Returns:
87
        process_info (str): Process information displayed in progress bar.
88
    """
89
    crop_size = opt['crop_size']
90
    step = opt['step']
91
    thresh_size = opt['thresh_size']
92
    img_name, extension = osp.splitext(osp.basename(path))
93

94
    # remove the x2, x3, x4 and x8 in the filename for DIV2K
95
    img_name = img_name.replace('x2', '').replace('x3', '').replace('x4', '').replace('x8', '')
96

97
    img = cv2.imread(path, cv2.IMREAD_UNCHANGED)
98

99
    h, w = img.shape[0:2]
100
    h_space = np.arange(0, h - crop_size + 1, step)
101
    if h - (h_space[-1] + crop_size) > thresh_size:
102
        h_space = np.append(h_space, h - crop_size)
103
    w_space = np.arange(0, w - crop_size + 1, step)
104
    if w - (w_space[-1] + crop_size) > thresh_size:
105
        w_space = np.append(w_space, w - crop_size)
106

107
    index = 0
108
    for x in h_space:
109
        for y in w_space:
110
            index += 1
111
            cropped_img = img[x:x + crop_size, y:y + crop_size, ...]
112
            cropped_img = np.ascontiguousarray(cropped_img)
113
            cv2.imwrite(
114
                osp.join(opt['save_folder'], f'{img_name}_s{index:03d}{extension}'), cropped_img,
115
                [cv2.IMWRITE_PNG_COMPRESSION, opt['compression_level']])
116
    process_info = f'Processing {img_name} ...'
117
    return process_info
118

119

120
if __name__ == '__main__':
121
    parser = argparse.ArgumentParser()
122
    parser.add_argument('--input', type=str, default='datasets/DF2K/DF2K_HR', help='Input folder')
123
    parser.add_argument('--output', type=str, default='datasets/DF2K/DF2K_HR_sub', help='Output folder')
124
    parser.add_argument('--crop_size', type=int, default=480, help='Crop size')
125
    parser.add_argument('--step', type=int, default=240, help='Step for overlapped sliding window')
126
    parser.add_argument(
127
        '--thresh_size',
128
        type=int,
129
        default=0,
130
        help='Threshold size. Patches whose size is lower than thresh_size will be dropped.')
131
    parser.add_argument('--n_thread', type=int, default=20, help='Thread number.')
132
    parser.add_argument('--compression_level', type=int, default=3, help='Compression level')
133
    args = parser.parse_args()
134

135
    main(args)
136