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 __getitem__(self, index):
# 从文件中读出HR图片，返回HR、LR图片对
HR_path, LR_path = None, None
HR_path = self.HR_paths[index]
scale = self.scale[self.idx_scale]
patch_size = self.patch_size
img_HR = util.read_img_(HR_path, self.n_colors)
# 归一化
img_HR = util.unit2single(img_HR)
img_HR = util.modcrop(img_HR, self.scale)
if self.opt['phase'] == 'train':
l_max = 50
theta = np.pi * np.random.rand(1)
l1 = 0.1 + l_max * np.random.rand(1)
l2 = 0.1 + (l1 - 0.1) * np.random.rand(1)
kernel = util.anisotropic_Gaussian(ksize=self.ksize,
theta=theta[0],
l1=l1[0],
l2=l2[0])
else:
kernel = util.anisotropic_Gaussian(ksize=self.ksize,
theta=np.pi,
l1=0.1,
l2=0.1)
k = np.reshape(kernel, (-1), order="F")
k_reduced = np.dot(self.p, k)
k_reduced = torch.from_numpy(k_reduced).float()
H, W, _ = img_HR.shape
img_LR = util.srmd_degradation(img_HR, kernel, self.scale)
if not self.opt.test_only:
patch_H, patch_L = util.paired_random_crop(img_HR, img_LR,
patch_size, scale)
# augmentation - flip, rotate
img_HR, img_LR = util.augment([patch_H, patch_L],
self.opt.use_flip, self.opt.use_rot)
if random.random < 0.1:
noise_level = np.zeros(1).float()
else:
noise_level = np.random.uniform(0, 15)
else:
noise_level = np.zeros(self.sigma_test).float
img_LR = img_LR + np.random.normal(0, noise_level / 255., img_LR.shape)
img_HR = util.img_single2tensor(img_HR)
img_LR = util.img_single2tensor(img_LR)
M_vector = torch.cat((k_reduced, noise_level),
0).unsqueeze(1).unsqueeze(1)
M = M_vector.repeat(1, img_LR.size()[-2], img_LR.size()[-1])
img_LR = torch.cat((img_LR, M), 0)
if LR_path is None:
LR_path = HR_path
return img_LR, img_HR, LR_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)
# 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):
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 __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):
# 从文件中读出HR图片，返回HR、LR图片对
HR_path, LR_path = None, None
HR_path = self.HR_paths[index]
img_HR = util.read_img_(HR_path, self.n_colors)
# 归一化
img_HR = util.unit2single(img_HR)
img_HR = util.modcrop(img_HR, self.scale)
kernel = fspecial('gaussian', 15, 2.6)
H, W, _ = img_HR.shape
img_LR = util.srmd_degradation(img_HR, kernel, self.scale)
if self.opt.phase == 'train':
patch_H, patch_L = util.paired_random_crop(img_HR, img_LR,
self.patch_size,
self.scale)
# augmentation - flip, rotate
img_HR, img_LR = util.augment([patch_H, patch_L],
self.opt.use_flip, self.opt.use_rot)
img_HR = util.img_single2tensor(img_HR)
img_LR = util.img_single2tensor(img_LR)
if random.random() < 0.1:
noise_level = torch.zeros(1).float()
else:
noise_level = torch.FloatTensor([
np.random.uniform(self.sigma_min, self.sigma_max)
]) / 255.0
else:
img_HR = util.img_single2tensor(img_HR)
img_LR = util.img_single2tensor(img_LR)
noise_level = torch.FloatTensor([self.sigma_test])
noise = torch.randn(img_LR.size()).mul_(noise_level).float()
img_LR.add_(noise)
kernel = util.single2tensor3(np.expand_dims(np.float32(kernel),
axis=2))
noise_level = torch.FloatTensor([noise_level]).view([1, 1, 1])
if LR_path is None:
LR_path = HR_path
return {
'LR': img_LR,
'HR': img_HR,
'k': kernel,
'sigma': noise_level,
'sf': self.scale,
'LR_path': LR_path,
'HR_path': HR_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)
# modcrop
img_LR = util.modcrop(img_LR, 2)
# 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 _scan(self):
path = self.dir_hr
names_hr = []
for dirpath, _, fnames in sorted(os.walk(path)):
for fname in sorted(fnames):
img_path = dirpath + '/' + fname
names_hr.append(img_path)
# 读取HR图片，生成 LR/BlurMap
names_lr = [[] for _ in self.scale]
for f in names_hr:
img_HR = util.read_img_(f, self.n_colors)
# 归一化
img_HR = util.unit2single(img_HR)
img_HR = util.modcrop(img_HR, self.scale)
if self.opt['phase'] == 'train':
l_max = 50
theta = np.pi * np.random.rand(1)
l1 = 0.1 + l_max * np.random.rand(1)
l2 = 0.1 + (l1 - 0.1) * np.random.rand(1)
kernel = util.anisotropic_Gaussian(ksize=self.ksize,
theta=theta[0],
l1=l1[0],
l2=l2[0])
else:
kernel = util.anisotropic_Gaussian(ksize=self.ksize,
theta=np.pi,
l1=0.1,
l2=0.1)
k = np.reshape(kernel, (-1), order="F")
k_reduced = np.dot(self.p, k)
k_reduced = torch.from_numpy(k_reduced).float()
H, W, _ = img_HR.shape
img_LR = util.srmd_degradation(img_HR, kernel, self.scale)
if not self.opt.test_only:
patch_H, patch_L = util.paired_random_crop(
img_HR, img_LR, patch_size, scale)
# augmentation - flip, rotate
img_HR, img_LR = util.augment([patch_H, patch_L],
self.opt.use_flip,
self.opt.use_rot)
if random.random < 0.1:
noise_level = np.zeros(1).float()
else:
noise_level = np.random.uniform(0, 15)
else:
noise_level = np.zeros(self.sigma_test).float
img_LR = img_LR + np.random.normal(0, noise_level / 255.,
img_LR.shape)
img_HR = util.img_single2tensor(img_HR)
img_LR = util.img_single2tensor(img_LR)
M_vector = torch.cat((k_reduced, noise_level),
0).unsqueeze(1).unsqueeze(1)
M = M_vector.repeat(1, img_LR.size()[-2], img_LR.size()[-1])
img_LR = torch.cat((img_LR, M), 0)
filename, _ = os.path.splitext(os.path.basename(f))
for si, s in enumerate(self.scale):
# 1018测试的图片格式与训练图片格式不符，记得修改！
if self.name == 'DIV2K':
names_lr[si].append(self.dir_lr + '/' + 'X{}'.format(s) +
'/' + 'sig2.6_' + filename + '.jpg')
elif self.name == 'Set5':
names_lr[si].append(self.dir_lr + '/' + 'X{}'.format(s) +
'/' + 'sig2.6_' + filename + '.png')
return names_hr, names_lr
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
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['color']: # change color space if necessary
img_GT = util.channel_convert(img_GT.shape[2], self.opt['color'],
[img_GT])[0]
# get LQ image
LQ_path = self.paths_LQ[index]
resolution = [int(s) for s in self.sizes_LQ[index].split('_')
] if self.data_type == 'lmdb' else None
img_LQ = util.read_img(self.LQ_env, LQ_path, resolution)
# get condition vector
LQ_name = os.path.basename(LQ_path)
cond_str = re.split(r"[_,.]", LQ_name)[-2]
cond_num_list = self.opt['cond_norm']
cond_list = []
for ind in range(len(cond_num_list)):
cond = int(cond_str[ind * 2:(ind + 1) * 2]) / cond_num_list[ind]
cond_list.append(cond)
# modcrop in the validation / test phase
if self.opt['phase'] != 'train':
img_GT = util.modcrop(img_GT, 2)
img_LQ = util.modcrop(img_LQ, 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)
img_LQ = cv2.resize(img_LQ, (GT_size, GT_size),
interpolation=cv2.INTER_LINEAR)
# randomly crop
rnd_h = random.randint(0, max(0, H - GT_size))
rnd_w = random.randint(0, max(0, W - GT_size))
img_LQ = img_LQ[rnd_h:rnd_h + GT_size, rnd_w:rnd_w + GT_size, :]
img_GT = img_GT[rnd_h:rnd_h + GT_size, rnd_w:rnd_w + GT_size, :]
# augmentation - flip, rotate
img_LQ, img_GT = util.augment([img_LQ, img_GT],
self.opt['use_flip'],
self.opt['use_rot'])
if self.opt['color']: # change color space if necessary
_, _, C = img_LQ.shape
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()
cond = torch.Tensor(cond_list).view(len(cond_list))
if LQ_path is None:
LQ_path = GT_path
return {
'LQ': img_LQ,
'GT': img_GT,
'cond': cond,
'LQ_path': LQ_path,
'GT_path': GT_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):
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):
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.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']
# 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']
# get HR image
if self.opt['phase'] == 'train' and \
random.choice(list(range(self.ratio))) == 0: # read background images
bg_index = random.randint(0, len(self.paths_HR_bg) - 1)
HR_path = self.paths_HR_bg[bg_index]
img_HR = util.read_img(self.HR_env_bg, HR_path)
seg = torch.FloatTensor(8, img_HR.shape[0],
img_HR.shape[1]).fill_(0)
seg[0, :, :] = 1 # background
else:
HR_path = self.paths_HR[index]
img_HR = util.read_img(self.HR_env, HR_path)
seg = torch.load(
HR_path.replace('/img/', '/bicseg/').replace(
'.JPG', '_bic.pth').replace('.png', '.pth'))
# read segmentatin files, you should change it to your settings.
# modcrop in the validation / test phase
if self.opt['phase'] != 'train':
img_HR = util.modcrop(img_HR, 8)
seg = np.transpose(seg.numpy(), (1, 2, 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, _ = seg.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)
seg = cv2.resize(np.copy(seg), (W_s, H_s),
interpolation=cv2.INTER_NEAREST)
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, :]
seg = seg[rnd_h_HR:rnd_h_HR + HR_size,
rnd_w_HR:rnd_w_HR + HR_size, :]
# augmentation - flip, rotate
img_LR, img_HR, seg = util.augment([img_LR, img_HR, seg],
self.opt['use_flip'],
self.opt['use_rot'])
# category
if 'building' in HR_path:
category = 1
elif 'plant' in HR_path:
category = 2
elif 'mountain' in HR_path:
category = 3
elif 'water' in HR_path:
category = 4
elif 'sky' in HR_path:
category = 5
elif 'grass' in HR_path:
category = 6
elif 'animal' in HR_path:
category = 7
else:
category = 0 # background
else:
category = -1 # during val, useless
# 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()
seg = torch.from_numpy(
np.ascontiguousarray(np.transpose(seg, (2, 0, 1)))).float()
if LR_path is None:
LR_path = HR_path
return {
'LR': img_LR,
'HR': img_HR,
'seg': seg,
'category': category,
'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 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
}
[224, 102, 102], # 7, animal
[255, 255, 255], # 8/255, void
])).float()
lookup_table /= 255
print('seg testing...')
idx = 0
for path in glob.glob(test_img_folder + '**/*.TIF'):
idx += 1
basename = os.path.basename(path)
base = os.path.splitext(basename)[0]
print(idx, base)
# read image
img = cv2.imread(path, cv2.IMREAD_UNCHANGED)
img = modcrop(img, 8)
if img.ndim == 2:
img = np.expand_dims(img, axis=2)
img = torch.from_numpy(np.transpose(img, (2, 0, 1))).float()
# matlab imresize
# the implementation is slower than matlab, can use matlab to generate first
img_LR = imresize(img / 255, 1 / 4, antialiasing=True)
if args.thirty:
img = imresize(img_LR, 8, antialiasing=True) * 255
# from 256 to 512
else:
img = imresize(img_LR, 4, antialiasing=True) * 255
img[0] -= 103.939
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
}
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.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
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):
# 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']
# 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']
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
}
downsample_and_quantize=False,
chroma_mode=True,
block_size=8)
jpeg_compressor_16 = JPEG.JPEG(compress=True,
downsample_and_quantize=True,
downsample_only=True,
chroma_mode=True,
block_size=16)
jpeg_extractor = JPEG.JPEG(compress=False, chroma_mode=True, block_size=16)
jpeg_compressor_16.Set_Q_Table(torch.tensor(90))
jpeg_compressor_8.Set_Q_Table(torch.tensor(90))
jpeg_extractor.Set_Q_Table(torch.tensor(90))
rmse_NN, rmse_interp, rmse_NN_orig, rmse_JPEG = 0, 0, 0, 0
for im_name in tqdm(images_list):
image = cv2.imread(os.path.join(dataset_folder, im_name))
image = bgr2ycbcr(modcrop(image, 16), only_y=False).astype(float)
im_shape = list(image.shape[:2])
subsampled_chroma = np.array(image)[::2, ::2, 1:]
recovered_image_NN = np.tile(
np.expand_dims(np.expand_dims(subsampled_chroma, 2), 1),
[1, 2, 1, 2, 1]).reshape(im_shape + [-1])
recovered_image_NN = 255 * ycbcr2rgb(
np.concatenate([np.expand_dims(image[..., 0], -1), recovered_image_NN],
-1) / 255)
recovered_image_interp = cv2.resize(subsampled_chroma,
tuple(im_shape[::-1]),
interpolation=cv2.INTER_LINEAR)
recovered_image_interp = 255 * ycbcr2rgb(
np.concatenate(
[np.expand_dims(image[..., 0], -1), recovered_image_interp], -1) /
255)
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']
#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):
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图片，返回HR、LR图片对
HR_path, LR_path = None, None
HR_path = self.HR_paths[index]
img_HR = util.read_img_(HR_path, self.n_colors)
# 归一化
img_HR = util.unit2single(img_HR)
img_HR = util.modcrop(img_HR, self.scale)
if self.opt.phase == 'train':
l_max = 50
theta = np.pi * np.random.rand(1)
l1 = 0.1 + l_max * np.random.rand(1)
l2 = 0.1 + (l1 - 0.1) * np.random.rand(1)
kernel = util.anisotropic_Gaussian(ksize=self.ksize,
theta=theta[0],
l1=l1[0],
l2=l2[0])
else:
kernel = util.anisotropic_Gaussian(ksize=self.ksize,
theta=np.pi,
l1=0.1,
l2=0.1)
k = np.reshape(kernel, (-1), order="F")
k_reduced = np.dot(self.p, k)
k_reduced = torch.from_numpy(k_reduced).float()
H, W, _ = img_HR.shape
img_LR = util.srmd_degradation(img_HR, kernel, self.scale)
if self.opt.phase == 'train':
patch_H, patch_L = util.paired_random_crop(img_HR, img_LR,
self.patch_size,
self.scale)
# augmentation - flip, rotate
img_HR, img_LR = util.augment([patch_H, patch_L],
self.opt.use_flip, self.opt.use_rot)
img_HR = util.img_single2tensor(img_HR)
img_LR = util.img_single2tensor(img_LR)
if random.random() < 0.1:
noise_level = torch.zeros(1).float()
else:
noise_level = torch.FloatTensor([
np.random.uniform(self.sigma_min, self.sigma_max)
]) / 255.0
else:
img_HR = util.img_single2tensor(img_HR)
img_LR = util.img_single2tensor(img_LR)
noise_level = torch.FloatTensor([self.sigma_test])
noise = torch.randn(img_LR.size()).mul_(noise_level).float()
img_LR.add_(noise)
M_vector = torch.cat((k_reduced, noise_level),
0).unsqueeze(1).unsqueeze(1)
M = M_vector.repeat(1, img_LR.size()[-2], img_LR.size()[-1])
img_LR = torch.cat((img_LR, M), 0)
if LR_path is None:
LR_path = HR_path
return {
'LR': img_LR,
'HR': img_HR,
'LR_path': LR_path,
'HR_path': HR_path
}
