import argparse from os import path as osp from utils.utils_video import scandir from utils.utils_lmdb import make_lmdb_from_imgs def create_lmdb_for_div2k(): """Create lmdb files for DIV2K dataset. Usage: Before run this script, please run `extract_subimages.py`. Typically, there are four folders to be processed for DIV2K dataset. DIV2K_train_HR_sub DIV2K_train_LR_bicubic/X2_sub DIV2K_train_LR_bicubic/X3_sub DIV2K_train_LR_bicubic/X4_sub Remember to modify opt configurations according to your settings. """ # HR images folder_path = 'trainsets/DIV2K/DIV2K_train_HR_sub' lmdb_path = 'trainsets/DIV2K/DIV2K_train_HR_sub.lmdb' img_path_list, keys = prepare_keys_div2k(folder_path) make_lmdb_from_imgs(folder_path, lmdb_path, img_path_list, keys) # LRx2 images folder_path = 'trainsets/DIV2K/DIV2K_train_LR_bicubic/X2_sub' lmdb_path = 'trainsets/DIV2K/DIV2K_train_LR_bicubic_X2_sub.lmdb' img_path_list, keys = prepare_keys_div2k(folder_path) make_lmdb_from_imgs(folder_path, lmdb_path, img_path_list, keys) # LRx3 images folder_path = 'trainsets/DIV2K/DIV2K_train_LR_bicubic/X3_sub' lmdb_path = 'trainsets/DIV2K/DIV2K_train_LR_bicubic_X3_sub.lmdb' img_path_list, keys = prepare_keys_div2k(folder_path) make_lmdb_from_imgs(folder_path, lmdb_path, img_path_list, keys) # LRx4 images folder_path = 'trainsets/DIV2K/DIV2K_train_LR_bicubic/X4_sub' lmdb_path = 'trainsets/DIV2K/DIV2K_train_LR_bicubic_X4_sub.lmdb' img_path_list, keys = prepare_keys_div2k(folder_path) make_lmdb_from_imgs(folder_path, lmdb_path, img_path_list, keys) def prepare_keys_div2k(folder_path): """Prepare image path list and keys for DIV2K dataset. Args: folder_path (str): Folder path. Returns: list[str]: Image path list. list[str]: Key list. """ print('Reading image path list ...') img_path_list = sorted(list(scandir(folder_path, suffix='png', recursive=False))) keys = [img_path.split('.png')[0] for img_path in sorted(img_path_list)] return img_path_list, keys def create_lmdb_for_reds(): """Create lmdb files for REDS dataset. Usage: Before run this script, please run `regroup_reds_dataset.py`. We take three folders for example: train_sharp train_sharp_bicubic train_blur (for video deblurring) Remember to modify opt configurations according to your settings. """ # train_sharp folder_path = 'trainsets/REDS/train_sharp' lmdb_path = 'trainsets/REDS/train_sharp_with_val.lmdb' img_path_list, keys = prepare_keys_reds(folder_path) make_lmdb_from_imgs(folder_path, lmdb_path, img_path_list, keys, multiprocessing_read=True) # train_sharp_bicubic folder_path = 'trainsets/REDS/train_sharp_bicubic' lmdb_path = 'trainsets/REDS/train_sharp_bicubic_with_val.lmdb' img_path_list, keys = prepare_keys_reds(folder_path) make_lmdb_from_imgs(folder_path, lmdb_path, img_path_list, keys, multiprocessing_read=True) # train_blur (for video deblurring) folder_path = 'trainsets/REDS_blur/train_blur' lmdb_path = 'trainsets/REDS_blur/train_blur_with_val.lmdb' img_path_list, keys = prepare_keys_reds(folder_path) make_lmdb_from_imgs(folder_path, lmdb_path, img_path_list, keys, multiprocessing_read=True) # train_blur_bicubic (for video deblurring-sr) folder_path = 'trainsets/REDS_blur_bicubic/train_blur_bicubic' lmdb_path = 'trainsets/REDS_blur_bicubic/train_blur_bicubic_with_val.lmdb' img_path_list, keys = prepare_keys_reds(folder_path) make_lmdb_from_imgs(folder_path, lmdb_path, img_path_list, keys, multiprocessing_read=True) def prepare_keys_reds(folder_path): """Prepare image path list and keys for REDS dataset. Args: folder_path (str): Folder path. Returns: list[str]: Image path list. list[str]: Key list. """ print('Reading image path list ...') img_path_list = sorted(list(scandir(folder_path, suffix='png', recursive=True))) keys = [v.split('.png')[0] for v in img_path_list] # example: 000/00000000 return img_path_list, keys def create_lmdb_for_vimeo90k(): """Create lmdb files for Vimeo90K dataset. Usage: Remember to modify opt configurations according to your settings. """ # GT folder_path = 'trainsets/vimeo90k/vimeo_septuplet/sequences' lmdb_path = 'trainsets/vimeo90k/vimeo90k_train_GT_only4th.lmdb' train_list_path = 'trainsets/vimeo90k/vimeo_septuplet/sep_trainlist.txt' img_path_list, keys = prepare_keys_vimeo90k(folder_path, train_list_path, 'gt') make_lmdb_from_imgs(folder_path, lmdb_path, img_path_list, keys, multiprocessing_read=True) # LQ folder_path = 'trainsets/vimeo90k/vimeo_septuplet_matlabLRx4/sequences' lmdb_path = 'trainsets/vimeo90k/vimeo90k_train_LR7frames.lmdb' train_list_path = 'trainsets/vimeo90k/vimeo_septuplet/sep_trainlist.txt' img_path_list, keys = prepare_keys_vimeo90k(folder_path, train_list_path, 'lq') make_lmdb_from_imgs(folder_path, lmdb_path, img_path_list, keys, multiprocessing_read=True) def create_lmdb_for_vimeo90k_bd(): """Create lmdb files for Vimeo90K dataset (blur-downsampled lr only). Usage: Remember to modify opt configurations according to your settings. """ # LQ (blur-downsampled, BD) folder_path = 'trainsets/vimeo90k/vimeo_septuplet_BDLRx4/sequences' lmdb_path = 'trainsets/vimeo90k/vimeo90k_train_BDLR7frames.lmdb' train_list_path = 'trainsets/vimeo90k/vimeo_septuplet/sep_trainlist.txt' img_path_list, keys = prepare_keys_vimeo90k(folder_path, train_list_path, 'lq') make_lmdb_from_imgs(folder_path, lmdb_path, img_path_list, keys, multiprocessing_read=True) def prepare_keys_vimeo90k(folder_path, train_list_path, mode): """Prepare image path list and keys for Vimeo90K dataset. Args: folder_path (str): Folder path. train_list_path (str): Path to the official train list. mode (str): One of 'gt' or 'lq'. Returns: list[str]: Image path list. list[str]: Key list. """ print('Reading image path list ...') with open(train_list_path, 'r') as fin: train_list = [line.strip() for line in fin] img_path_list = [] keys = [] for line in train_list: folder, sub_folder = line.split('/') img_path_list.extend([osp.join(folder, sub_folder, f'im{j + 1}.png') for j in range(7)]) keys.extend([f'{folder}/{sub_folder}/im{j + 1}' for j in range(7)]) if mode == 'gt': print('Only keep the 4th frame for the gt mode.') img_path_list = [v for v in img_path_list if v.endswith('im4.png')] keys = [v for v in keys if v.endswith('/im4')] return img_path_list, keys def create_lmdb_for_dvd(): """Create lmdb files for DVD dataset. Usage: We take two folders for example: GT input Remember to modify opt configurations according to your settings. """ # train_sharp folder_path = 'trainsets/DVD/train_GT' lmdb_path = 'trainsets/DVD/train_GT.lmdb' img_path_list, keys = prepare_keys_dvd(folder_path) make_lmdb_from_imgs(folder_path, lmdb_path, img_path_list, keys, multiprocessing_read=True) # train_sharp_bicubic folder_path = 'trainsets/DVD/train_GT_blurred' lmdb_path = 'trainsets/DVD/train_GT_blurred.lmdb' img_path_list, keys = prepare_keys_dvd(folder_path) make_lmdb_from_imgs(folder_path, lmdb_path, img_path_list, keys, multiprocessing_read=True) def prepare_keys_dvd(folder_path): """Prepare image path list and keys for DVD dataset. Args: folder_path (str): Folder path. Returns: list[str]: Image path list. list[str]: Key list. """ print('Reading image path list ...') img_path_list = sorted(list(scandir(folder_path, suffix='jpg', recursive=True))) keys = [v.split('.jpg')[0] for v in img_path_list] # example: 000/00000000 return img_path_list, keys def create_lmdb_for_gopro(): """Create lmdb files for GoPro dataset. Usage: We take two folders for example: GT input Remember to modify opt configurations according to your settings. """ # train_sharp folder_path = 'trainsets/GoPro/train_GT' lmdb_path = 'trainsets/GoPro/train_GT.lmdb' img_path_list, keys = prepare_keys_gopro(folder_path) make_lmdb_from_imgs(folder_path, lmdb_path, img_path_list, keys, multiprocessing_read=True) # train_sharp_bicubic folder_path = 'trainsets/GoPro/train_GT_blurred' lmdb_path = 'trainsets/GoPro/train_GT_blurred.lmdb' img_path_list, keys = prepare_keys_gopro(folder_path) make_lmdb_from_imgs(folder_path, lmdb_path, img_path_list, keys, multiprocessing_read=True) def prepare_keys_gopro(folder_path): """Prepare image path list and keys for GoPro dataset. Args: folder_path (str): Folder path. Returns: list[str]: Image path list. list[str]: Key list. """ print('Reading image path list ...') img_path_list = sorted(list(scandir(folder_path, suffix='png', recursive=True))) keys = [v.split('.png')[0] for v in img_path_list] # example: 000/00000000 return img_path_list, keys def create_lmdb_for_davis(): """Create lmdb files for DAVIS dataset. Usage: We take one folders for example: GT Remember to modify opt configurations according to your settings. """ # train_sharp folder_path = 'trainsets/DAVIS/train_GT' lmdb_path = 'trainsets/DAVIS/train_GT.lmdb' img_path_list, keys = prepare_keys_davis(folder_path) make_lmdb_from_imgs(folder_path, lmdb_path, img_path_list, keys, multiprocessing_read=True) def prepare_keys_davis(folder_path): """Prepare image path list and keys for DAVIS dataset. Args: folder_path (str): Folder path. Returns: list[str]: Image path list. list[str]: Key list. """ print('Reading image path list ...') img_path_list = sorted(list(scandir(folder_path, suffix='jpg', recursive=True))) keys = [v.split('.jpg')[0] for v in img_path_list] # example: 000/00000000 return img_path_list, keys def create_lmdb_for_ldv(): """Create lmdb files for LDV dataset. Usage: We take two folders for example: GT input Remember to modify opt configurations according to your settings. """ # training_raw folder_path = 'trainsets/LDV/training_raw' lmdb_path = 'trainsets/LDV/training_raw.lmdb' img_path_list, keys = prepare_keys_ldv(folder_path) make_lmdb_from_imgs(folder_path, lmdb_path, img_path_list, keys, multiprocessing_read=True) # training_fixed-QP folder_path = 'trainsets/LDV/training_fixed-QP' lmdb_path = 'trainsets/LDV/training_fixed-QP.lmdb' img_path_list, keys = prepare_keys_ldv(folder_path) make_lmdb_from_imgs(folder_path, lmdb_path, img_path_list, keys, multiprocessing_read=True) # training_fixed-rate folder_path = 'trainsets/LDV/training_fixed-rate' lmdb_path = 'trainsets/LDV/training_fixed-rate.lmdb' img_path_list, keys = prepare_keys_ldv(folder_path) make_lmdb_from_imgs(folder_path, lmdb_path, img_path_list, keys, multiprocessing_read=True) def prepare_keys_ldv(folder_path): """Prepare image path list and keys for LDV dataset. Args: folder_path (str): Folder path. Returns: list[str]: Image path list. list[str]: Key list. """ print('Reading image path list ...') img_path_list = sorted(list(scandir(folder_path, suffix='png', recursive=True))) keys = [v.split('.png')[0] for v in img_path_list] # example: 000/00000000 return img_path_list, keys def create_lmdb_for_reds_orig(): """Create lmdb files for REDS_orig dataset (120 fps). Usage: Before run this script, please run `regroup_reds_dataset.py`. We take one folders for example: train_orig Remember to modify opt configurations according to your settings. """ # train_sharp folder_path = 'trainsets/REDS_orig/train_orig' lmdb_path = 'trainsets/REDS_orig/train_orig_with_val.lmdb' img_path_list, keys = prepare_keys_reds_orig(folder_path) make_lmdb_from_imgs(folder_path, lmdb_path, img_path_list, keys, multiprocessing_read=True) def prepare_keys_reds_orig(folder_path): """Prepare image path list and keys for REDS_orig dataset (120 fps). Args: folder_path (str): Folder path. Returns: list[str]: Image path list. list[str]: Key list. """ print('Reading image path list ...') img_path_list = sorted(list(scandir(folder_path, suffix='png', recursive=True))) keys = [v.split('.png')[0] for v in img_path_list] # example: 000/00000000 return img_path_list, keys if __name__ == '__main__': parser = argparse.ArgumentParser() parser.add_argument( '--dataset', type=str, help=("Options: 'DIV2K', 'REDS', 'Vimeo90K', 'Vimeo90K_BD', 'DVD', 'GoPro'," "'DAVIS', 'LDV', 'REDS_orig' " 'You may need to modify the corresponding configurations in codes.')) args = parser.parse_args() dataset = args.dataset.lower() if dataset == 'div2k': create_lmdb_for_div2k() elif dataset == 'reds': create_lmdb_for_reds() elif dataset == 'vimeo90k': create_lmdb_for_vimeo90k() elif dataset == 'vimeo90k_bd': create_lmdb_for_vimeo90k_bd() elif dataset == 'dvd': create_lmdb_for_dvd() elif dataset == 'gopro': create_lmdb_for_gopro() elif dataset == 'davis': create_lmdb_for_davis() elif dataset == 'ldv': create_lmdb_for_ldv() elif dataset == 'reds_orig': create_lmdb_for_reds_orig() else: raise ValueError('Wrong dataset.')