Python torch.utils.data.util() Examples
The following are 27
code examples of torch.utils.data.util().
You can vote up the ones you like or vote down the ones you don't like,
and go to the original project or source file by following the links above each example.
You may also want to check out all available functions/classes of the module
torch.utils.data
, or try the search function
.
.
Example #1
| Source File: video_test_dataset.py From EDVR with Apache License 2.0 | 6 votes |
|
def __getitem__(self, index):
# path_LQ = self.data_info['path_LQ'][index]
# path_GT = self.data_info['path_GT'][index]
folder = self.data_info['folder'][index]
idx, max_idx = self.data_info['idx'][index].split('/')
idx, max_idx = int(idx), int(max_idx)
border = self.data_info['border'][index]
if self.cache_data:
select_idx = util.index_generation(idx, max_idx, self.opt['N_frames'],
padding=self.opt['padding'])
imgs_LQ = self.imgs_LQ[folder].index_select(0, torch.LongTensor(select_idx))
img_GT = self.imgs_GT[folder][idx]
else:
pass # TODO
return {
'LQs': imgs_LQ,
'GT': img_GT,
'folder': folder,
'idx': self.data_info['idx'][index],
'border': border
}
Example #2
| Source File: LR_dataset.py From BasicSR with Apache License 2.0 | 6 votes |
|
def __getitem__(self, index):
LR_path = None
# get LR image
LR_path = self.paths_LR[index]
img_LR = util.read_img(self.LR_env, LR_path)
H, W, C = img_LR.shape
# change color space if necessary
if self.opt['color']:
img_LR = util.channel_convert(C, self.opt['color'], [img_LR])[0]
# BGR to RGB, HWC to CHW, numpy to tensor
if img_LR.shape[2] == 3:
img_LR = img_LR[:, :, [2, 1, 0]]
img_LR = torch.from_numpy(np.ascontiguousarray(np.transpose(img_LR, (2, 0, 1)))).float()
return {'LR': img_LR, 'LR_path': LR_path}
Example #3
| Source File: LQGT_dataset.py From BasicSR with Apache License 2.0 | 6 votes |
|
def __init__(self, opt):
super(LQGTDataset, self).__init__()
self.opt = opt
self.data_type = self.opt['data_type']
self.paths_LQ, self.paths_GT = None, None
self.sizes_LQ, self.sizes_GT = None, None
self.LQ_env, self.GT_env = None, None # environment for lmdb
self.paths_GT, self.sizes_GT = util.get_image_paths(self.data_type, opt['dataroot_GT'])
self.paths_LQ, self.sizes_LQ = util.get_image_paths(self.data_type, opt['dataroot_LQ'])
assert self.paths_GT, 'Error: GT path is empty.'
if self.paths_LQ and self.paths_GT:
assert len(self.paths_LQ) == len(
self.paths_GT
), 'GT and LQ datasets have different number of images - {}, {}.'.format(
len(self.paths_LQ), len(self.paths_GT))
self.random_scale_list = [1]
Example #4
| Source File: LRHR_seg_bg_dataset.py From real-world-sr with MIT License | 6 votes |
|
def __init__(self, opt):
super(LRHRSeg_BG_Dataset, self).__init__()
self.opt = opt
self.paths_LR = None
self.paths_HR = None
self.paths_HR_bg = None # HR images for background scenes
self.LR_env = None # environment for lmdb
self.HR_env = None
self.HR_env_bg = None
# read image list from lmdb or image files
self.HR_env, self.paths_HR = util.get_image_paths(opt['data_type'], opt['dataroot_GT'])
self.LR_env, self.paths_LR = util.get_image_paths(opt['data_type'], opt['dataroot_LR'])
self.HR_env_bg, self.paths_HR_bg = util.get_image_paths(opt['data_type'],
opt['dataroot_GT_bg'])
assert self.paths_HR, 'Error: HR path is empty.'
if self.paths_LR and self.paths_HR:
assert len(self.paths_LR) == len(self.paths_HR), \
'HR and LR datasets have different number of images - {}, {}.'.format(
len(self.paths_LR), len(self.paths_HR))
self.random_scale_list = [1, 0.9, 0.8, 0.7, 0.6, 0.5]
self.ratio = 10 # 10 OST data samples and 1 DIV2K general data samples(background)
Example #5
| Source File: video_test_dataset.py From mmsr with Apache License 2.0 | 6 votes |
|
def __getitem__(self, index):
# path_LQ = self.data_info['path_LQ'][index]
# path_GT = self.data_info['path_GT'][index]
folder = self.data_info['folder'][index]
idx, max_idx = self.data_info['idx'][index].split('/')
idx, max_idx = int(idx), int(max_idx)
border = self.data_info['border'][index]
if self.cache_data:
select_idx = util.index_generation(idx, max_idx, self.opt['N_frames'],
padding=self.opt['padding'])
imgs_LQ = self.imgs_LQ[folder].index_select(0, torch.LongTensor(select_idx))
img_GT = self.imgs_GT[folder][idx]
else:
pass # TODO
return {
'LQs': imgs_LQ,
'GT': img_GT,
'folder': folder,
'idx': self.data_info['idx'][index],
'border': border
}
Example #6
| Source File: LR_dataset.py From real-world-sr with MIT License | 6 votes |
|
def __getitem__(self, index):
LR_path = None
# get LR image
LR_path = self.paths_LR[index]
img_LR = util.read_img(self.LR_env, LR_path)
H, W, C = img_LR.shape
# change color space if necessary
if self.opt['color']:
img_LR = util.channel_convert(C, self.opt['color'], [img_LR])[0]
# BGR to RGB, HWC to CHW, numpy to tensor
if img_LR.shape[2] == 3:
img_LR = img_LR[:, :, [2, 1, 0]]
img_LR = torch.from_numpy(np.ascontiguousarray(np.transpose(img_LR, (2, 0, 1)))).float()
return {'LR': img_LR, 'LR_path': LR_path}
Example #7
| Source File: LQGT_dataset.py From real-world-sr with MIT License | 6 votes |
|
def __init__(self, opt):
super(LQGTDataset, self).__init__()
self.opt = opt
self.data_type = self.opt['data_type']
self.paths_LQ, self.paths_GT = None, None
self.sizes_LQ, self.sizes_GT = None, None
self.LQ_env, self.GT_env = None, None # environment for lmdb
self.paths_GT, self.sizes_GT = util.get_image_paths(self.data_type, opt['dataroot_GT'])
self.paths_LQ, self.sizes_LQ = util.get_image_paths(self.data_type, opt['dataroot_LQ'])
assert self.paths_GT, 'Error: GT path is empty.'
if self.paths_LQ and self.paths_GT:
assert len(self.paths_LQ) == len(
self.paths_GT
), 'GT and LQ datasets have different number of images - {}, {}.'.format(
len(self.paths_LQ), len(self.paths_GT))
self.random_scale_list = [1]
Example #8
| Source File: LQ_dataset.py From IKC with Apache License 2.0 | 6 votes |
|
def __init__(self, opt):
super(LQDataset, self).__init__()
self.opt = opt
self.opt_P = opt
self.LR_paths = None
self.LR_sizes = None # environment for lmdb
self.LR_env = None
self.LR_size = opt['LR_size']
self.real_ker_path = '/mnt/yjchai/SR_data/Flickr2K/kermap.pt'
self.real_ker_map_list = util.load_ker_map_list(self.real_ker_path)
# read image list from lmdb or image files
if opt['data_type'] == 'lmdb':
self.LR_paths, self.LR_sizes = util.get_image_paths(opt['data_type'], opt['dataroot_LQ'])
elif opt['data_type'] == 'img':
self.LR_paths = util.get_image_paths(opt['data_type'], opt['dataroot_LQ']) #LR_list
else:
print('Error: data_type is not matched in Dataset')
assert self.LR_paths, 'Error: LR paths are empty.'
Example #9
| Source File: LQ_dataset.py From EDVR with Apache License 2.0 | 6 votes |
|
def __getitem__(self, index):
if self.data_type == 'lmdb' and self.LQ_env is None:
self._init_lmdb()
LQ_path = None
# get LQ image
LQ_path = self.LQ_path[index]
resolution = [int(s) for s in self.sizes_LQ[index].split('_')
] if self.data_type == 'lmdb' else None
img_LQ = util.read_img(self.LQ_env, LQ_path, resolution)
H, W, C = img_LQ.shape
if self.opt['color']: # change color space if necessary
img_LQ = util.channel_convert(C, self.opt['color'], [img_LQ])[0]
# BGR to RGB, HWC to CHW, numpy to tensor
if img_LQ.shape[2] == 3:
img_LQ = img_LQ[:, :, [2, 1, 0]]
img_LQ = torch.from_numpy(np.ascontiguousarray(np.transpose(img_LQ, (2, 0, 1)))).float()
return {'LQ': img_LQ, 'LQ_path': LQ_path}
Example #10
| Source File: LRHR_seg_bg_dataset.py From IKC with Apache License 2.0 | 6 votes |
|
def __init__(self, opt):
super(LRHRSeg_BG_Dataset, self).__init__()
self.opt = opt
self.paths_LR = None
self.paths_HR = None
self.paths_HR_bg = None # HR images for background scenes
self.LR_env = None # environment for lmdb
self.HR_env = None
self.HR_env_bg = None
# read image list from lmdb or image files
self.HR_env, self.paths_HR = util.get_image_paths(opt['data_type'], opt['dataroot_GT'])
self.LR_env, self.paths_LR = util.get_image_paths(opt['data_type'], opt['dataroot_LR'])
self.HR_env_bg, self.paths_HR_bg = util.get_image_paths(opt['data_type'],
opt['dataroot_GT_bg'])
assert self.paths_HR, 'Error: HR path is empty.'
if self.paths_LR and self.paths_HR:
assert len(self.paths_LR) == len(self.paths_HR), \
'HR and LR datasets have different number of images - {}, {}.'.format(
len(self.paths_LR), len(self.paths_HR))
self.random_scale_list = [1, 0.9, 0.8, 0.7, 0.6, 0.5]
self.ratio = 10 # 10 OST data samples and 1 DIV2K general data samples(background)
Example #11
| Source File: LQ_dataset.py From mmsr with Apache License 2.0 | 6 votes |
|
def __getitem__(self, index):
if self.data_type == 'lmdb' and self.LQ_env is None:
self._init_lmdb()
LQ_path = None
# get LQ image
LQ_path = self.paths_LQ[index]
resolution = [int(s) for s in self.sizes_LQ[index].split('_')
] if self.data_type == 'lmdb' else None
img_LQ = util.read_img(self.LQ_env, LQ_path, resolution)
H, W, C = img_LQ.shape
if self.opt['color']: # change color space if necessary
img_LQ = util.channel_convert(C, self.opt['color'], [img_LQ])[0]
# BGR to RGB, HWC to CHW, numpy to tensor
if img_LQ.shape[2] == 3:
img_LQ = img_LQ[:, :, [2, 1, 0]]
img_LQ = torch.from_numpy(np.ascontiguousarray(np.transpose(img_LQ, (2, 0, 1)))).float()
return {'LQ': img_LQ, 'LQ_path': LQ_path}
Example #12
| Source File: LRHR_seg_bg_dataset.py From BasicSR with Apache License 2.0 | 6 votes |
|
def __init__(self, opt):
super(LRHRSeg_BG_Dataset, self).__init__()
self.opt = opt
self.paths_LR = None
self.paths_HR = None
self.paths_HR_bg = None # HR images for background scenes
self.LR_env = None # environment for lmdb
self.HR_env = None
self.HR_env_bg = None
# read image list from lmdb or image files
self.HR_env, self.paths_HR = util.get_image_paths(opt['data_type'], opt['dataroot_GT'])
self.LR_env, self.paths_LR = util.get_image_paths(opt['data_type'], opt['dataroot_LR'])
self.HR_env_bg, self.paths_HR_bg = util.get_image_paths(opt['data_type'],
opt['dataroot_GT_bg'])
assert self.paths_HR, 'Error: HR path is empty.'
if self.paths_LR and self.paths_HR:
assert len(self.paths_LR) == len(self.paths_HR), \
'HR and LR datasets have different number of images - {}, {}.'.format(
len(self.paths_LR), len(self.paths_HR))
self.random_scale_list = [1, 0.9, 0.8, 0.7, 0.6, 0.5]
self.ratio = 10 # 10 OST data samples and 1 DIV2K general data samples(background)
Example #13
| Source File: LQGTker_dataset.py From IKC with Apache License 2.0 | 6 votes |
|
def __init__(self, opt):
super(LQGTKerDataset, self).__init__()
self.opt = opt
self.opt_F = opt
self.opt_P = opt
self.opt_C = opt
self.LR_paths, self.GT_paths = None, None
self.LR_env, self.GT_env = None, None # environment for lmdb
self.LR_size, self.GT_size = opt['LR_size'], opt['GT_size']
# read image list from lmdb or image files
if opt['data_type'] == 'lmdb':
self.LR_paths, self.LR_sizes = util.get_image_paths(opt['data_type'], opt['dataroot_LQ'])
self.GT_paths, self.GT_sizes = util.get_image_paths(opt['data_type'], opt['dataroot_GT'])
elif opt['data_type'] == 'img':
self.LR_paths = util.get_image_paths(opt['data_type'], opt['dataroot_LQ']) # LR list
self.GT_paths = util.get_image_paths(opt['data_type'], opt['dataroot_GT']) # GT list
else:
print('Error: data_type is not matched in Dataset')
assert self.GT_paths, 'Error: GT paths are empty.'
if self.LR_paths and self.GT_paths:
assert len(self.LR_paths) == len(self.GT_paths), 'GT and LR datasets have different number of images - {}, {}.'.format(len(self.LR_paths), len(self.GT_paths))
self.random_scale_list = [1]
Example #14
| Source File: video_test_dataset.py From mmsr with Apache License 2.0 | 5 votes |
|
def __init__(self, opt):
super(VideoTestDataset, self).__init__()
self.opt = opt
self.cache_data = opt['cache_data']
self.half_N_frames = opt['N_frames'] // 2
self.GT_root, self.LQ_root = opt['dataroot_GT'], opt['dataroot_LQ']
self.data_type = self.opt['data_type']
self.data_info = {'path_LQ': [], 'path_GT': [], 'folder': [], 'idx': [], 'border': []}
if self.data_type == 'lmdb':
raise ValueError('No need to use LMDB during validation/test.')
#### Generate data info and cache data
self.imgs_LQ, self.imgs_GT = {}, {}
if opt['name'].lower() in ['vid4', 'reds4']:
subfolders_LQ = util.glob_file_list(self.LQ_root)
subfolders_GT = util.glob_file_list(self.GT_root)
for subfolder_LQ, subfolder_GT in zip(subfolders_LQ, subfolders_GT):
subfolder_name = osp.basename(subfolder_GT)
img_paths_LQ = util.glob_file_list(subfolder_LQ)
img_paths_GT = util.glob_file_list(subfolder_GT)
max_idx = len(img_paths_LQ)
assert max_idx == len(
img_paths_GT), 'Different number of images in LQ and GT folders'
self.data_info['path_LQ'].extend(img_paths_LQ)
self.data_info['path_GT'].extend(img_paths_GT)
self.data_info['folder'].extend([subfolder_name] * max_idx)
for i in range(max_idx):
self.data_info['idx'].append('{}/{}'.format(i, max_idx))
border_l = [0] * max_idx
for i in range(self.half_N_frames):
border_l[i] = 1
border_l[max_idx - i - 1] = 1
self.data_info['border'].extend(border_l)
if self.cache_data:
self.imgs_LQ[subfolder_name] = util.read_img_seq(img_paths_LQ)
self.imgs_GT[subfolder_name] = util.read_img_seq(img_paths_GT)
elif opt['name'].lower() in ['vimeo90k-test']:
pass # TODO
else:
raise ValueError(
'Not support video test dataset. Support Vid4, REDS4 and Vimeo90k-Test.')
Example #15
| Source File: REDS_dataset.py From EDVR with Apache License 2.0 | 5 votes |
|
def __init__(self, opt):
super(REDSDataset, self).__init__()
self.opt = opt
# temporal augmentation
self.interval_list = opt['interval_list']
self.random_reverse = opt['random_reverse']
logger.info('Temporal augmentation interval list: [{}], with random reverse is {}.'.format(
','.join(str(x) for x in opt['interval_list']), self.random_reverse))
self.half_N_frames = opt['N_frames'] // 2
self.GT_root, self.LQ_root = opt['dataroot_GT'], opt['dataroot_LQ']
self.data_type = self.opt['data_type']
self.LR_input = False if opt['GT_size'] == opt['LQ_size'] else True # low resolution inputs
#### directly load image keys
if self.data_type == 'lmdb':
self.paths_GT, _ = util.get_image_paths(self.data_type, opt['dataroot_GT'])
logger.info('Using lmdb meta info for cache keys.')
elif opt['cache_keys']:
logger.info('Using cache keys: {}'.format(opt['cache_keys']))
self.paths_GT = pickle.load(open(opt['cache_keys'], 'rb'))['keys']
else:
raise ValueError(
'Need to create cache keys (meta_info.pkl) by running [create_lmdb.py]')
# remove the REDS4 for testing
self.paths_GT = [
v for v in self.paths_GT if v.split('_')[0] not in ['000', '011', '015', '020']
]
assert self.paths_GT, 'Error: GT path is empty.'
if self.data_type == 'lmdb':
self.GT_env, self.LQ_env = None, None
elif self.data_type == 'mc': # memcached
self.mclient = None
elif self.data_type == 'img':
pass
else:
raise ValueError('Wrong data type: {}'.format(self.data_type))
Example #16
| Source File: LQ_dataset.py From mmsr with Apache License 2.0 | 5 votes |
|
def __init__(self, opt):
super(LQDataset, self).__init__()
self.opt = opt
self.data_type = self.opt['data_type']
self.paths_LQ, self.paths_GT = None, None
self.LQ_env = None # environment for lmdb
self.paths_LQ, self.sizes_LQ = util.get_image_paths(self.data_type, opt['dataroot_LQ'])
assert self.paths_LQ, 'Error: LQ paths are empty.'
Example #17
| Source File: LQker_dataset.py From IKC with Apache License 2.0 | 5 votes |
|
def __getitem__(self, index):
if self.opt['data_type'] == 'lmdb':
if self.LR_env is None:
self._init_lmdb()
LR_size = self.LR_size
# get LR image, kernel map
LR_path = self.LR_paths[index]
ker_map = self.ker_maps[index]
if self.opt['data_type'] == 'lmdb':
resolution = [int(s) for s in self.LR_sizes[index].split('_')]
else:
resolution = None
img_LR = util.read_img(self.LR_env, LR_path, resolution)
H, W, C = img_LR.shape
if self.opt['phase'] == 'train':
#randomly crop
rnd_h = random.randint(0, max(0, H - LR_size))
rnd_w = random.randint(0, max(0, W - LR_size))
img_LR = img_LR[rnd_h:rnd_h + LR_size, rnd_w:rnd_w + LR_size, :]
# augmentation - flip, rotate
img_LR = util.augment(img_LR, self.opt['use_flip'], self.opt['use_rot'], self.opt['mode'])
# change color space if necessary
if self.opt['color']:
img_LR = util.channel_convert(C, self.opt['color'], [img_LR])[0]
# BGR to RGB, HWC to CHW, numpy to tensor
if img_LR.shape[2] == 3:
img_LR = img_LR[:, :, [2, 1, 0]]
img_LR = torch.from_numpy(np.ascontiguousarray(np.transpose(img_LR, (2, 0, 1)))).float()
return {'LQ': img_LR, 'ker': ker_map, 'LQ_path': LR_path}
Example #18
| Source File: REDS_dataset.py From mmsr with Apache License 2.0 | 5 votes |
|
def __init__(self, opt):
super(REDSDataset, self).__init__()
self.opt = opt
# temporal augmentation
self.interval_list = opt['interval_list']
self.random_reverse = opt['random_reverse']
logger.info('Temporal augmentation interval list: [{}], with random reverse is {}.'.format(
','.join(str(x) for x in opt['interval_list']), self.random_reverse))
self.half_N_frames = opt['N_frames'] // 2
self.GT_root, self.LQ_root = opt['dataroot_GT'], opt['dataroot_LQ']
self.data_type = self.opt['data_type']
self.LR_input = False if opt['GT_size'] == opt['LQ_size'] else True # low resolution inputs
#### directly load image keys
if self.data_type == 'lmdb':
self.paths_GT, _ = util.get_image_paths(self.data_type, opt['dataroot_GT'])
logger.info('Using lmdb meta info for cache keys.')
elif opt['cache_keys']:
logger.info('Using cache keys: {}'.format(opt['cache_keys']))
self.paths_GT = pickle.load(open(opt['cache_keys'], 'rb'))['keys']
else:
raise ValueError(
'Need to create cache keys (meta_info.pkl) by running [create_lmdb.py]')
# remove the REDS4 for testing
self.paths_GT = [
v for v in self.paths_GT if v.split('_')[0] not in ['000', '011', '015', '020']
]
assert self.paths_GT, 'Error: GT path is empty.'
if self.data_type == 'lmdb':
self.GT_env, self.LQ_env = None, None
elif self.data_type == 'mc': # memcached
self.mclient = None
elif self.data_type == 'img':
pass
else:
raise ValueError('Wrong data type: {}'.format(self.data_type))
Example #19
| Source File: video_test_dataset.py From EDVR with Apache License 2.0 | 5 votes |
|
def __init__(self, opt):
super(VideoTestDataset, self).__init__()
self.opt = opt
self.cache_data = opt['cache_data']
self.half_N_frames = opt['N_frames'] // 2
self.GT_root, self.LQ_root = opt['dataroot_GT'], opt['dataroot_LQ']
self.data_type = self.opt['data_type']
self.data_info = {'path_LQ': [], 'path_GT': [], 'folder': [], 'idx': [], 'border': []}
if self.data_type == 'lmdb':
raise ValueError('No need to use LMDB during validation/test.')
#### Generate data info and cache data
self.imgs_LQ, self.imgs_GT = {}, {}
if opt['name'].lower() in ['vid4', 'reds4']:
subfolders_LQ = util.glob_file_list(self.LQ_root)
subfolders_GT = util.glob_file_list(self.GT_root)
for subfolder_LQ, subfolder_GT in zip(subfolders_LQ, subfolders_GT):
subfolder_name = osp.basename(subfolder_GT)
img_paths_LQ = util.glob_file_list(subfolder_LQ)
img_paths_GT = util.glob_file_list(subfolder_GT)
max_idx = len(img_paths_LQ)
assert max_idx == len(
img_paths_GT), 'Different number of images in LQ and GT folders'
self.data_info['path_LQ'].extend(img_paths_LQ)
self.data_info['path_GT'].extend(img_paths_GT)
self.data_info['folder'].extend([subfolder_name] * max_idx)
for i in range(max_idx):
self.data_info['idx'].append('{}/{}'.format(i, max_idx))
border_l = [0] * max_idx
for i in range(self.half_N_frames):
border_l[i] = 1
border_l[max_idx - i - 1] = 1
self.data_info['border'].extend(border_l)
if self.cache_data:
self.imgs_LQ[subfolder_name] = util.read_img_seq(img_paths_LQ)
self.imgs_GT[subfolder_name] = util.read_img_seq(img_paths_GT)
elif opt['name'].lower() in ['vimeo90k-test']:
pass # TODO
else:
raise ValueError(
'Not support video test dataset. Support Vid4, REDS4 and Vimeo90k-Test.')
Example #20
| Source File: LQ_dataset.py From EDVR with Apache License 2.0 | 5 votes |
|
def __init__(self, opt):
super(LQDataset, self).__init__()
self.opt = opt
self.paths_LQ, self.paths_GT = None, None
self.LQ_env = None # environment for lmdb
self.paths_LQ, self.sizes_LQ = util.get_image_paths(self.data_type, opt['dataroot_LQ'])
assert self.paths_LQ, 'Error: LQ paths are empty.'
Example #21
| Source File: SRker_dataset.py From IKC with Apache License 2.0 | 5 votes |
|
def __init__(self, opt, ker_map_list, SR_img_list):
super(SRKerDataset, self).__init__()
self.opt = opt
self.opt_C = opt
self.LR_paths = None
self.LR_sizes = None # environment for lmdb
self.LR_env = None
self.LR_size = opt['LR_size']
self.SR_env = None
self.SR_img_list = SR_img_list
self.SR_size = opt['GT_size']
self.ker_map_list = ker_map_list
self.real_ker_path = '/mnt/yjchai/SR_data/Flickr2K/kermap.pt'
self.real_ker_map_list = util.load_ker_map_list(self.real_ker_path)
# read image list from lmdb or image files
#if opt['data_type'] == 'lmdb':
# self.LR_paths, self.LR_sizes = util.get_image_paths(opt['data_type'], opt['dataroot_LQ'])
#elif opt['data_type'] == 'img':
# self.LR_paths = util.get_image_paths(opt['data_type'], opt['dataroot_LQ']) #LR_list
#else:
# print('Error: data_type is not matched in Dataset')
#assert self.LR_paths, 'Error: LR paths are empty.'
Example #22
| Source File: LQ_dataset.py From IKC with Apache License 2.0 | 5 votes |
|
def __getitem__(self, index):
if self.opt['data_type'] == 'lmdb':
if self.LR_env is None:
self._init_lmdb()
LR_size = self.LR_size
# get real kernel map
real_ker_map = self.real_ker_map_list[index].float()
# get LR image
LR_path = self.LR_paths[index]
if self.opt['data_type'] == 'lmdb':
resolution = [int(s) for s in self.LR_sizes[index].split('_')]
else:
resolution = None
img_LR = util.read_img(self.LR_env, LR_path, resolution)
H, W, C = img_LR.shape
if self.opt['phase'] == 'train':
#randomly crop
rnd_h = random.randint(0, max(0, H - LR_size))
rnd_w = random.randint(0, max(0, W - LR_size))
img_LR = img_LR[rnd_h:rnd_h + LR_size, rnd_w:rnd_w + LR_size, :]
# augmentation - flip, rotate
img_LR = util.augment(img_LR, self.opt['use_flip'], self.opt['use_rot'], self.opt['mode'])
# change color space if necessary
if self.opt['color']:
img_LR = util.channel_convert(C, self.opt['color'], [img_LR])[0]
# BGR to RGB, HWC to CHW, numpy to tensor
if img_LR.shape[2] == 3:
img_LR = img_LR[:, :, [2, 1, 0]]
img_LR = torch.from_numpy(np.ascontiguousarray(np.transpose(img_LR, (2, 0, 1)))).float()
return {'LQ': img_LR, 'LQ_path': LR_path, 'real_ker': real_ker_map}
Example #23
| Source File: LR_dataset.py From real-world-sr with MIT License | 5 votes |
|
def __init__(self, opt):
super(LRDataset, self).__init__()
self.opt = opt
self.paths_LR = None
self.LR_env = None # environment for lmdb
# read image list from lmdb or image files
self.LR_env, self.paths_LR = util.get_image_paths(opt['data_type'], opt['dataroot_LR'])
assert self.paths_LR, 'Error: LR paths are empty.'
Example #24
| Source File: LR_dataset.py From BasicSR with Apache License 2.0 | 5 votes |
|
def __init__(self, opt):
super(LRDataset, self).__init__()
self.opt = opt
self.paths_LR = None
self.LR_env = None # environment for lmdb
# read image list from lmdb or image files
self.LR_env, self.paths_LR = util.get_image_paths(opt['data_type'], opt['dataroot_LR'])
assert self.paths_LR, 'Error: LR paths are empty.'
Example #25
| Source File: Vimeo90K_dataset.py From EDVR with Apache License 2.0 | 4 votes |
|
def __init__(self, opt):
super(Vimeo90KDataset, self).__init__()
self.opt = opt
# temporal augmentation
self.interval_list = opt['interval_list']
self.random_reverse = opt['random_reverse']
logger.info('Temporal augmentation interval list: [{}], with random reverse is {}.'.format(
','.join(str(x) for x in opt['interval_list']), self.random_reverse))
self.GT_root, self.LQ_root = opt['dataroot_GT'], opt['dataroot_LQ']
self.data_type = self.opt['data_type']
self.LR_input = False if opt['GT_size'] == opt['LQ_size'] else True # low resolution inputs
#### determine the LQ frame list
'''
N | frames
1 | 4
3 | 3,4,5
5 | 2,3,4,5,6
7 | 1,2,3,4,5,6,7
'''
self.LQ_frames_list = []
for i in range(opt['N_frames']):
self.LQ_frames_list.append(i + (9 - opt['N_frames']) // 2)
#### directly load image keys
if self.data_type == 'lmdb':
self.paths_GT, _ = util.get_image_paths(self.data_type, opt['dataroot_GT'])
logger.info('Using lmdb meta info for cache keys.')
elif opt['cache_keys']:
logger.info('Using cache keys: {}'.format(opt['cache_keys']))
self.paths_GT = pickle.load(open(opt['cache_keys'], 'rb'))['keys']
else:
raise ValueError(
'Need to create cache keys (meta_info.pkl) by running [create_lmdb.py]')
assert self.paths_GT, 'Error: GT path is empty.'
if self.data_type == 'lmdb':
self.GT_env, self.LQ_env = None, None
elif self.data_type == 'mc': # memcached
self.mclient = None
elif self.data_type == 'img':
pass
else:
raise ValueError('Wrong data type: {}'.format(self.data_type))
Example #26
| Source File: SRker_dataset.py From IKC with Apache License 2.0 | 4 votes |
|
def __getitem__(self, index):
if self.opt['data_type'] == 'lmdb':
if self.LR_env is None:
self._init_lmdb()
LR_size = self.LR_size
SR_size = self.SR_size
scale = self.opt['scale']
# get real kernel map
real_ker_map = self.real_ker_map_list[index]
# get each kernel map
ker_map = self.ker_map_list[index]
# get LR image
LR_path = self.LR_paths[index]
if self.opt['data_type'] == 'lmdb':
resolution = [int(s) for s in self.LR_sizes[index].split('_')]
else:
resolution = None
img_LR = util.read_img(self.LR_env, LR_path, resolution)
H, W, C = img_LR.shape
#get SR image
img_SR = self.SR_img_list[index]
if self.opt['phase'] == 'train':
#randomly crop
rnd_h = random.randint(0, max(0, H - LR_size))
rnd_w = random.randint(0, max(0, W - LR_size))
rnd_h_SR, rnd_w_SR = int(rnd_h * scale), int(rnd_w * scale)
img_SR = img_SR[rnd_h_SR:rnd_h_SR + SR_size, rnd_w_SR:rnd_w_SR + SR_size, :]
# augmentation - flip, rotate
img_SR = util.augment(img_SR, self.opt['use_flip'], self.opt['use_rot'], self.opt['mode'])
# change color space if necessary
if self.opt['color']:
img_SR = util.channel_convert(C, self.opt['color'], [img_SR])[0]
# BGR to RGB, HWC to CHW, numpy to tensor
if img_SR.shape[2] == 3:
img_SR = img_SR[:, :, [2, 1, 0]]
img_SR = torch.from_numpy(np.ascontiguousarray(np.transpose(img_SR, (2, 0, 1)))).float()
return {'SR': img_SR, 'real_ker': real_ker_map, 'ker': ker_map}
Example #27
| Source File: Vimeo90K_dataset.py From mmsr with Apache License 2.0 | 4 votes |
|
def __init__(self, opt):
super(Vimeo90KDataset, self).__init__()
self.opt = opt
# temporal augmentation
self.interval_list = opt['interval_list']
self.random_reverse = opt['random_reverse']
logger.info('Temporal augmentation interval list: [{}], with random reverse is {}.'.format(
','.join(str(x) for x in opt['interval_list']), self.random_reverse))
self.GT_root, self.LQ_root = opt['dataroot_GT'], opt['dataroot_LQ']
self.data_type = self.opt['data_type']
self.LR_input = False if opt['GT_size'] == opt['LQ_size'] else True # low resolution inputs
#### determine the LQ frame list
'''
N | frames
1 | 4
3 | 3,4,5
5 | 2,3,4,5,6
7 | 1,2,3,4,5,6,7
'''
self.LQ_frames_list = []
for i in range(opt['N_frames']):
self.LQ_frames_list.append(i + (9 - opt['N_frames']) // 2)
#### directly load image keys
if self.data_type == 'lmdb':
self.paths_GT, _ = util.get_image_paths(self.data_type, opt['dataroot_GT'])
logger.info('Using lmdb meta info for cache keys.')
elif opt['cache_keys']:
logger.info('Using cache keys: {}'.format(opt['cache_keys']))
self.paths_GT = pickle.load(open(opt['cache_keys'], 'rb'))['keys']
else:
raise ValueError(
'Need to create cache keys (meta_info.pkl) by running [create_lmdb.py]')
assert self.paths_GT, 'Error: GT path is empty.'
if self.data_type == 'lmdb':
self.GT_env, self.LQ_env = None, None
elif self.data_type == 'mc': # memcached
self.mclient = None
elif self.data_type == 'img':
pass
else:
raise ValueError('Wrong data type: {}'.format(self.data_type))
