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):
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 __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 __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):
# 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):
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.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):
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 __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
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()
if self.opt['size'] is not None:
img_LR = TF.to_pil_image(img_LR)
img_LR = T.CenterCrop(self.opt['size'])(img_LR)
img_LR = TF.to_tensor(img_LR)
return {'LQ': img_LR, 'LQ_path': LR_path}
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}
def __getitem__(self, index):
LR_path = None
if self.opt['znorm']:
if self.opt['znorm'] == True:
self.znorm = True # Alternative: images are z-normalized to the [-1,1] range
# get LR image
LR_path = self.paths_LR[index]
#img_LR = util.read_img(self.LR_env, LR_path)
img_LR = util.read_img(self.LR_env, LR_path, znorm=self.znorm)
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]]
elif img_LR.shape[2] == 4:
img_LR = img_LR[:, :, [2, 1, 0, 3]]
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' and (self.GT_env is None or self.LQ_UX4_env is None):
self._init_lmdb()
GT_path, LQ_UX4_path, Ref_path, SR_path = None, None, None, None
scale = self.opt['scale']
GT_size = self.opt['GT_size']
# get GT image
GT_path = self.paths_GT[index]
resolution = [int(s) for s in self.sizes_GT[index].split('_')
] if self.data_type == 'lmdb' else None
img_GT = util.read_img(self.GT_env, GT_path, resolution)
if self.opt['phase'] != 'train': # modcrop in the validation / test phase
img_GT = util.modcrop(img_GT, scale)
#if self.opt['color']: # change color space if necessary
# img_GT = util.channel_convert(img_GT.shape[2], self.opt['color'], [img_GT])[0]
# get Ref image
Ref_path = self.paths_Ref[index]
resolution = [int(s) for s in self.sizes_Ref[index].split('_')
] if self.data_type == 'lmdb' else None
img_Ref = util.read_img(self.Ref_env, Ref_path, resolution)
if self.opt['Ref_color']: # change color space if necessary
img_Ref = util.channel_convert(img_Ref.shape[2], self.opt['Ref_color'], [img_Ref])[0]
SR_path = self.paths_SR[index]
resolution = [int(s) for s in self.sizes_SR[index].split('_')
] if self.data_type == 'lmdb' else None
img_SR = util.read_img(self.SR_env, SR_path, resolution)
# get LQ_UX4 image
if self.paths_LQ_UX4:
LQ_UX4_path = self.paths_LQ_UX4[index]
resolution = [int(s) for s in self.sizes_LQ_UX4[index].split('_')
] if self.data_type == 'lmdb' else None
img_LQ_UX4 = util.read_img(self.LQ_UX4_env, LQ_UX4_path, resolution)
if self.opt['phase'] == 'train':
H, W, C = img_LQ_UX4.shape
LQ_size = GT_size
# randomly crop
rnd_h = random.randint(0, max(0, H - LQ_size))
rnd_w = random.randint(0, max(0, W - LQ_size))
img_LQ_UX4 = img_LQ_UX4[rnd_h:rnd_h + LQ_size, rnd_w:rnd_w + LQ_size, :]
img_Ref = img_Ref[rnd_h:rnd_h + LQ_size, rnd_w:rnd_w + LQ_size, :]
img_SR = img_SR[rnd_h:rnd_h + LQ_size, rnd_w:rnd_w + LQ_size, :]
img_GT = img_GT[rnd_h:rnd_h + LQ_size, rnd_w:rnd_w + LQ_size, :]
# augmentation - flip, rotate
img_LQ_UX4, img_Ref, img_SR, img_GT = util.augment([img_LQ_UX4, img_Ref, img_SR, img_GT], self.opt['use_flip'],
self.opt['use_rot'])
# BGR to RGB, HWC to CHW, numpy to tensor
if img_GT.shape[2] == 3:
img_GT = img_GT[:, :, [2, 1, 0]]
img_LQ_UX4 = img_LQ_UX4[:, :, [2, 1, 0]]
img_SR = img_SR[:, :, [2, 1, 0]]
if img_Ref.shape[2] == 3:
img_Ref = img_Ref[:, :, [2, 1, 0]]
img_GT = torch.from_numpy(np.ascontiguousarray(np.transpose(img_GT, (2, 0, 1)))).float()
img_LQ_UX4 = torch.from_numpy(np.ascontiguousarray(np.transpose(img_LQ_UX4, (2, 0, 1)))).float()
img_SR = torch.from_numpy(np.ascontiguousarray(np.transpose(img_SR, (2, 0, 1)))).float()
img_Ref = torch.from_numpy(np.ascontiguousarray(np.transpose(img_Ref, (2, 0, 1)))).float()
return {'LQ_UX4': img_LQ_UX4,'Ref': img_Ref, 'SR': img_SR, 'GT': img_GT, 'LQ_UX4_path': LQ_UX4_path, 'Ref_path': Ref_path, 'SR_path': SR_path, 'GT_path': GT_path}
def __getitem__(self, index):
HR_path, LR_path = 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 validation phase
if self.opt['phase'] != 'train':
img_HR = util.modcrop(img_HR, scale)
# get LR image
if self.paths_LR:
LR_path = self.paths_LR[index]
img_LR = util.read_img(self.LR_env, LR_path)
else: # down-sampling on-the-fly
# randomly scale during training
if self.opt['phase'] == 'train':
random_scale = random.choice(self.random_scale_list)
H_s, W_s, _ = img_HR.shape
def _mod(n, random_scale, scale, thres):
rlt = int(n * random_scale)
rlt = (rlt // scale) * scale
return thres if rlt < thres else rlt
H_s = _mod(H_s, random_scale, scale, HR_size)
W_s = _mod(W_s, random_scale, scale, HR_size)
img_HR = cv2.resize(np.copy(img_HR), (W_s, H_s),
interpolation=cv2.INTER_LINEAR)
H, W, _ = img_HR.shape
# using matlab imresize
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
if self.opt['phase'] == 'train':
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_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_LR, img_HR = util.augment([img_LR, img_HR], self.opt['use_flip'], \
self.opt['use_rot'])
# channel conversion
if self.opt['color']:
img_LR, img_HR = util.channel_convert(C, self.opt['color'],
[img_LR, img_HR])
# 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 __getitem__(self, index):
HR_path, LR_path = 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]
if self.opt['resize'] < 1:
print('Resize by %.2f' % self.opt['resize'])
img_HR = cv2.resize(img_HR, None, fx=self.opt['resize'], fy=self.opt['resize'], interpolation=cv2.INTER_LINEAR)
# print(img_HR.shape)
# get LR image
if self.paths_LR:
LR_path = self.paths_LR[index]
img_LR = util.read_img(self.LR_env, LR_path)
else: # down-sampling on-the-fly
# randomly scale during training
if self.opt['phase'] == 'train':
random_scale = random.choice(self.random_scale_list)
H_s, W_s, _ = img_HR.shape
def _mod(n, random_scale, scale, thres):
rlt = int(n * random_scale)
rlt = (rlt // scale) * scale
return thres if rlt < thres else rlt
H_s = _mod(H_s, random_scale, scale, HR_size)
W_s = _mod(W_s, random_scale, scale, HR_size)
img_HR = cv2.resize(np.copy(img_HR), (W_s, H_s), interpolation=cv2.INTER_LINEAR)
# force to 3 channels
if img_HR.ndim == 2:
img_HR = cv2.cvtColor(img_HR, cv2.COLOR_GRAY2BGR)
H, W, _ = img_HR.shape
if self.opt['downsample'] == 'cubic':
img_LR = cv2.resize(img_HR, dsize=None, fx=1.0/scale, fy=1.0/scale, interpolation=cv2.INTER_CUBIC)
elif self.opt['downsample'] == 'numpy':
img_LR = util.imresize_np(img_HR, 1 / scale, True)
elif self.opt['downsample'] == 'linear':
img_LR = cv2.resize(img_HR, dsize=None, fx=1.0 / scale, fy=1.0 / scale, interpolation=cv2.INTER_LINEAR)
if img_LR.ndim == 2:
img_LR = np.expand_dims(img_LR, axis=2)
# print(img_HR.shape)
# print(img_LR.shape)
if self.opt['phase'] == 'train':
# if the image size is too small
H, W, _ = img_HR.shape
if H < HR_size or W < HR_size or not self.opt['crop']:
img_HR = cv2.resize(
np.copy(img_HR), (HR_size, HR_size), interpolation=cv2.INTER_LINEAR)
# using matlab imresize
if self.opt['downsample'] == 'cubic':
img_LR = cv2.resize(img_HR, dsize=None, fx=1.0 / scale, fy=1.0 / scale,
interpolation=cv2.INTER_CUBIC)
elif self.opt['downsample'] == 'numpy':
img_LR = util.imresize_np(img_HR, 1 / scale, True)
elif self.opt['downsample'] == 'linear':
img_LR = cv2.resize(img_HR, dsize=None, fx=1.0 / scale, fy=1.0 / scale,
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 = HR_size // scale
# randomly crop
if self.opt['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, :]
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_LR, img_HR = util.augment([img_LR, img_HR], self.opt['use_flip'], \
self.opt['use_rot'])
if self.LR_random_fuzzy is not None:
random_fuzzy = random.choice(self.LR_random_fuzzy)
assert self.opt['downsample'] == 'numpy'
init_LR = np.copy(img_LR)
img_LR = util.imresize(img_LR, random_fuzzy, True)
img_LR = util.imresize(img_LR, 1 / random_fuzzy, True)
if img_LR.shape[0] != LR_size or img_LR.shape[1] != LR_size:
print('Warning: LR shape changed after random fuzzy. Using initial one.', img_LR.shape[0], img_LR.shape[1], LR_size)
img_LR = init_LR
# change color space if necessary
if self.opt['color']:
img_LR = util.channel_convert(C, self.opt['color'], [img_LR])[0] # TODO during val no definetion
# 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()
# return: 0-1
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 __getitem__(self, index):
HR_path, LR_path = None, None
scale = self.opt['scale']
HR_size = self.opt['HR_size']
# get HR image
index1 = index % len(self.paths_HR)
HR_path = self.paths_HR[index1]
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]
# get LR image
if self.paths_LR:
index2 = index % len(self.paths_LR)
LR_path = self.paths_LR[index2]
img_LR = util.read_img(self.LR_env, LR_path)
if self.opt['phase'] == 'train':
# if the image size is too small
H, W, _ = img_HR.shape
if H < HR_size or W < HR_size:
img_HR = cv2.resize(
np.copy(img_HR), (HR_size, HR_size), interpolation=cv2.INTER_LINEAR)
# using matlab imresize
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 = 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_LR = img_LR[rnd_h:rnd_h + LR_size, rnd_w:rnd_w + LR_size, :]
H, W, C = img_HR.shape
# randomly crop
rnd_h = random.randint(0, max(0, H - HR_size))
rnd_w = random.randint(0, max(0, W - HR_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'])
# change color space if necessary
if self.opt['color']:
img_LR = util.channel_convert(C, self.opt['color'], [img_LR])[0] # TODO during val no definetion
# 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 __getitem__(self, index):
HR_path, LR_path = None, None
scale = self.opt['scale']
HR_size = self.opt['HR_size']
if HR_size:
LR_size = HR_size // scale
self.znorm = False # Default case: images are in the [0,1] range
if self.opt['znorm']:
if self.opt['znorm'] == True:
# Alternative: images are z-normalized to the [-1,1] range
self.znorm = True
######## Read the images ########
# Check if LR Path is provided
if self.paths_LR:
# If LR is provided, check if 'rand_flip_LR_HR' is enabled
if self.opt['rand_flip_LR_HR'] and self.opt['phase'] == 'train':
LRHRchance = random.uniform(0, 1)
if self.opt['flip_chance']:
flip_chance = self.opt['flip_chance']
else:
flip_chance = 0.05
#print("Random Flip Enabled")
# Normal case, no flipping:
else:
LRHRchance = 0.
flip_chance = 0.
#print("No Random Flip")
# get HR and LR images
# If enabled, random chance that LR and HR images are flipped
# Normal case, no flipping
# If img_LR (LR_path) doesn't exist, use img_HR (HR_path)
if LRHRchance < (1 - flip_chance):
HR_path = self.paths_HR[index]
LR_path = self.paths_LR[index]
if LR_path is None:
LR_path = HR_path
#print("HR kept")
# Flipped case:
# If img_HR (LR_path) doesn't exist, use img_HR (LR_path)
else:
HR_path = self.paths_LR[index]
LR_path = self.paths_HR[index]
if HR_path is None:
HR_path = LR_path
#print("HR flipped")
# Read the LR and HR images from the provided paths
img_LR = util.read_img(self.LR_env, LR_path, znorm=self.znorm)
img_HR = util.read_img(self.HR_env, HR_path, znorm=self.znorm)
# Even if LR dataset is provided, force to generate aug_downscale % of downscales OTF from HR
# The code will later make sure img_LR has the correct size
if self.opt['aug_downscale']:
aug_downscale = self.opt['aug_downscale']
if np.random.rand() < aug_downscale:
img_LR = img_HR
# If LR is not provided, use HR and modify on the fly
else:
HR_path = self.paths_HR[index]
img_HR = util.read_img(self.HR_env, HR_path, znorm=self.znorm)
img_LR = img_HR
######## Modify the images ########
# HR modcrop in the validation / test phase
if self.opt['phase'] != 'train':
img_HR = util.modcrop(img_HR, scale)
# change color space if necessary
# Note: Changing the LR colorspace here could make it so some colors are introduced when
# doing the augmentations later (ie: with Gaussian or Speckle noise), may be good if the
# model can learn to remove color noise in grayscale images, otherwise move to before
# converting to tensors
# self.opt['color'] For both LR and HR as in the the original code, kept for compatibility
# self.opt['color_HR'] and self.opt['color_LR'] for independent control
if self.opt['color_HR'] or self.opt['color']: # Only change HR
img_HR = util.channel_convert(img_HR.shape[2], self.opt['color'],
[img_HR])[0]
if self.opt['color_LR'] or self.opt['color']: # Only change LR
img_LR = util.channel_convert(img_LR.shape[2], self.opt['color'],
[img_LR])[0]
######## Augmentations ########
# Augmentations during training
if self.opt['phase'] == 'train':
# Validate there's an img_LR, if not, use img_HR
if img_LR is None:
img_LR = img_HR
print("Image LR: ", LR_path, (
"was not loaded correctly, using HR pair to downscale on the fly."
))
# Check that HR and LR have the same dimensions ratio, else, generate new LR from HR
if img_HR.shape[0] // img_LR.shape[0] != img_HR.shape[
1] // img_LR.shape[1]:
print(
"Warning: img_LR dimensions ratio does not match img_HR dimensions ratio for: ",
HR_path)
img_LR = img_HR
# Random Crop (reduce computing cost and adjust images to correct size first)
if img_HR.shape[0] > HR_size or img_HR.shape[1] > HR_size:
# Here the scale should be in respect to the images, not to the training scale (in case they are being scaled on the fly)
scaleor = img_HR.shape[0] // img_LR.shape[0]
img_HR, img_LR = augmentations.random_crop_pairs(
img_HR, img_LR, HR_size, scaleor)
# Or if the HR images are too small, Resize to the HR_size size and fit LR pair to LR_size too
if img_HR.shape[0] < HR_size or img_HR.shape[1] < HR_size:
print("Warning: Image: ", HR_path,
" size does not match HR size: (", HR_size,
"). The image size is: ", img_HR.shape)
# rescale HR image to the HR_size
img_HR, _ = augmentations.resize_img(np.copy(img_HR),
crop_size=(HR_size,
HR_size),
algo=cv2.INTER_LINEAR)
# rescale LR image to the LR_size (The original code discarded the img_LR and generated a new one on the fly from img_HR)
img_LR, _ = augmentations.resize_img(np.copy(img_LR),
crop_size=(LR_size,
LR_size),
algo=cv2.INTER_LINEAR)
# Randomly scale LR from HR during training if :
# - LR dataset is not provided
# - LR dataset is not in the correct scale
# - Also to check if LR is not at the correct scale already (if img_LR was changed to img_HR)
if img_LR.shape[0] != LR_size or img_LR.shape[1] != LR_size:
ds_algo = 777 # default to matlab-like bicubic downscale
# if manually set and scale algorithms are provided, then:
if self.opt['lr_downscale']:
if self.opt['lr_downscale_types']:
ds_algo = self.opt['lr_downscale_types']
else: # else, if for some reason img_LR is too large, default to matlab-like bicubic downscale
# if not self.opt['aug_downscale']: #only print the warning if not being forced to use HR images instead of LR dataset (which is a known case)
print(
"LR image is too large, auto generating new LR for: ",
LR_path)
img_LR, scale_interpol_algo = augmentations.scale_img(
img_LR, scale, algo=ds_algo)
if self.znorm:
# The generated LR sometimes get slightly out of the [-1,1] range
np.clip(img_LR, -1., 1., out=img_LR)
else:
# The generated LR sometimes get slightly out of the [0,1] range
np.clip(img_LR, 0., 1., out=img_LR)
# """
# Rotations. 'use_flip' = 180 or 270 degrees (mirror), 'use_rot' = 90 degrees, 'HR_rrot' = random rotations +-45 degrees
if (self.opt['use_flip']
or self.opt['use_rot']) and self.opt['hr_rrot']:
if np.random.rand() > 0.5:
img_LR, img_HR = util.augment([img_LR, img_HR],
self.opt['use_flip'],
self.opt['use_rot'])
else:
if np.random.rand(
) > 0.5: # randomize the random rotations, so half the images are the original
img_HR, img_LR = augmentations.random_rotate_pairs(
img_HR, img_LR, HR_size, scale)
elif (self.opt['use_flip']
or self.opt['use_rot']) and not self.opt['hr_rrot']:
# augmentation - flip, rotate
img_LR, img_HR = util.augment([img_LR, img_HR],
self.opt['use_flip'],
self.opt['use_rot'])
elif self.opt['hr_rrot']:
if np.random.rand(
) > 0.5: # randomize the random rotations, so half the images are the original
img_HR, img_LR = augmentations.random_rotate_pairs(
img_HR, img_LR, HR_size, scale)
# Final checks
# if the resulting HR image size so far is too large or too small, resize HR to the correct size and downscale to generate a new LR on the fly
if img_HR.shape[0] != HR_size or img_HR.shape[1] != HR_size:
print("Image: ", HR_path, " size does not match HR size: (",
HR_size, "). The image size is: ", img_HR.shape)
# rescale HR image to the HR_size
img_HR, _ = augmentations.resize_img(np.copy(img_HR),
crop_size=(HR_size,
HR_size),
algo=cv2.INTER_LINEAR)
# if manually provided and scale algorithms are provided, then:
if self.opt['lr_downscale_types']:
ds_algo = self.opt['lr_downscale_types']
else:
# using matlab imresize to generate LR pair
ds_algo = 777
img_LR, _ = augmentations.scale_img(img_HR,
scale,
algo=ds_algo)
# if the resulting LR so far does not have the correct dimensions, also generate a new HR-LR image pair on the fly
if img_LR.shape[0] != LR_size or img_LR.shape[0] != LR_size:
print("Image: ", LR_path, " size does not match LR size: (",
HR_size // scale, "). The image size is: ", img_LR.shape)
# rescale HR image to the HR_size (should not be needed, but something went wrong before, just for sanity)
img_HR, _ = augmentations.resize_img(np.copy(img_HR),
crop_size=(HR_size,
HR_size),
algo=cv2.INTER_LINEAR)
# if manually provided and scale algorithms are provided, then:
if self.opt['lr_downscale_types']:
ds_algo = self.opt['lr_downscale_types']
else:
# using matlab imresize to generate LR pair
ds_algo = 777
img_LR, _ = augmentations.scale_img(img_HR,
scale,
algo=ds_algo)
# Below are the LR On The Fly augmentations
# Add noise to HR if enabled AND noise types are provided (for noise2noise and similar)
if self.opt['hr_noise']:
if self.opt['hr_noise_types']:
img_HR, hr_noise_algo = augmentations.noise_img(
img_HR, noise_types=self.opt['hr_noise_types'])
else:
print(
"Noise types 'hr_noise_types' not defined. Skipping OTF noise for HR."
)
# Create color fringes
# Caution: Can easily destabilize a model
# Only applied to a small % of the images. Around 20% and 50% appears to be stable.
if self.opt['lr_fringes']:
lr_fringes_chance = self.opt['lr_fringes_chance'] if self.opt[
'lr_fringes_chance'] else 0.4
if np.random.rand() > (1. - lr_fringes_chance):
img_LR = augmentations.translate_chan(img_LR)
# """
# v LR blur AND blur types are provided, else will skip
if self.opt['lr_blur']:
if self.opt['lr_blur_types']:
img_LR, blur_algo, blur_kernel_size = augmentations.blur_img(
img_LR, blur_algos=self.opt['lr_blur_types'])
else:
print(
"Blur types 'lr_blur_types' not defined. Skipping OTF blur."
)
# """
# """
# v LR primary noise: Add noise to LR if enabled AND noise types are provided, else will skip
if self.opt['lr_noise']:
if self.opt['lr_noise_types']:
img_LR, noise_algo = augmentations.noise_img(
img_LR, noise_types=self.opt['lr_noise_types'])
else:
print(
"Noise types 'lr_noise_types' not defined. Skipping OTF noise."
)
# v LR secondary noise: Add additional noise to LR if enabled AND noise types are provided, else will skip
if self.opt['lr_noise2']:
if self.opt['lr_noise_types2']:
img_LR, noise_algo2 = augmentations.noise_img(
img_LR, noise_types=self.opt['lr_noise_types2'])
else:
print(
"Noise types 'lr_noise_types2' not defined. Skipping OTF secondary noise."
)
# """
# """
# v LR cutout / LR random erasing (for inpainting/classification tests)
if self.opt['lr_cutout'] and (self.opt['lr_erasing'] != True):
img_LR = augmentations.cutout(img_LR, img_LR.shape[0] // 2)
elif self.opt['lr_erasing'] and (self.opt['lr_cutout'] != True):
img_LR = augmentations.random_erasing(img_LR)
elif self.opt['lr_cutout'] and self.opt['lr_erasing']:
if np.random.rand(
) > 0.5: # only do cutout or erasing, not both at the same time
img_LR = augmentations.cutout(img_LR,
img_LR.shape[0] // 2,
p=0.5)
else:
img_LR = augmentations.random_erasing(img_LR,
p=0.5,
modes=[3])
# """
# Apply "auto levels" to images
# Randomize for augmentation
rand_levels = (1 - self.opt['rand_auto_levels']
) if self.opt['rand_auto_levels'] else 1
if self.opt['auto_levels'] and np.random.rand() > rand_levels:
if self.opt['auto_levels'] == 'HR':
img_HR = augmentations.simplest_cb(img_HR,
znorm=self.znorm)
elif self.opt['auto_levels'] == 'LR':
img_LR = augmentations.simplest_cb(img_LR,
znorm=self.znorm)
elif self.opt['auto_levels'] == True or self.opt[
'auto_levels'] == 'Both':
img_HR = augmentations.simplest_cb(img_HR,
znorm=self.znorm)
img_LR = augmentations.simplest_cb(img_LR,
znorm=self.znorm)
# Apply unsharpening mask to HR images
# img_HR1 = img_HR
# Randomize for augmentation
rand_unsharp = (
1 -
self.opt['rand_unsharp']) if self.opt['rand_unsharp'] else 1
if self.opt['unsharp_mask'] and np.random.rand() > rand_unsharp:
img_HR = augmentations.unsharp_mask(img_HR, znorm=self.znorm)
# For testing and validation
if self.opt['phase'] != 'train':
# Randomly downscale LR if enabled
if self.opt['lr_downscale']:
if self.opt['lr_downscale_types']:
img_LR, scale_interpol_algo = augmentations.scale_img(
img_LR, scale, algo=self.opt['lr_downscale_types'])
else: # Default to matlab-like bicubic downscale
img_LR, scale_interpol_algo = augmentations.scale_img(
img_LR, scale, algo=777)
# Alternative position for changing the colorspace of LR.
# if self.opt['color_LR']: # Only change LR
# img_LR = util.channel_convert(img_LR.shape[2], self.opt['color'], [img_LR])[0]
# Debug
# Save img_LR and img_HR images to a directory to visualize what is the result of the on the fly augmentations
# DO NOT LEAVE ON DURING REAL TRAINING
#self.output_sample_imgs = True
if self.opt['phase'] == 'train':
if self.output_sample_imgs:
import os
# LR_dir, im_name = os.path.split(LR_path)
HR_dir, im_name = os.path.split(HR_path)
baseHRdir, _ = os.path.split(HR_dir)
debugpath = os.path.join(baseHRdir, 'sampleOTFimgs')
print('debugpathhhhhhhh', debugpath)
# debugpath = os.path.join(os.path.split(LR_dir)[0], 'sampleOTFimgs')
# debugpath = os.path.join('D:/tmp_test', 'sampleOTFimgs')
# print(debugpath)
if not os.path.exists(debugpath):
os.makedirs(debugpath)
if self.opt[
'znorm']: # Back from [-1,1] range to [0,1] range for OpenCV2
img_LRn = (img_LR + 1.0) / 2.0
img_HRn = (img_HR + 1.0) / 2.0
# img_HRn1 = (img_HR1 + 1.0) / 2.0
else: # Already in the [0,1] range for OpenCV2
img_LRn = img_LR
img_HRn = img_HR
# img_HRn1 = img_HR1
import uuid
hex = uuid.uuid4().hex
# random name to save + had to multiply by 255, else getting all black image
cv2.imwrite(os.path.join(debugpath, im_name + hex + '_LR.png'),
img_LRn * 255)
# random name to save + had to multiply by 255, else getting all black image
cv2.imwrite(os.path.join(debugpath, im_name + hex + '_HR.png'),
img_HRn * 255)
# cv2.imwrite(debugpath+"\\"+im_name+hex+'_HR1.png',img_HRn1*255) #random name to save + had to multiply by 255, else getting all black image
######## Convert images to PyTorch Tensors ########
"""
if (img_HR.min() < -1):
print("HR.min :", img_HR.min())
print(HR_path)
if (img_HR.max() > 1):
print("HR.max :", img_HR.max())
print(HR_path)
if (img_LR.min() < -1):
print("LR.min :", img_LR.min())
print(LR_path)
if (img_LR.max() > 1):
print("LR.max :", img_LR.max())
print(LR_path)
#"""
# 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]]
# BGRA to RGBA, HWC to CHW, numpy to tensor
elif img_LR.shape[2] == 4:
img_HR = img_HR[:, :, [2, 1, 0, 3]]
img_LR = img_LR[:, :, [2, 1, 0, 3]]
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 __getitem__(self, index):
# self.block_size = 8
Uncomp_size = self.opt['patch_size']
# get Uncomp image
Uncomp_path = self.paths_Uncomp[index]
img_Uncomp = util.read_img(self.Uncomp_env, Uncomp_path)
# modcrop in the validation / test phase
if self.opt['phase'] != 'train':
img_Uncomp = util.modcrop(img_Uncomp, self.block_size)
# change color space if necessary
img_Uncomp = 255 * util.channel_convert(
img_Uncomp.shape[2],
'ycbcr' if 'chroma' in self.opt['mode'] else 'y', [img_Uncomp])[0]
if 'chroma' in self.opt['mode'] and img_Uncomp.shape[
2] == 1: #For the case of loading a grayscale image in JPEG format during chroma training:
img_Uncomp = np.tile(img_Uncomp, [1, 1, 3])
# img_Uncomp = util.channel_convert(img_Uncomp.shape[2], 'ycbcr', [img_Uncomp])[0]
# randomly scale during training
if self.opt['phase'] == 'train':
QF = self.quality_factors[np.random.choice(len(
self.quality_factors),
p=self.QF_probs)]
if isinstance(QF, list):
QF = np.random.randint(low=QF[0], high=QF[1])
# raise Exception('QF range is unsupported yet')
random_scale = random.choice(self.random_scale_list)
if random_scale != 1:
H_s, W_s, _ = img_Uncomp.shape
def _mod(n, random_scale, thres):
rlt = int(n * random_scale)
rlt = (rlt // self.block_size) * self.block_size
return thres if rlt < thres else rlt
H_s = _mod(H_s, random_scale, Uncomp_size)
W_s = _mod(W_s, random_scale, Uncomp_size)
img_Uncomp = cv2.resize(np.copy(img_Uncomp), (W_s, H_s),
interpolation=cv2.INTER_LINEAR)
# force to 3 channels
if img_Uncomp.ndim == 2:
img_Uncomp = np.expand_dims(img_Uncomp, -1)
# img_Uncomp = cv2.cvtColor(img_Uncomp, cv2.COLOR_GRAY2BGR)
else:
QF = self.per_index_QF[index]
H, W, _ = img_Uncomp.shape
# # using CEM imresize:
# img_LR = imresize(img_HR,scale_factor=[1/float(scale)],kernel=self.kernel)
if self.opt['phase'] == 'train':
# if the image size is too small
H, W, _ = img_Uncomp.shape
# randomly crop
rnd_h_Uncomp = random.randint(0, max(0, H - Uncomp_size))
rnd_w_Uncomp = random.randint(0, max(0, W - Uncomp_size))
img_Uncomp = img_Uncomp[rnd_h_Uncomp:rnd_h_Uncomp + Uncomp_size,
rnd_w_Uncomp:rnd_w_Uncomp + Uncomp_size, :]
# augmentation - flip, rotate
img_Uncomp = util.augment([img_Uncomp], self.opt['use_flip'],
self.opt['use_rot'])
img_Uncomp = img_Uncomp[0]
# BGR to RGB, HWC to CHW, numpy to tensor
# if img_Uncomp.shape[2] == 3:
# img_Uncomp = img_Uncomp[:, :, [2, 1, 0]]
img_Uncomp = torch.from_numpy(
np.ascontiguousarray(np.transpose(img_Uncomp, (2, 0, 1)))).float()
return {'Uncomp': img_Uncomp, 'Uncomp_path': Uncomp_path, 'QF': QF}
def __getitem__(self, index):
HR_path, LR_path = None, None
scale = self.opt['scale']
HR_size = self.opt['HR_size']
#v
if self.opt['rand_flip_LR_HR'] and self.LR_scale and self.opt['phase'] == 'train':
LRHRchance = random.uniform(0, 1)
if self.opt['flip_chance']:
flip_chance = self.opt['flip_chance']
else:
flip_chance = 0.05
#print("Random Flip Enabled")
else:
LRHRchance = 0.
flip_chance = 0.
#print("No Random Flip")
# get HR image
if LRHRchance < (1- flip_chance):
HR_path = self.paths_HR[index]
#print("HR kept")
else:
HR_path = self.paths_LR[index]
#print("HR flipped")
#v
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]
#v
if self.HR_crop and (self.HR_rrot != True):
crop_size = (HR_size, HR_size)
img_HR, _ = augmentations.random_resize_img(img_HR, crop_size)
elif self.HR_rrot and (self.HR_crop != True):
img_HR, _ = augmentations.random_rotate(img_HR)
elif self.HR_crop and self.HR_rrot:
if np.random.rand() > 0.5:
crop_size = (HR_size, HR_size)
img_HR, _ = augmentations.random_resize_img(img_HR, crop_size)
else:
img_HR, _ = augmentations.random_rotate(img_HR)
#v
#v
if self.HR_noise:
img_HR, hr_noise_algo = augmentations.noise_img(img_HR, self.hr_noise_types)
#v
# get LR image
if self.paths_LR:
if self.HR_crop or self.HR_rrot: #v
img_LR = img_HR
else:
if LRHRchance < (1- flip_chance):
LR_path = self.paths_LR[index]
#print("LR kept")
else:
LR_path = self.paths_HR[index]
#print("LR flipped")
img_LR = util.read_img(self.LR_env, LR_path)
#"""
#v scale
if self.LR_scale:
img_LR, scale_interpol_algo = augmentations.scale_img(img_LR, scale)
#"""
#"""
#v blur
if self.LR_blur:
img_LR, blur_algo, blur_kernel_size = augmentations.blur_img(img_LR, self.blur_algos)
#"""
#"""
#v noise
if self.LR_noise:
img_LR, noise_algo = augmentations.noise_img(img_LR, self.noise_types)
if self.LR_noise2:
img_LR, noise_algo2 = augmentations.noise_img(img_LR, self.noise_types2)
#"""
#"""
#v LR cutout / LR random erasing
if self.LR_cutout and (self.LR_erasing != True):
img_LR = augmentations.cutout(img_LR, img_LR.shape[0] // 2)
elif self.LR_erasing and (self.LR_cutout != True): #only do cutout or erasing, not both
img_LR = augmentations.random_erasing(img_LR)
elif self.LR_cutout and self.LR_erasing:
if np.random.rand() > 0.5:
img_LR = augmentations.cutout(img_LR, img_LR.shape[0] // 2, p=0.5)
else:
img_LR = augmentations.random_erasing(img_LR, p=0.5, modes=[3])
#"""
else: # down-sampling on-the-fly
# randomly scale during training
if self.opt['phase'] == 'train':
random_scale = random.choice(self.random_scale_list)
H_s, W_s, _ = img_HR.shape
def _mod(n, random_scale, scale, thres):
rlt = int(n * random_scale)
rlt = (rlt // scale) * scale
return thres if rlt < thres else rlt
H_s = _mod(H_s, random_scale, scale, HR_size)
W_s = _mod(W_s, random_scale, scale, HR_size)
img_HR = cv2.resize(np.copy(img_HR), (W_s, H_s), interpolation=cv2.INTER_LINEAR)
# 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 = util.imresize_np(img_HR, 1 / scale, True)
if img_LR.ndim == 2:
img_LR = np.expand_dims(img_LR, axis=2)
if self.opt['phase'] == 'train':
# if the image size is too small
H, W, _ = img_HR.shape
if H < HR_size or W < HR_size:
img_HR = cv2.resize(
np.copy(img_HR), (HR_size, HR_size), interpolation=cv2.INTER_LINEAR)
# using matlab imresize
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 = 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_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_LR, img_HR = util.augment([img_LR, img_HR], self.opt['use_flip'], \
self.opt['use_rot'])
# change color space if necessary
if self.opt['color']:
#img_LR = util.channel_convert(C, self.opt['color'], [img_LR])[0] # TODO during val no definetion
img_LR = util.channel_convert(img_LR.shape[2], self.opt['color'], [img_LR])[0] # v appears to work ok
# 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 __getitem__(self, index):
HR_path, LR_path = None, None
scale = self.opt['scale']
HR_size_w = self.opt['HR_size_w']
HR_size_h = self.opt['HR_size_h']
# 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]
# get LR image
if self.paths_LR:
LR_path = self.paths_LR[index]
img_LR = util.read_img(self.LR_env, LR_path)
if True:#self.opt['phase'] == 'train':
# if the image size is too small
H, W, _ = img_HR.shape
if H < HR_size_h or W < HR_size_w:
img_HR = cv2.resize(
np.copy(img_HR), (HR_size_w, HR_size_h), interpolation=cv2.INTER_LINEAR)
img_LR = cv2.resize(
np.copy(img_LR), (HR_size_w//scale, HR_size_h//scale), interpolation=cv2.INTER_LINEAR)
else:
H, W, C = img_LR.shape
LR_size_h = HR_size_h // scale
LR_size_w = HR_size_w // scale
# randomly crop
rnd_h = random.randint(0, max(0, H - LR_size_h))
rnd_w = random.randint(0, max(0, W - LR_size_w))
img_LR = img_LR[rnd_h:rnd_h + LR_size_h, rnd_w:rnd_w + LR_size_w, :]
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_h, rnd_w_HR:rnd_w_HR + HR_size_w, :]
#augmentation - flip, rotate
#img_LR, img_HR = util.augment([img_LR, img_HR], self.opt['use_flip'], \
# self.opt['use_rot'])
#print(img_LR.shape,img_HR.shape)
# change color space if necessary
if self.opt['color']:
img_LR = util.channel_convert(C, self.opt['color'], [img_LR])[0] # TODO during val no definetion
# 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 __getitem__(self, index):
# print("EJecuta GeT ITEM")
HR_path, LR_path = None, None
scale = 0 #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':
# print("Opcion: ", self.opt["phase"])
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]
# get LR image
if self.paths_LR:
LR_path = self.paths_LR[index]
img_LR = util.read_img(self.LR_env, LR_path)
else: # down-sampling on-the-fly
# randomly scale during training
if self.opt['phase'] == 'train':
random_scale = random.choice(self.random_scale_list)
H_s, W_s, _ = img_HR.shape
def _mod(n, random_scale, scale, thres):
rlt = int(n * random_scale)
rlt = (rlt // scale) * scale
return thres if rlt < thres else rlt
H_s = _mod(H_s, random_scale, scale, HR_size)
W_s = _mod(W_s, random_scale, scale, HR_size)
img_HR = cv2.resize(np.copy(img_HR), (W_s, H_s), interpolation=cv2.INTER_LINEAR)
# 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 = util.imresize_np(img_HR, 1 / scale, True)
if img_LR.ndim == 2:
img_LR = np.expand_dims(img_LR, axis=2)
if self.opt['phase'] == 'train':
# if the image size is too small
# print("1...........")
H, W, _ = img_HR.shape
if H < HR_size or W < HR_size:
img_HR = cv2.resize(
np.copy(img_HR), (HR_size, HR_size), interpolation=cv2.INTER_LINEAR)
# using matlab imresize
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 = 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_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_LR, img_HR = util.augment([img_LR, img_HR], self.opt['use_flip'], \
self.opt['use_rot'])
# change color space if necessary
if self.opt['color']:
img_LR = util.channel_convert(C, self.opt['color'], [img_LR])[0] # TODO during val no definetion
# 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()
# print("2.................................")
#
# print(img_HR.shape)
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 __getitem__(self, index):
if self.data_type == 'lmdb' and (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]
resolution = [int(s) for s in self.sizes_GT[index].split('_')
] if self.data_type == 'lmdb' else None
img_GT = util.read_img(self.GT_env, GT_path, resolution)
if self.opt['phase'] != 'train': # modcrop in the validation / test phase
img_GT = util.modcrop(img_GT, scale)
if self.opt['color']: # change color space if necessary
img_GT = util.channel_convert(img_GT.shape[2], self.opt['color'], [img_GT])[0]
# get LQ image
if self.paths_LQ:
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)
else: # down-sampling on-the-fly
# randomly scale during training
random_scale = random.choice(self.random_scale_list)
H_s, W_s, _ = img_GT.shape
def _mod(n, random_scale, scale, thres):
rlt = int(n * random_scale)
rlt = (rlt // scale) * scale
return thres if rlt < thres else rlt
H_s = _mod(H_s, random_scale, scale, GT_size)
W_s = _mod(W_s, random_scale, scale, GT_size)
img_GT = cv2.resize(img_GT, (W_s, H_s), interpolation=cv2.INTER_LINEAR)
if img_GT.ndim == 2:
img_GT = cv2.cvtColor(img_GT, cv2.COLOR_GRAY2BGR)
H, W, _ = img_GT.shape
# using matlab imresize
img_LQ = util.imresize_np(img_GT, 1 / scale, True)
if img_LQ.ndim == 2:
img_LQ = np.expand_dims(img_LQ, axis=2)
# get lr_large image
img_LQ_Large = cv2.resize(img_LQ, (GT_size, GT_size), cv2.INTER_CUBIC)
# get the pts, heatmaps and masks
f_anno = osp.join(self.landmarks_folder_256, GT_path + '.txt')
# load the landmarks
pts, point_set = util.anno_parser(f_anno, 68)
# if self.opt['phase'] == 'train':
# # if the image size is too small
# H, W, _ = img_GT.shape
# if H < GT_size or W < GT_size:
# img_GT = cv2.resize(img_GT, (GT_size, GT_size), interpolation=cv2.INTER_LINEAR)
# # using matlab imresize
# img_LQ = util.imresize_np(img_GT, 1 / scale, True)
# if img_LQ.ndim == 2:
# img_LQ = np.expand_dims(img_LQ, axis=2)
H, W, C = img_LQ.shape
# augmentation - flip, rotate
imgs = []
imgs.append(img_LQ)
imgs.append(img_GT)
rlt, pts = util.augment_imgs_landmarks(GT_size, imgs, pts, self.opt['use_flip'], self.opt['use_rot'])
img_LQ = rlt[0]
img_GT = rlt[-1]
if self.opt['color']: # change color space if necessary
img_LQ = util.channel_convert(C, self.opt['color'],
[img_LQ])[0] # TODO during val no definition
# 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]]
img_LQ_Large = img_LQ_Large[:, :, [2, 1, 0]]
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()
img_LQ_Large = torch.from_numpy(np.ascontiguousarray(np.transpose(img_LQ_Large, (2, 0, 1)))).float()
pts = util.apply_bound(pts, 256, 256)
# H*W*C
GT_heatmaps, GT_mask = util.generate_label_map(pts, 16, 16, self.sigma, 16.0, self.heatmap_type)
GT_heatmaps = torch.from_numpy(GT_heatmaps.transpose((2, 0, 1))).type(torch.FloatTensor)
GT_mask = torch.from_numpy(GT_mask.transpose((2,0,1))).type(torch.ByteTensor)
if LQ_path is None:
LQ_path = GT_path
return {'LQ': img_LQ, 'GT': img_GT, 'LQ_Large': img_LQ_Large, 'GT_heatmaps': GT_heatmaps, 'GT_mask': GT_mask}
def __getitem__(self, index):
if self.opt['data_type'] == 'lmdb':
if (self.GT_env is None) or (self.LR_env is None):
self._init_lmdb()
GT_path, LR_path = None, None
scale = self.opt['scale']
GT_size = self.opt['GT_size']
LR_size = self.opt['LR_size']
# get GT image
GT_path = self.GT_paths[index]
if self.opt['data_type'] == 'lmdb':
resolution = [int(s) for s in self.GT_sizes[index].split('_')]
else:
resolution = None
img_GT = util.read_img(
self.GT_env, GT_path,
resolution) #return: Numpy float32, HWC, BGR, [0,1]
# modcrop in the validation / test phase
if self.opt['phase'] != 'train':
img_GT = util.modcrop(img_GT, scale)
# get LR image
if self.LR_paths: # LR exist
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)
else:
## down-sampling on-the-fly, randomly scale during training
if self.opt['phase'] == 'train':
random_scale = random.choice(self.random_scale_list)
H_s, W_s, _ = img_GT.shape
def _mod(n, random_scale, scale, thres):
rlt = int(n * random_scale)
rlt = (rlt // scale) * scale
return thres if rlt < thres else rlt
H_s = _mod(H_s, random_scale, scale, GT_size)
W_s = _mod(W_s, random_scale, scale, GT_size)
img_GT = cv2.resize(np.copy(img_GT), (W_s, H_s),
interpolation=cv2.INTER_LINEAR)
# force to 3 channels
if img_GT.ndim == 2:
img_GT = cv2.cvtColor(img_GT, cv2.COLOR_GRAY2BGR)
if self.opt['phase'] == 'train':
H, W, C = img_LR.shape
assert LR_size == GT_size // scale, 'GT size does not match LR 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, :]
rnd_h_GT, rnd_w_GT = int(rnd_h * scale), int(rnd_w * scale)
img_GT = img_GT[rnd_h_GT:rnd_h_GT + GT_size,
rnd_w_GT:rnd_w_GT + GT_size, :]
# augmentation - flip, rotate
img_LR, img_GT = util.augment([img_LR, img_GT],
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] # TODO during val no definition
img_GT = util.channel_convert(img_GT.shape[2], self.opt['color'],
[img_GT])[0]
# BGR to RGB, HWC to CHW, numpy to tensor
if img_GT.shape[2] == 3:
img_GT = img_GT[:, :, [2, 1, 0]]
img_LR = img_LR[:, :, [2, 1, 0]]
img_GT = torch.from_numpy(
np.ascontiguousarray(np.transpose(img_GT, (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 = GT_path
return {
'LQ': img_LR,
'GT': img_GT,
'LQ_path': LR_path,
'GT_path': GT_path
}
def __getitem__(self, index):
# HR_path, LR_path = None, None fake real
scale = self.opt['scale']
HR_size = self.opt['HR_size']
index_fake, index_real = index, np.random.randint(0, len(self.paths_real_LR))
# get LR image
LR_fake_path = self.paths_fake_LR[index_fake]
LR_real_path = self.paths_real_LR[index_real]
# fake_w_path = self.paths_fake_weights[index_fake]
# real_w_path = self.paths_real_weights[index_real]
img_LR_fake = util.read_img(self.LR_env, LR_fake_path)
img_LR_real = util.read_img(self.LR_env, LR_real_path)
# fake_w = np.load(fake_w_path)[0].transpose((1, 2, 0))
# real_w = np.load(real_w_path)[0].transpose((1, 2, 0)) # 1, H, W
# fake_w = cv2.resize(fake_w, (img_LR_fake.shape[1], img_LR_fake.shape[0]), interpolation=cv2.INTER_LINEAR)
# real_w = cv2.resize(real_w, (img_LR_real.shape[1], img_LR_real.shape[0]), interpolation=cv2.INTER_LINEAR)
# fake_w = np.reshape(fake_w, list(fake_w.shape)+[1])
# real_w = np.reshape(real_w, list(real_w.shape)+[1]) # H, W, 1
# get HR image
HR_path = self.paths_HR[index_fake]
index_unpair = np.random.randint(0, len(self.paths_HR))
HR_unpair = self.paths_HR[index_unpair]
img_HR = util.read_img(self.HR_env, HR_path)
img_unpair_HR = util.read_img(self.HR_env, HR_unpair)
# 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]
if self.opt['phase'] == 'train':
# if the image size is too small
H, W, _ = img_HR.shape
if H < HR_size or W < HR_size:
img_HR = cv2.resize(np.copy(img_HR), (HR_size, HR_size), interpolation=cv2.INTER_LINEAR)
# using matlab imresize
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_fake.shape
H_r, W_r, C = img_LR_real.shape
LR_size = HR_size // scale
# randomly crop
rnd_h_fake = random.randint(0, max(0, H - LR_size))
rnd_w_fake = random.randint(0, max(0, W - LR_size))
rnd_h_real = random.randint(0, max(0, H_r - LR_size))
rnd_w_real = random.randint(0, max(0, W_r - LR_size))
img_LR_fake = img_LR_fake[rnd_h_fake:rnd_h_fake + LR_size, rnd_w_fake:rnd_w_fake + LR_size, :]
img_LR_real = img_LR_real[rnd_h_real:rnd_h_real + LR_size, rnd_w_real:rnd_w_real + LR_size, :]
# fake_w = fake_w[rnd_h_fake:rnd_h_fake + LR_size, rnd_w_fake:rnd_w_fake + LR_size, :]
# real_w = real_w[rnd_h_real:rnd_h_real + LR_size, rnd_w_real:rnd_w_real + LR_size, :]
H, W, C = img_HR.shape
H_real, W_real, C_real = img_unpair_HR.shape
rnd_h = int(rnd_h_fake*scale)
rnd_w = int(rnd_w_fake*scale)
rnd_h_real = random.randint(0, max(0, H_real - HR_size))
rnd_w_real = random.randint(0, max(0, W_real - HR_size))
img_HR = img_HR[rnd_h:rnd_h + HR_size, rnd_w:rnd_w + HR_size, :]
img_unpair_HR = img_unpair_HR[rnd_h_real:rnd_h_real + HR_size, rnd_w_real:rnd_w_real + HR_size, :]
# augmentation - flip, rotate
img_LR_fake, img_LR_real, img_HR, img_unpair_HR \
= util.augment([img_LR_fake, img_LR_real, img_HR, img_unpair_HR],
self.opt['use_flip'], self.opt['use_rot'])
# if self.paths_real_weights:
# real_w = util.augment([real_w],
# self.opt['use_flip'], self.opt['use_rot'])
# change color space if necessary
if self.opt['color']:
img_LR = util.channel_convert(C, self.opt['color'], [img_LR])[0] # TODO during val no definetion
# BGR to RGB, HWC to CHW, numpy to tensor
if img_HR.shape[2] == 3:
img_HR = img_HR[:, :, [2, 1, 0]]
img_unpair_HR = img_unpair_HR[:, :, [2, 1, 0]]
img_LR_real = img_LR_real[:, :, [2, 1, 0]]
img_LR_fake = img_LR_fake[:, :, [2, 1, 0]]
img_HR = torch.from_numpy(np.ascontiguousarray(np.transpose(img_HR, (2, 0, 1)))).float()
img_unpair_HR = torch.from_numpy(np.ascontiguousarray(np.transpose(img_unpair_HR, (2, 0, 1)))).float()
img_LR_real = torch.from_numpy(np.ascontiguousarray(np.transpose(img_LR_real, (2, 0, 1)))).float()
img_LR_fake = torch.from_numpy(np.ascontiguousarray(np.transpose(img_LR_fake, (2, 0, 1)))).float()
return {'LR_real': img_LR_real, 'LR_fake': img_LR_fake, 'HR': img_HR, 'HR_unpair': img_unpair_HR,
'LR_real_path': LR_real_path, 'LR_fake_path': LR_fake_path, 'HR_path': HR_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)
# change color space if necessary
if self.opt['color']:
img_GT = util.channel_convert(img_GT.shape[2], self.opt['color'],
[img_GT])[0]
# get LQ image
if self.paths_LQ:
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)
else: # down-sampling on-the-fly
# randomly scale during training
if self.opt['phase'] == 'train':
random_scale = random.choice(self.random_scale_list)
H_s, W_s, _ = img_GT.shape
def _mod(n, random_scale, scale, thres):
rlt = int(n * random_scale)
rlt = (rlt // scale) * scale
return thres if rlt < thres else rlt
H_s = _mod(H_s, random_scale, scale, GT_size)
W_s = _mod(W_s, random_scale, scale, GT_size)
img_GT = cv2.resize(np.copy(img_GT), (W_s, H_s),
interpolation=cv2.INTER_LINEAR)
# force to 3 channels
if img_GT.ndim == 2:
img_GT = cv2.cvtColor(img_GT, cv2.COLOR_GRAY2BGR)
H, W, _ = img_GT.shape
# using matlab imresize
img_LQ = util.imresize_np(img_GT, 1 / scale, True)
if img_LQ.ndim == 2:
img_LQ = np.expand_dims(img_LQ, axis=2)
if self.opt['phase'] == 'train':
# if the image size is too small
H, W, _ = 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)
# using matlab imresize
img_LQ = util.imresize_np(img_GT, 1 / scale, True)
if img_LQ.ndim == 2:
img_LQ = np.expand_dims(img_LQ, axis=2)
H, W, C = img_LQ.shape
LQ_size = GT_size // scale
# randomly crop
rnd_h = random.randint(0, max(0, H - LQ_size))
rnd_w = random.randint(0, max(0, W - LQ_size))
img_LQ = img_LQ[rnd_h:rnd_h + LQ_size, rnd_w:rnd_w + LQ_size, :]
rnd_h_GT, rnd_w_GT = int(rnd_h * scale), int(rnd_w * scale)
img_GT = img_GT[rnd_h_GT:rnd_h_GT + GT_size,
rnd_w_GT:rnd_w_GT + GT_size, :]
# augmentation - flip, rotate
img_LQ, img_GT = util.augment([img_LQ, img_GT],
self.opt['use_flip'],
self.opt['use_rot'])
# change color space if necessary
if self.opt['color']:
img_LQ = util.channel_convert(
C, self.opt['color'],
[img_LQ])[0] # TODO during val no definition
# 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]]
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):
# HR_path, LR_path = None, None fake real
isReal = False
scale = self.opt['scale']
HR_size = self.opt['HR_size']
# print(self.paths_LR[-10:])
# print(self.paths_HR[-10:])
# get LR image
LR_path = self.paths_LR[index] # len(LR) > len(HR)
img_LR = util.read_img(self.LR_env, LR_path)
# get HR image
if self.opt['prefix'] not in os.path.basename(LR_path):
HR_path = self.paths_HR[index]
img_HR = util.read_img(self.HR_env, HR_path)
weights_path = self.paths_weights[index] # load domain distance weights .npy
img_weights = np.load(weights_path)[0].transpose((1, 2, 0))
img_weights = cv2.resize(img_weights, (img_HR.shape[1], img_HR.shape[0]), interpolation=cv2.INTER_LINEAR)
else:
isReal = True
numofHR = len(self.paths_HR)
index_HR = np.random.randint(0, numofHR)
HR_path = self.paths_HR[index_HR]
img_HR = util.read_img(self.HR_env, HR_path)
img_weights = np.ones_like(img_HR)
# 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]
if self.opt['phase'] == 'train':
# if the image size is too small
H, W, _ = img_HR.shape
if H < HR_size or W < HR_size:
img_HR = cv2.resize(np.copy(img_HR), (HR_size, HR_size), interpolation=cv2.INTER_LINEAR)
# using matlab imresize
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
# Href, Wref, Cref = img_TransRefer.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_LR = img_LR[rnd_h:rnd_h + LR_size, rnd_w:rnd_w + LR_size, :]
if isReal:
H_real, W_real, C_real = img_HR.shape
rnd_h_real = random.randint(0, max(0, H_real - HR_size))
rnd_w_real = random.randint(0, max(0, W_real - HR_size))
img_HR = img_HR[rnd_h_real:rnd_h_real + HR_size, rnd_w_real:rnd_w_real + HR_size, :]
img_weights = img_weights[rnd_h_real:rnd_h_real + HR_size, rnd_w_real:rnd_w_real + HR_size, :]
else:
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, :]
img_weights = img_weights[rnd_h_HR:rnd_h_HR + HR_size, rnd_w_HR:rnd_w_HR + HR_size, :]
# rnd_href = random.randint(0, max(0, Href - HR_size))
# rnd_wref = random.randint(0, max(0, Wref - HR_size))
# img_TransRefer = img_TransRefer[rnd_href:rnd_href + HR_size, rnd_wref:rnd_wref + HR_size, :]
# augmentation - flip, rotate
img_LR, img_HR, img_weights = util.augment([img_LR, img_HR, img_weights], self.opt['use_flip'], \
self.opt['use_rot'])
# change color space if necessary
if self.opt['color']:
img_LR = util.channel_convert(C, self.opt['color'], [img_LR])[0] # TODO during val no definetion
# 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_weights = torch.from_numpy(np.ascontiguousarray(np.transpose(img_weights, (2, 0, 1)))).float()
img_LR = torch.from_numpy(np.ascontiguousarray(np.transpose(img_LR, (2, 0, 1)))).float()
# img_TransRefer = torch.from_numpy(np.ascontiguousarray(np.transpose(img_TransRefer, (2, 0, 1)))).float()
# Wavelet Trans
if self.opt['phase'] == 'train':
img_HR = img_HR.reshape([1, img_HR.shape[0], img_HR.shape[1], img_HR.shape[2]])
LL, H_cat = self.DWT2(img_HR)
LL = LL[0]*0.5
H_cat = H_cat[0][0]*0.5+0.5
w_c, C, H, W = H_cat.shape
H_cat = H_cat.reshape([w_c*C, H, W])
img_HR = img_HR[0].detach()
if LR_path is None:
LR_path = HR_path
return {'LR': img_LR, 'HR': img_HR, 'weights': img_weights,'LR_path': LR_path, 'HR_path': HR_path,
"isReal": isReal, 'wt_LL': LL.detach(), 'Ht_cat': H_cat.detach()}
def __getitem__(self, index):
if self.data_type == 'lmdb' and (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']
img_GT_mask = None
# get GT image
GT_path = self.paths_GT[index]
resolution = [int(s) for s in self.sizes_GT[index].split('_')
] if self.data_type == 'lmdb' else None
img_GT = util.read_img(self.GT_env, GT_path, resolution)
if self.opt[
'phase'] != 'train': # modcrop in the validation / test phase
img_GT = util.modcrop(img_GT, scale)
if self.opt['color']: # change color space if necessary
img_GT = util.channel_convert(img_GT.shape[2], self.opt['color'],
[img_GT])[0]
# get LQ image
if self.paths_LQ:
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)
else: # down-sampling on-the-fly
# randomly scale during training
if self.opt['phase'] == 'train':
random_scale = random.choice(self.random_scale_list)
H_s, W_s, _ = img_GT.shape
def _mod(n, random_scale, scale, thres):
rlt = int(n * random_scale)
rlt = (rlt // scale) * scale
return thres if rlt < thres else rlt
H_s = _mod(H_s, random_scale, scale, GT_size)
W_s = _mod(W_s, random_scale, scale, GT_size)
img_GT = cv2.resize(img_GT, (W_s, H_s),
interpolation=cv2.INTER_LINEAR)
if img_GT.ndim == 2:
img_GT = cv2.cvtColor(img_GT, cv2.COLOR_GRAY2BGR)
H, W, _ = img_GT.shape
# using matlab imresize
img_LQ = util.imresize_np(img_GT, 1 / scale, True)
if img_LQ.ndim == 2:
img_LQ = np.expand_dims(img_LQ, axis=2)
if self.opt['phase'] == 'train':
# if the image size is too small
H, W, _ = img_GT.shape
if H < GT_size or W < GT_size:
img_GT = cv2.resize(img_GT, (GT_size, GT_size),
interpolation=cv2.INTER_LINEAR)
# using matlab imresize
img_LQ = util.imresize_np(img_GT, 1 / scale, True)
if img_LQ.ndim == 2:
img_LQ = np.expand_dims(img_LQ, axis=2)
H, W, C = img_LQ.shape
LQ_size = GT_size // scale
# randomly crop
rnd_h = random.randint(0, max(0, H - LQ_size))
rnd_w = random.randint(0, max(0, W - LQ_size))
img_LQ = img_LQ[rnd_h:rnd_h + LQ_size, rnd_w:rnd_w + LQ_size, :]
rnd_h_GT, rnd_w_GT = int(rnd_h * scale), int(rnd_w * scale)
img_GT = img_GT[rnd_h_GT:rnd_h_GT + GT_size,
rnd_w_GT:rnd_w_GT + GT_size, :]
# augmentation - flip, rotate
img_LQ, img_GT = util.augment([img_LQ, img_GT],
self.opt['use_flip'],
self.opt['use_rot'])
# create edge mask for GT
low_threshold = 25
ratio = 3
kernel_size = 5
img_GT_grayscale = cv2.cvtColor(img_GT, cv2.COLOR_BGR2GRAY)
img_GT_blurred = cv2.blur(img_GT_grayscale, (kernel_size, kernel_size))
img_GT_blurred = 255 * img_GT_blurred # Now scale by 255
img_GT_blurred = img_GT_blurred.astype(np.uint8)
detected_edges = cv2.Canny(img_GT_blurred, low_threshold,
low_threshold * ratio, kernel_size)
mask = detected_edges != 0
img_GT_mask = img_GT_grayscale * (mask[:, :].astype(
img_GT_grayscale.dtype))
img_GT_mask = cv2.GaussianBlur(img_GT_mask, (3, 3), 0)
img_GT_mask = np.expand_dims(img_GT_mask, 0)
# cv2.imshow("window", img_GT_mask)
# cv2.waitKey()
if self.opt['color']: # change color space if necessary
img_LQ = util.channel_convert(
C, self.opt['color'],
[img_LQ])[0] # TODO during val no definition
# 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]]
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,
'GT_mask': img_GT_mask
}
def __getitem__(self, index):
HR_path, LR_path = None, None
scale = self.opt['scale']
HR_size = self.opt['HR_size']
if HR_size:
LR_size = HR_size // scale
# Check if LR Path is provided
if self.paths_LR:
#If LR is provided, check if 'rand_flip_LR_HR' is enabled (Only will work if HR and LR images have the same initial size) during training
if self.opt['rand_flip_LR_HR'] and (
self.LR_scale
or scale == 1) and self.opt['phase'] == 'train':
LRHRchance = random.uniform(0, 1)
if self.opt['flip_chance']:
flip_chance = self.opt['flip_chance']
else:
flip_chance = 0.05
#print("Random Flip Enabled")
# Normal case, no flipping:
else:
LRHRchance = 0.
flip_chance = 0.
#print("No Random Flip")
# get HR and LR images
# If enabled, random chance that LR and HR images are flipped
# Normal case, no flipping
# If img_LR (LR_path) doesn't exist, use img_HR (HR_path)
if LRHRchance < (1 - flip_chance):
HR_path = self.paths_HR[index]
LR_path = self.paths_LR[index]
if LR_path is None:
LR_path = HR_path
#print("HR kept")
# Flipped case:
# If img_HR (LR_path) doesn't exist, use img_HR (LR_path)
else:
HR_path = self.paths_LR[index]
LR_path = self.paths_HR[index]
if HR_path is None:
HR_path = LR_path
#print("HR flipped")
# Read the LR and HR images from the provided paths
img_LR = util.read_img(self.LR_env, LR_path)
img_HR = util.read_img(self.HR_env, HR_path)
# Even if LR dataset is provided, force to generate aug_downscale % of downscales OTF from HR
# The code will later make sure img_LR has the correct size
if self.opt['aug_downscale']:
aug_downscale = self.opt['aug_downscale']
if np.random.rand() < aug_downscale:
img_LR = img_HR
# If LR is not provided, use HR and modify on the fly
else:
HR_path = self.paths_HR[index]
img_HR = util.read_img(self.HR_env, HR_path)
img_LR = img_HR
# HR 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]
img_LR = util.channel_convert(img_LR.shape[2], self.opt['color'],
[img_LR])[0]
#Augmentations during training
if self.opt['phase'] == 'train':
#Validate there's an img_LR,
if img_LR is None:
img_LR = img_HR
print("Image LR: ", LR_path, (
"was not loaded correctly, using HR pair to downscale on the fly."
))
#Random Crop (reduce computing cost and adjust images to correct size first)
if img_HR.shape[0] > HR_size or img_HR.shape[1] > HR_size:
#Here the scale should be in respect to the images, not to the training scale (in case they are being scaled on the fly)
if img_HR.shape[0] // img_LR.shape[0] is not img_HR.shape[
1] // img_LR.shape[1]:
print(
"Warning: img_LR dimensions ratio does not match img_HR dimensions ratio for: ",
HR_path)
img_LR = img_HR
scaleor = img_HR.shape[0] // img_LR.shape[0]
img_HR, img_LR = augmentations.random_crop_pairs(
img_HR, img_LR, HR_size, scaleor)
#or if the HR images are too small, Resize to the HR_size size and fit LR pair to LR_size too
if img_HR.shape[0] < HR_size or img_HR.shape[1] < HR_size:
print("Warning: Image: ", HR_path,
" size does not match HR size: (", HR_size,
"). The image size is: ", img_HR.shape)
# rescale HR image to the HR_size
img_HR, _ = augmentations.resize_img(np.copy(img_HR),
crop_size=(HR_size,
HR_size),
algo=cv2.INTER_LINEAR)
# rescale LR image to the LR_size
img_LR, _ = augmentations.resize_img(np.copy(img_LR),
crop_size=(LR_size,
LR_size),
algo=cv2.INTER_LINEAR)
#"""
# randomly scale LR from HR during training if LR dataset is not provided
# Also check if LR is not at the correct scale already
if img_LR.shape[0] is not LR_size and img_LR.shape[
1] is not LR_size:
if self.LR_scale: # if manually provided and scale algorithms are provided, then:
if self.scale_algos:
ds_algo = self.scale_algos
else:
ds_algo = 777
else: # else, if for some reason img_LR is too large, default to matlab-like bicubic downscale
#if not self.opt['aug_downscale']: #only print the warning if not being forced to use HR images instead of LR dataset (which is a known case)
ds_algo = 777
print(
"LR image is too large, auto generating new LR for: ",
LR_path)
img_LR, scale_interpol_algo = augmentations.scale_img(
img_LR, scale, algo=ds_algo)
#"""
# Rotations. 'use_flip' = 180 or 270 degrees (mirror), 'use_rot' = 90 degrees, 'HR_rrot' = random rotations +-45 degrees
if self.opt['use_flip'] and self.opt['use_rot'] and self.HR_rrot:
if np.random.rand() > 0.5:
img_LR, img_HR = util.augment([img_LR, img_HR], self.opt['use_flip'], \
self.opt['use_rot'])
else:
if np.random.rand(
) > 0.5: # randomize the random rotations, so half the images are the original
img_HR, img_LR = augmentations.random_rotate_pairs(
img_HR, img_LR, HR_size, scale)
elif (self.opt['use_flip']
or self.opt['use_rot']) and not self.HR_rrot:
# augmentation - flip, rotate
img_LR, img_HR = util.augment([img_LR, img_HR], self.opt['use_flip'], \
self.opt['use_rot'])
elif self.HR_rrot:
if np.random.rand(
) > 0.5: # randomize the random rotations, so half the images are the original
img_HR, img_LR = augmentations.random_rotate_pairs(
img_HR, img_LR, HR_size, scale)
# Final checks
# if the resulting HR image size so far is too large or too small, resize HR to the correct size and downscale to generate a new LR on the fly
if img_HR.shape[0] is not HR_size or img_HR.shape[1] is not HR_size:
print("Image: ", HR_path, " size does not match HR size: (",
HR_size, "). The image size is: ", img_HR.shape)
# rescale HR image to the HR_size
img_HR, _ = augmentations.resize_img(np.copy(img_HR),
crop_size=(HR_size,
HR_size),
algo=cv2.INTER_LINEAR)
if self.scale_algos: # if manually provided and scale algorithms are provided, then:
ds_algo = self.scale_algos
else:
## using matlab imresize to generate LR pair
ds_algo = 777
img_LR, _ = augmentations.scale_img(img_HR,
scale,
algo=ds_algo)
# Final checks
# if the resulting LR so far does not have the correct dimensions, also generate a new HR- LR image pair on the fly
if img_LR.shape[0] is not LR_size or img_LR.shape[0] is not LR_size:
print("Image: ", LR_path, " size does not match LR size: (",
HR_size // scale, "). The image size is: ", img_LR.shape)
# rescale HR image to the HR_size
img_HR, _ = augmentations.resize_img(np.copy(img_HR),
crop_size=(HR_size,
HR_size),
algo=cv2.INTER_LINEAR)
if self.scale_algos: # if manually provided and scale algorithms are provided, then:
ds_algo = self.scale_algos
else:
## using matlab imresize to generate LR pair
ds_algo = 777
img_LR, _ = augmentations.scale_img(img_HR,
scale,
algo=ds_algo)
# Add noise to HR if enabled
if self.HR_noise:
img_HR, hr_noise_algo = augmentations.noise_img(
img_HR, noise_types=self.hr_noise_types)
# Below are the LR On The Fly augmentations
#"""
#v LR blur
if self.LR_blur:
img_LR, blur_algo, blur_kernel_size = augmentations.blur_img(
img_LR, blur_algos=self.blur_algos)
#"""
#"""
#v LR primary noise
if self.LR_noise:
img_LR, noise_algo = augmentations.noise_img(
img_LR, noise_types=self.noise_types)
#v LR secondary noise
if self.LR_noise2:
img_LR, noise_algo2 = augmentations.noise_img(
img_LR, noise_types=self.noise_types2)
#"""
#"""
#v LR cutout / LR random erasing
if self.LR_cutout and (self.LR_erasing != True):
img_LR = augmentations.cutout(img_LR, img_LR.shape[0] // 2)
elif self.LR_erasing and (self.LR_cutout != True
): #only do cutout or erasing, not both
img_LR = augmentations.random_erasing(img_LR)
elif self.LR_cutout and self.LR_erasing:
if np.random.rand() > 0.5:
img_LR = augmentations.cutout(img_LR,
img_LR.shape[0] // 2,
p=0.5)
else:
img_LR = augmentations.random_erasing(img_LR,
p=0.5,
modes=[3])
#"""
#For testing and validation
if self.opt['phase'] != 'train':
#"""
#v randomly downscale LR if enabled
if self.LR_scale:
img_LR, scale_interpol_algo = augmentations.scale_img(
img_LR, scale, algo=self.scale_algos)
#"""
# Debug
# Save img_LR and img_HR images to a directory to visualize what is the result of the on the fly augmentations
# DO NOT LEAVE ON DURING REAL TRAINING
# self.output_sample_imgs = True
if self.opt['phase'] == 'train':
if self.output_sample_imgs:
import os
LR_dir, im_name = os.path.split(LR_path)
#baseHRdir, _ = os.path.split(HR_dir)
#debugpath = os.path.join(baseHRdir, os.sep, 'sampleOTFimgs')
debugpath = os.path.join(
os.path.split(LR_dir)[0], 'sampleOTFimgs')
#print(debugpath)
if not os.path.exists(debugpath):
os.makedirs(debugpath)
import uuid
hex = uuid.uuid4().hex
cv2.imwrite(
debugpath + "\\" + im_name + hex + '_LR.png', img_LR * 255
) #random name to save + had to multiply by 255, else getting all black image
cv2.imwrite(
debugpath + "\\" + im_name + hex + '_HR.png', img_HR * 255
) #random name to save + had to multiply by 255, else getting all black image
# 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]]
# BGRA to RGBA, HWC to CHW, numpy to tensor
elif img_LR.shape[2] == 4:
img_HR = img_HR[:, :, [2, 1, 0, 3]]
img_LR = img_LR[:, :, [2, 1, 0, 3]]
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
}