def __getitem__(self, index):
if self.HQ_envs is None:
self._init_lmdb()
HQ_size = self.opt["HQ_size"]
env_idx, key = self.paths_HQ[index]
name_a, name_b = key.split("_")
target_frame_idx = int(name_b)
# determine the neighbor frames
# ensure not exceeding the borders
neighbor_list = [target_frame_idx]
name_b = "{:08d}".format(neighbor_list[0])
# get the HQ image (as the center frame)
img_HQ_l = []
for v in neighbor_list:
img_HQ = util.read_img(self.HQ_envs[env_idx],
"{}_{:08d}".format(name_a,
v), (3, 720, 1280))
img_HQ_l.append(img_HQ)
# get LQ images
img_LQ = util.read_img(self.LQ_envs[env_idx],
"{}_{:08d}".format(name_a, neighbor_list[-1]),
(3, 720, 1280))
if self.opt["phase"] == "train":
_, H, W = 3, 720, 1280 # LQ size
# randomly crop
rnd_h = random.randint(0, max(0, H - HQ_size))
rnd_w = random.randint(0, max(0, W - HQ_size))
img_LQ = img_LQ[rnd_h:rnd_h + HQ_size, rnd_w:rnd_w + HQ_size, :]
img_HQ_l = [
v[rnd_h:rnd_h + HQ_size, rnd_w:rnd_w + HQ_size, :]
for v in img_HQ_l
]
# augmentation - flip, rotate
img_HQ_l.append(img_LQ)
rlt = util.augment(img_HQ_l, self.opt["use_flip"],
self.opt["use_rot"])
img_HQ_l = rlt[0:-1]
img_LQ = rlt[-1]
# stack LQ images to NHWC, N is the frame number
img_HQs = np.stack(img_HQ_l, axis=0)
# BGR to RGB, HWC to CHW, numpy to tensor
img_LQ = img_LQ[:, :, [2, 1, 0]]
img_HQs = img_HQs[:, :, :, [2, 1, 0]]
img_LQ = torch.from_numpy(
np.ascontiguousarray(np.transpose(img_LQ, (2, 0, 1)))).float()
img_HQs = torch.from_numpy(
np.ascontiguousarray(np.transpose(img_HQs, (0, 3, 1, 2)))).float()
# print(img_LQ.shape, img_HQs.shape)
if self.use_identical and np.random.randint(0, 10) == 0:
img_LQ = img_HQs[-1, :, :, :]
return {"LQ": img_LQ, "HQs": img_HQs, "identical_w": 10}
return {"LQ": img_LQ, "HQs": img_HQs, "identical_w": 0}
def read_imgs(self, root_dir, name_a, name_b, is_gt=False):
if not isinstance(name_b, (tuple, list)):
if self.data_type == 'lmdb':
paths = [name_a + '_' + name_b]
else:
paths = [osp.join(root_dir, name_a, name_b + '.png')]
else:
paths = []
for name in name_b:
if not isinstance(name, str):
name = "{:08d}".format(name)
if self.data_type == 'lmdb':
paths.append(name_a + '_' + name)
else:
paths.append(osp.join(root_dir, name_a, name + '.png'))
imgs = []
for path in paths:
if self.data_type == 'lmdb':
if is_gt:
img = util.read_img(self.GT_env, path, self.GT_size_tuple)
else:
img = util.read_img(self.LQ_env, path, self.LQ_size_tuple)
else:
img = util.read_img(None, path)
imgs.append(img)
return imgs
def __getitem__(self, index):
GT_path, LQ_path = None, None
# get GT image
GT_path = self.paths_GT[index]
LQ_path = self.paths_LQ[index]
img_GT = util.read_img(self.GT_env, GT_path)
img_LQ = util.read_img(self.LQ_env, LQ_path)
if self.opt['color']:
img_GT = util.channel_convert(img_GT.shape[2], self.opt['color'],
[img_GT])[0]
img_LQ = util.channel_convert(img_LQ.shape[2], self.opt['color'],
[img_LQ])[0]
# BGR to RGB, HWC to CHW, numpy to tensor
if img_GT.shape[2] == 3:
img_GT = img_GT[:, :, [2, 1, 0]]
img_LQ = img_LQ[:, :, [2, 1, 0]]
H, W, _ = img_LQ.shape
img_GT = torch.from_numpy(
np.ascontiguousarray(np.transpose(img_GT, (2, 0, 1)))).float()
img_LQ = torch.from_numpy(
np.ascontiguousarray(np.transpose(img_LQ, (2, 0, 1)))).float()
if LQ_path is None:
LQ_path = GT_path
return {
'LQ': img_LQ,
'GT': img_GT,
'LQ_path': LQ_path,
'GT_path': GT_path
}
def __getitem__(self, index):
path, target = self.samples[index]
img_dis = read_img(env=None, path=path[0])
img_ref = read_img(env=None, path=path[1])
'''H, W, _ = img_ref.shape
crop_size = 224
rnd_h = random.randint(0, max(0, (H - crop_size) // 2))
rnd_w = random.randint(0, max(0, (W - crop_size) // 2))
img_dis = img_dis[rnd_h:rnd_h + crop_size, rnd_w:rnd_w + crop_size, :]
img_ref = img_ref[rnd_h:rnd_h + crop_size, rnd_w:rnd_w + crop_size, :]
# augmentation - flip, rotate
img_dis, img_ref = augment([img_dis, img_ref], self.opt['use_flip'], rot=False)'''
if img_ref.shape[2] == 3:
img_ref = img_ref[:, :, [2, 1, 0]]
img_dis = img_dis[:, :, [2, 1, 0]]
img_ref = torch.from_numpy(
np.ascontiguousarray(np.transpose(img_ref, (2, 0, 1)))).float()
img_dis = torch.from_numpy(
np.ascontiguousarray(np.transpose(img_dis, (2, 0, 1)))).float()
img_dis = self.transform(img_dis)
img_ref = self.transform(img_ref)
return {
'Dis': img_dis,
'Ref': img_ref,
'Label': target,
'Dis_path': path[0]
}
def __getitem__(self, index):
HR_path, LR_path = None, None
scale = self.opt['scale']
# get HR image
HR_path = self.paths_HR[index]
img_HR = util.read_img(self.HR_env, HR_path)
# get LR image
if self.paths_LR:
LR_path = self.paths_LR[index]
D1_path = self.paths_D1[index]
D2_path = self.paths_D2[index]
D3_path = self.paths_D3[index]
img_LR = util.read_img(self.LR_env, LR_path)
img_D1 = util.read_img(self.D1_env, D1_path)
img_D2 = util.read_img(self.D2_env, D2_path)
img_D3 = util.read_img(self.D3_env, D3_path)
else: # down-sampling on-the-fly
if self.opt['phase'] == 'train':
# force to 3 channels
if img_HR.ndim == 2:
img_HR = cv2.cvtColor(img_HR, cv2.COLOR_GRAY2BGR)
H, W, _ = img_HR.shape
# using matlab imresize
img_LR = cv2.resize(img_HR, dsize=(int(H / scale), int(W / scale)), interpolation=cv2.INTER_CUBIC)
img_D1 = cv2.resize(img_HR, dsize=(int(H / scale * 2), int(W / scale * 2)), interpolation=cv2.INTER_CUBIC)
img_D2 = cv2.resize(img_HR, dsize=(int(H / scale * 4), int(W / scale * 4)), interpolation=cv2.INTER_CUBIC)
img_D3 = cv2.resize(img_HR, dsize=(int(H / scale * 8), int(W / scale * 8)), interpolation=cv2.INTER_CUBIC)
if img_LR.ndim == 2:
img_LR = np.expand_dims(img_LR, axis=2)
img_D1 = np.expand_dims(img_D1, axis=2)
img_D2 = np.expand_dims(img_D2, axis=2)
img_D3 = np.expand_dims(img_D3, axis=2)
if self.opt['phase'] == 'train':
# augmentation - flip, rotate
img_LR, img_HR, img_D1, img_D2, img_D3 = util.augment([img_LR, img_HR, img_D1, img_D2, img_D3],
self.opt['use_flip'], self.opt['use_rot'])
# BGR to RGB, HWC to CHW, numpy to tensor
if img_HR.shape[2] == 3:
img_HR = img_HR[:, :, [2, 1, 0]]
img_LR = img_LR[:, :, [2, 1, 0]]
img_D1 = img_D1[:, :, [2, 1, 0]]
img_D2 = img_D2[:, :, [2, 1, 0]]
img_D3 = img_D3[:, :, [2, 1, 0]]
img_HR = torch.from_numpy(np.ascontiguousarray(np.transpose(img_HR, (2, 0, 1)))).float()
img_LR = torch.from_numpy(np.ascontiguousarray(np.transpose(img_LR, (2, 0, 1)))).float()
img_D1 = torch.from_numpy(np.ascontiguousarray(np.transpose(img_D1, (2, 0, 1)))).float()
img_D2 = torch.from_numpy(np.ascontiguousarray(np.transpose(img_D2, (2, 0, 1)))).float()
img_D3 = torch.from_numpy(np.ascontiguousarray(np.transpose(img_D3, (2, 0, 1)))).float()
if LR_path is None:
LR_path = HR_path
return {'LR': img_LR, 'HR': img_HR, 'LR_path': LR_path, 'HR_path': HR_path,
'D1': img_D1, 'D2': img_D2, 'D3': img_D3, 'D1_path': D1_path, 'D2_path': D2_path, 'D3_path': D3_path}
def __getitem__(self, index):
HR_path, LR_path, MR_path = None, None, None
scale = self.opt['scale']
HR_size = self.opt['HR_size']
# get HR image
HR_path = self.paths_HR[index]
img_HR = util.read_img(self.HR_env, HR_path)
# modcrop in the validation / test phase
# if self.opt['phase'] != 'train':
# img_HR = util.modcrop(img_HR, scale)
# change color space if necessary
if self.opt['color']:
img_HR = util.channel_convert(img_HR.shape[2], self.opt['color'], [img_HR])[0]
LR_path = self.paths_LR[index]
img_LR = util.read_img(self.LR_env, LR_path)
MR_path = self.paths_MR[index]
img_MR = util.read_img(self.MR_env, MR_path)
if self.opt['noise_gt']:
img_MR = img_LR - img_MR
if self.opt['phase'] == 'train':
# if the image size is too small
H, W, C = img_LR.shape
LR_size = HR_size // scale
# randomly crop
rnd_h = random.randint(0, max(0, H - LR_size))
rnd_w = random.randint(0, max(0, W - LR_size))
img_MR = img_MR[rnd_h:rnd_h + LR_size, rnd_w:rnd_w + LR_size, :]
img_LR = img_LR[rnd_h:rnd_h + LR_size, rnd_w:rnd_w + LR_size, :]
rnd_h_HR, rnd_w_HR = int(rnd_h * scale), int(rnd_w * scale)
img_HR = img_HR[rnd_h_HR:rnd_h_HR + HR_size, rnd_w_HR:rnd_w_HR + HR_size, :]
# augmentation - flip, rotate
img_MR, img_LR, img_HR = util.augment([img_MR, img_LR, img_HR], self.opt['use_flip'], \
self.opt['use_rot'])
# channel conversion
if self.opt['color']:
# img_HR, img_LR, img_MR = util.channel_convert(C, self.opt['color'], [img_HR, img_LR, img_MR])
img_LR = util.channel_convert(C, self.opt['color'], [img_LR])[0]
img_MR = util.channel_convert(C, self.opt['color'], [img_MR])[0]
# BGR to RGB, HWC to CHW, numpy to tensor
if img_HR.shape[2] == 3:
img_HR = img_HR[:, :, [2, 1, 0]]
img_LR = img_LR[:, :, [2, 1, 0]]
img_MR = img_MR[:, :, [2, 1, 0]]
img_HR = torch.from_numpy(np.ascontiguousarray(np.transpose(img_HR, (2, 0, 1)))).float()
img_LR = torch.from_numpy(np.ascontiguousarray(np.transpose(img_LR, (2, 0, 1)))).float()
img_MR = torch.from_numpy(np.ascontiguousarray(np.transpose(img_MR, (2, 0, 1)))).float()
return {'HR': img_HR, 'LR': img_LR, 'MR': img_MR, 'HR_path': HR_path, 'MR_path': MR_path, 'LR_path': LR_path}
def read_img(self, key, is_gt=False):
if self.data_type == 'lmdb':
env = self.GT_env if is_gt else self.LQ_env
sizes = self.GT_size_tuple if is_gt else self.LQ_size_tuple
img = util.read_img(env, key, sizes)
else:
data_root = self.GT_root if is_gt else self.LQ_root
name_a, name_b = key.split('_')
im_path = osp.join(data_root, name_a, name_b + '.png')
img = util.read_img(None, path)
return img
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]
select_idx = util.index_generation(idx,
max_idx,
self.opt['N_frames'],
padding=self.opt['padding'])
if self.data_type == 'lmdb':
if self.GT_env is None or self.LQ_env is None:
self._init_lmdb()
key = self.lmdb_paths_GT[index]
name_a, name_b = key.split('_')
center_frame_idx = int(name_b)
GT_size_tuple = self.opt['GT_shape']
LQ_size_tuple = self.opt['LQ_shape']
img_GT = util.read_img(self.GT_env, key, GT_size_tuple)
img_LQ_l = []
for v in select_idx:
img_LQ = util.read_img(self.LQ_env,
'{}_{:08d}'.format(name_a,
v), LQ_size_tuple)
img_LQ_l.append(img_LQ)
# stack LQ images to NHWC, N is the frame number
img_LQs = np.stack(img_LQ_l, axis=0)
# BGR to RGB, HWC to CHW, numpy to tensor
img_GT = img_GT[:, :, [2, 1, 0]]
img_LQs = img_LQs[:, :, :, [2, 1, 0]]
img_GT = torch.from_numpy(
np.ascontiguousarray(np.transpose(img_GT, (2, 0, 1)))).float()
imgs_LQ = torch.from_numpy(
np.ascontiguousarray(np.transpose(img_LQs,
(0, 3, 1, 2)))).float()
elif self.cache_data:
imgs_LQ = self.imgs_LQ[folder].index_select(
0, torch.LongTensor(select_idx))
img_GT = self.imgs_GT[folder][idx]
else:
imgs_LQ = util.read_img_seq(self.imgs_LQ[folder]).index_select(
0, torch.LongTensor(select_idx))
img_GT = util.read_img_seq(self.imgs_GT[folder])[idx]
return {
'LQs': imgs_LQ,
'GT': img_GT,
'folder': folder,
'idx': self.data_info['idx'][index],
'border': border
}
def __getitem__(self, index):
HR_path, LR_path = None, None
scale = 4
HR_path = self.paths_HR[index]
img_HR = util.read_img(self.HR_env, HR_path)
if self.cnf['phase'] != 'train':
img_HR = util.modcrop(img_HR, scale)
LR_path = self.paths_LR[index]
img_LR = util.read_img(self.LR_env, LR_path)
if self.cnf['phase'] == 'train':
H, W, _ = img_HR.shape
if H < self.hr_sz or W < self.hr_sz:
img_HR = cv2.resize(np.copy(img_HR), (self.hr_sz, self.hr_sz),
interpolation=cv2.INTER_LINEAR)
img_LR = util.imresize_np(img_HR, 1 / scale, True)
if img_LR.ndim == 2:
img_LR = np.expand_dims(img_LR, axis=2)
H, W, C = img_LR.shape
LR_size = self.hr_sz // scale
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, :]
rnd_h_HR, rnd_w_HR = int(rnd_h * scale), int(rnd_w * scale)
img_HR = img_HR[rnd_h_HR:rnd_h_HR + self.hr_sz,
rnd_w_HR:rnd_w_HR + self.hr_sz, :]
img_LR, img_HR = util.augment([img_LR, img_HR], True, True)
if img_HR.shape[2] == 3:
img_HR = img_HR[:, :, [2, 1, 0]]
img_LR = img_LR[:, :, [2, 1, 0]]
img_HR = from_numpy(
np.ascontiguousarray(np.transpose(img_HR, (2, 0, 1)))).float()
img_LR = from_numpy(
np.ascontiguousarray(np.transpose(img_LR, (2, 0, 1)))).float()
if LR_path is None:
LR_path = HR_path
return {
'LR': img_LR,
'HR': img_HR,
'LR_path': LR_path,
'HR_path': HR_path
}
def read_data(self, index):
video_name = self.videos_name[index]
idx, max_idx = self.idx_maxidx[index].split('/')
idx, max_idx = int(idx), int(max_idx)
indices = util.index_generation(idx, max_idx, self.n_frames)
lr_path = osp.dirname(self.lr_paths[index])
filename = osp.basename(self.lr_paths[index])
lrs = []
for i in indices:
lr = util.read_img(None, osp.join(lr_path, '{:03d}.png'.format(i)))
lrs.append(lr)
gt = util.read_img(None, self.gt_paths[index])
return lrs, gt, video_name, filename
def __getitem__(self, index):
HR_path, LR_path = None, None
HR_size = self.opt['HR_size']
# get HR image
HR_path = self.paths_HR[index]
img_HR = util.read_img(self.HR_env, HR_path)
# get LR image
LR_path = self.paths_LR[index]
img_LR = util.read_img(self.LR_env, LR_path)
# modcrop in the validation / test phase
if self.opt['phase'] != 'train':
img_HR = util.modcrop(img_HR, 2)
img_LR = util.modcrop(img_LR, 2)
if self.opt['phase'] == 'train':
H, W, C = img_LR.shape
LR_size = HR_size
# 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, :]
img_HR = img_HR[rnd_h:rnd_h + HR_size, rnd_w:rnd_w + HR_size, :]
# augmentation - flip, rotate
img_LR, img_HR = util.augment([img_LR, img_HR], self.opt['use_flip'], \
self.opt['use_rot'])
# BGR to RGB, HWC to CHW, numpy to tensor
if img_HR.shape[2] == 3:
img_HR = img_HR[:, :, [2, 1, 0]]
img_LR = img_LR[:, :, [2, 1, 0]]
img_HR = torch.from_numpy(
np.ascontiguousarray(np.transpose(img_HR, (2, 0, 1)))).float()
img_LR = torch.from_numpy(
np.ascontiguousarray(np.transpose(img_LR, (2, 0, 1)))).float()
if LR_path is None:
LR_path = HR_path
return {
'LR': img_LR,
'HR': img_HR,
'LR_path': LR_path,
'HR_path': HR_path
}
def read_imgs(self, key):
imgs = []
if self.data_type == 'lmdb':
for idx in self.LQ_frames_list:
path = key[:-1] + str(idx)
imgs.append(
util.read_img(self.GT_env,
path,
self.GT_size_tuple,
dtype='uint8'))
else:
for idx in self.LQ_frames_list:
path = key[:-5] + str(idx) + key[-4:]
imgs.append(util.read_img(None, path, dtype='uint8'))
return np.stack(imgs)
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]
if self.data_type == 'lmdb':
resolution = [int(s) for s in self.sizes_LQ[index].split('_')]
else:
resolution = None
img_LQ = util.read_img(self.LQ_env, LQ_path, resolution)
H, W, C = img_LQ.shape
# change color space if necessary
if self.opt['color']:
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}
def read_sub_images(self, paths):
if not isinstance(paths, (list, tuple)):
paths = [paths]
imgs = []
for path in paths:
imgs.append(util.read_img(None, path, dtype='uint8'))
return np.stack(imgs)
def __getitem__(self, index):
# get full size image
full_path = self.paths_hq[index % len(self.paths_hq)]
loaded_img = util.read_img(None, full_path, None)
img_full1 = util.channel_convert(loaded_img.shape[2], 'RGB',
[loaded_img])[0]
img_full2 = util.augment([img_full1], True, True)[0]
img_full3 = get_square_image(img_full2)
# This error crops up from time to time. I suspect an issue with util.read_img.
if img_full3.shape[0] == 0 or img_full3.shape[1] == 0:
print(
"Error with image: %s. Loaded image shape: %s" %
(full_path, str(loaded_img.shape)), str(img_full1.shape),
str(img_full2.shape), str(img_full3.shape))
# Attempt to recover by just using a fixed array of zeros, which the downstream networks should be fine training against, within reason.
img_full3 = np.zeros((1024, 1024, 3), dtype=np.int)
img_full = cv2.resize(img_full3, (self.hq_size_cap, self.hq_size_cap),
interpolation=cv2.INTER_AREA)
patches_hq = [
cv2.resize(img_full, (self.tile_size, self.tile_size),
interpolation=cv2.INTER_AREA)
]
self.recursively_extract_patches(img_full, patches_hq, 1)
# Image corruption is applied against the full size image for this dataset.
img_corrupted = self.corruptor.corrupt_images([img_full])[0]
patches_hq_corrupted = [
cv2.resize(img_corrupted, (self.tile_size, self.tile_size),
interpolation=cv2.INTER_AREA)
]
self.recursively_extract_patches(img_corrupted, patches_hq_corrupted,
1)
# BGR to RGB, HWC to CHW, numpy to tensor
if patches_hq[0].shape[2] == 3:
patches_hq = [
cv2.cvtColor(p, cv2.COLOR_BGR2RGB) for p in patches_hq
]
patches_hq_corrupted = [
cv2.cvtColor(p, cv2.COLOR_BGR2RGB)
for p in patches_hq_corrupted
]
patches_hq = [
torch.from_numpy(np.ascontiguousarray(np.transpose(
p, (2, 0, 1)))).float() for p in patches_hq
]
patches_hq = torch.stack(patches_hq, dim=0)
patches_hq_corrupted = [
torch.from_numpy(np.ascontiguousarray(np.transpose(
p, (2, 0, 1)))).float() for p in patches_hq_corrupted
]
patches_lq = [
torch.nn.functional.interpolate(p.unsqueeze(0),
scale_factor=1 / self.scale,
mode='area').squeeze()
for p in patches_hq_corrupted
]
patches_lq = torch.stack(patches_lq, dim=0)
d = {'lq': patches_lq, 'hq': patches_hq, 'GT_path': full_path}
return d
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
##########################################################################
img_LR = cv2.resize(
np.copy(img_LR), (int(W / 2), int(H / 2)),
interpolation=cv2.INTER_LINEAR) # for avoiding out of memory
# force to 3 channels
if img_LR.ndim == 2:
img_LR = cv2.cvtColor(img_LR, cv2.COLOR_GRAY2BGR)
###########################################################################
# 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}
def __getitem__(self, index):
image_path = self.image_paths[index]
label = self.image_labels[index]
image = util.read_img(image_path, self.target_size)
mask = util.read_mask(
image_path.replace('images',
'segmentations').replace('jpg', 'png'),
self.target_size)
if image.shape[2] == 3:
image = image[:, :, [2, 1, 0]]
## data augmentation
if random.random() > 1.0:
aug = util.get_augmentor()
angle = random.randint(-25, 25)
image = util.rotate(image, angle)
image = aug.augment_image(image)
image -= np.array([123.68, 116.779, 103.939])
image /= np.array([58.393, 57.12, 57.375])
else:
raise Exception('The shape of {} is not 3.'.format(image_path))
image = torch.from_numpy(
np.ascontiguousarray(np.transpose(image, (2, 0, 1)))).float()
return {'inputs': image, 'labels': label, 'mask': mask}
def __getitem__(self, index):
# get LQ and HQ image
LQ_path = self.paths_LQ[index]
HQ_path = self.paths_HQ[index]
img_LQ = util.read_img(None, LQ_path, None)
img_HQ = util.read_img(None, HQ_path, None)
if self.opt['crop_size']:
# randomly crop
LH, LW, _ = img_LQ.shape
rnd_h = random.randint(0, max(0, LH - self.crop_size))
rnd_w = random.randint(0, max(0, LW - self.crop_size))
rnd_hh = rnd_h * self.opt['scale']
rnd_wh = rnd_w * self.opt['scale']
patch_size = self.crop_size * self.opt['scale']
img_LQ = img_LQ[rnd_h:rnd_h + self.crop_size, rnd_w:rnd_w + self.crop_size, :]
img_HQ = img_HQ[rnd_hh:rnd_hh + patch_size, rnd_wh:rnd_wh + patch_size, :]
if self.opt['phase'] == 'train':
# augmenttation cutclur
if self.opt['use_cutblur']:
img_LQ,img_HQ=util.argment_cutblur(img_LQ, img_HQ,self.opt['scale'])
if self.opt['use_rgbpermute']:
img_LQ, img_HQ=util.argment_rgb(img_LQ, img_HQ)
# augmentation - flip, rotate
img_LQ, img_HQ = util.augment([img_LQ, img_HQ], self.opt['use_flip'],
self.opt['use_rot'])
# BGR to RGB, HWC to CHW, numpy to tensor
if img_LQ.shape[2] == 3:
img_LQ = img_LQ[:, :, [2, 1, 0]]
if img_HQ.shape[2] == 3:
img_HQ = img_HQ[:, :, [2, 1, 0]]
img_LQ = torch.from_numpy(np.ascontiguousarray(np.transpose(img_LQ, (2, 0, 1)))).float()
img_HQ = torch.from_numpy(np.ascontiguousarray(np.transpose(img_HQ, (2, 0, 1)))).float()
return {'LQ': img_LQ, 'LQ_path': LQ_path, 'HQ': img_HQ, 'HQ_path': HQ_path}
def imgs2split_imgs(self, imgs_path):
# Used as counter variables
count, split_count= (0, 0)
for image_path in data_util.get_image_paths('img', imgs_path)[0]:
image = data_util.read_img(None, path=image_path)
# Saves the frames and split them every 100 frames
if split_count == count:
split_count += 100
util.mkdir(f'{self.imgs_cache}{split_count}')
cv2.imwrite(f"{self.imgs_cache}{split_count}/{count}.png", image)
count += 1
def read_imgs(self, root_dir, name_a, name_bs):
assert (len(name_bs) > 1)
imgs = []
if self.data_type == 'lmdb':
for name in name_bs:
if not isinstance(name, str):
name = "{:08d}".format(name)
path = name_a + '_' + name
imgs.append(
util.read_img(self.GT_env,
path,
self.GT_size_tuple,
dtype='uint8'))
else:
for name in name_bs:
if not isinstance(name, str):
name = "{:08d}".format(name)
path = osp.join(root_dir, name_a, name + '.png')
imgs.append(util.read_img(None, path, dtype='uint8'))
return np.stack(imgs)
def __getitem__(self, index):
# get LR image
LR_path = self.paths_LR[index]
img_LR = util.read_img(self.LR_env, LR_path)
# 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}
def __getitem__(self, index):
if self.data_type == 'lmdb':
if (self.GT_env is None) or (self.LQ_env is None):
self._init_lmdb()
GT_path, LQ_path = None, None
scale = self.opt['scale']
GT_size = self.opt['GT_size']
# get GT image
GT_path = self.paths_GT[index]
if self.data_type == 'lmdb':
resolution = [int(s) for s in self.sizes_GT[index].split('_')]
else:
resolution = None
img_GT = util.read_img(self.GT_env, GT_path, resolution)
# modcrop in the validation / test phase
if self.opt['phase'] != 'train':
img_GT = util.modcrop(img_GT, scale)
# print('val img_GT shape: ', img_GT.shape)
# change color space if necessary
if self.opt['color']:
img_GT = util.channel_convert(img_GT.shape[2], self.opt['color'],
[img_GT])[0]
# print('img_GT shape: ', img_GT.shape)
if img_GT.ndim == 2:
img_GT = np.expand_dims(img_GT, axis=2)
if self.opt['phase'] == 'train':
# if the image size is too small
H, W, C = img_GT.shape
if H < GT_size or W < GT_size:
img_GT = cv2.resize(np.copy(img_GT), (GT_size, GT_size),
interpolation=cv2.INTER_LINEAR)
# randomly crop - GT img shape: [H, W, 1]
rnd_h = random.randint(0, max(0, H - GT_size))
rnd_w = random.randint(0, max(0, W - GT_size))
img_GT = img_GT[rnd_h:rnd_h + GT_size, rnd_w:rnd_w + GT_size, :]
# augmentation - flip, rotate
img_GT = util.augment([img_GT], self.opt['use_flip'],
self.opt['use_rot'])[0]
# transform() - subsample k space - img_LF
# input img shape [H, W, 1] or [H, W, 3]
img_LF, img_GT = self.transform(img_GT, self.mask_func, self.seed)
if LQ_path is None:
LQ_path = GT_path
return {
'LQ': img_LF,
'GT': img_GT,
'LQ_path': LQ_path,
'GT_path': GT_path
}
def __getitem__(self, index):
if self.data_type == "mc":
self._ensure_memcached()
elif self.data_type == "lmdb" and (self.HQ_env is None
or self.LQ_env is None):
self._init_lmdb()
HQ_size = self.opt["HQ_size"]
key = self.paths_HQ[index]
name_a, name_b = key.split("_")
# get the HQ image
img_HQ = util.read_img(self.HQ_env, key, (3, 720, 1280))
# get the LQ image
img_LQ = util.read_img(self.LQ_env, key, (3, 720, 1280))
if self.opt["phase"] == "train":
_, H, W = 3, 720, 1280 # LQ size
# randomly crop
rnd_h = random.randint(0, max(0, H - HQ_size))
rnd_w = random.randint(0, max(0, W - HQ_size))
img_LQ = img_LQ[rnd_h:rnd_h + HQ_size, rnd_w:rnd_w + HQ_size, :]
img_HQ = img_HQ[rnd_h:rnd_h + HQ_size, rnd_w:rnd_w + HQ_size, :]
# augmentation - flip, rotate
imgs = [img_HQ, img_LQ]
rlt = util.augment(imgs, self.opt["use_flip"], self.opt["use_rot"])
img_HQ = rlt[0]
img_LQ = rlt[1]
# BGR to RGB, HWC to CHW, numpy to tensor
img_LQ = img_LQ[:, :, [2, 1, 0]]
img_HQ = img_HQ[:, :, [2, 1, 0]]
img_LQ = torch.from_numpy(
np.ascontiguousarray(np.transpose(img_LQ, (2, 0, 1)))).float()
img_HQ = torch.from_numpy(
np.ascontiguousarray(np.transpose(img_HQ, (2, 0, 1)))).float()
return {"LQ": img_LQ, "HQ": img_HQ}
def __getitem__(self, index):
if self.opt["data_type"] == "lmdb":
if self.LR_env is None:
self._init_lmdb()
LR_path = None
scale = self.opt["scale"]
LR_size = self.opt["LR_size"]
# 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
) # return: Numpy float32, HWC, BGR, [0,1]
# modcrop in the validation / test phase
if self.opt["phase"] != "train":
img_LR = util.modcrop(img_LR, scale)
if self.opt["phase"] == "train":
H, W, C = img_LR.shape
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(img_LR.shape[2], 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}
def __getitem__(self, index):
scale = self.opt['scale']
LQ_size = self.opt['LQ_size']
path_LQ = self.paths_LQ[index]
path_GT = self.paths_GT[index]
name_a, name_b = path_LQ.split('/')[-2], path_LQ.split('/')[-1]
# center LQ
img_LQ_l = []
for v in range(7):
# TODO
if v == 3:
img_LQ = util.read_img(
None,
osp.join(self.GT_root, name_a, name_b,
'im{}.png'.format(v + 1)))
img_LQ_l.append(img_LQ)
else:
img_LQ = util.read_img(
None,
osp.join(self.LQ_root, name_a, name_b,
'im{}.png'.format(v + 1)))
img_LQ_l.append(img_LQ)
# LQ_size_tuple = (3, 64, 112) if self.LR_input else (3, 256, 448)
# TODO
# stack LQ images to NHWC, N is the frame number
img_LQs = np.stack(img_LQ_l, axis=0)
# BGR to RGB, HWC to CHW, numpy to tensor
img_LQs = img_LQs[:, :, :, [2, 1, 0]]
img_LQs = torch.from_numpy(
np.ascontiguousarray(np.transpose(img_LQs, (0, 3, 1, 2)))).float()
img_LQs = img_LQs[:, :, :, 125:381]
return {
'LQs': img_LQs,
'key': name_a + '_' + name_b,
}
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}
def __getitem__(self, index):
LR_path = None
LR_path = self.paths_LR[index]
img_LR = util.read_img(self.LR_env, LR_path)
H, W, C = img_LR.shape
if self.opt['color']:
img_LR = util.channel_convert(C, self.opt['color'], [img_LR])[0]
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}
def read_imgs(self, root_dir, name_a, name_b, scale=None):
if scale is None:
paths = [osp.join(root_dir, name_a, name_b + '.png')]
else:
if not isinstance(name_b, (tuple, list)):
name_b = [name_b]
paths = []
for name in name_b:
if not isinstance(name, str):
name = "{:08d}".format(name)
paths.append(
osp.join(root_dir, 'X{:.02f}'.format(scale), name_a,
name + '.png'))
imgs = []
for path in paths:
imgs.append(util.read_img(None, path))
return imgs
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
# channel conversion
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}
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()
returned_dict = {'LR': img_LR, 'LR_path': LR_path}
if self.paths_kernel is not None:
returned_dict['kernel'] = scipy.io.loadmat(self.paths_kernel[index])['Kernel']
return returned_dict
