|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
import cv2 |
|
|
import numpy as np |
|
|
import os |
|
|
import sys |
|
|
from multiprocessing import Pool |
|
|
from os import path as osp |
|
|
from tqdm import tqdm |
|
|
|
|
|
from basicsr.utils import scandir_SIDD |
|
|
from basicsr.utils.create_lmdb import create_lmdb_for_SIDD |
|
|
|
|
|
|
|
|
def main(): |
|
|
opt = {} |
|
|
opt['n_thread'] = 20 |
|
|
opt['compression_level'] = 3 |
|
|
|
|
|
opt['input_folder'] = './datasets/SIDD/Data' |
|
|
opt['save_folder'] = './datasets/SIDD/train/input_crops' |
|
|
opt['crop_size'] = 512 |
|
|
opt['step'] = 384 |
|
|
opt['thresh_size'] = 0 |
|
|
opt['keywords'] = '_NOISY' |
|
|
extract_subimages(opt) |
|
|
|
|
|
opt['save_folder'] = './datasets/SIDD/train/gt_crops' |
|
|
opt['keywords'] = '_GT' |
|
|
extract_subimages(opt) |
|
|
|
|
|
create_lmdb_for_SIDD() |
|
|
|
|
|
|
|
|
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.') |
|
|
|
|
|
|
|
|
img_list = list(scandir_SIDD(input_folder, keywords=opt['keywords'], recursive=True, 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)) |
|
|
|
|
|
img_name = img_name.replace(opt['keywords'], '') |
|
|
|
|
|
img = cv2.imread(path, cv2.IMREAD_UNCHANGED) |
|
|
|
|
|
if img.ndim == 2: |
|
|
h, w = img.shape |
|
|
elif img.ndim == 3: |
|
|
h, w, c = img.shape |
|
|
else: |
|
|
raise ValueError(f'Image ndim should be 2 or 3, but got {img.ndim}') |
|
|
|
|
|
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__': |
|
|
main() |
|
|
|
