def __getitem__(self, index):
if self.data_type == 'lmdb' and self.GT_env is None:
self._init_lmdb()
key = self.paths_GT[index]
img_GT = self.read_img(key, True)
img_LQ = self.read_img(key, False)
if self.opt['phase'] == 'train':
H, W, _ = img_LQ.shape
rnd_h = random.randint(0, max(0, H - self.LQ_size))
rnd_w = random.randint(0, max(0, W - self.LQ_size))
img_LQ = img_LQ[rnd_h:rnd_h + self.LQ_size,
rnd_w:rnd_w + self.LQ_size, :]
rnd_h_HR, rnd_w_HR = int(rnd_h * self.scale), int(rnd_w *
self.scale)
img_GT = img_GT[rnd_h_HR : rnd_h_HR + self.GT_size, \
rnd_w_HR : rnd_w_HR + self.GT_size, :]
# augmentation - flip, rotate
img_LQ, img_GT = util.augment([img_LQ, img_GT],
self.opt['use_flip'],
self.opt['use_rot'])
# BGR ==> RGB
img_GT, img_LQ = img_GT[:, :, (2, 1, 0)], img_LQ[:, :, (2, 1, 0)]
# HWC ==> CHW
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()
return {'LQs': img_LQ, 'GT': img_GT, 'key': key, 'scale': self.scale}
def __getitem__(self, index):
# get HR image
# load frequence as BGR, HWDC
img_HR = self.HR_hdf5['data'][index] if self.HR_hdf5 else None
if 'data' in self.LR_hdf5:
img_LR = self.LR_hdf5['data'][index]
if self.opt['phase'] == 'train':
# augmentation - flip, rotate
img_LR, img_HR = util.augment(
[img_LR, img_HR], self.opt['use_flip'], self.opt['use_rot']
)
# BGR to RGB,
if self.HR_hdf5:
if img_HR.shape[3] == 3:
img_HR = img_HR[:, :, :, [2, 1, 0]]
# HWDC to CHWD, numpy to tensor
img_HR = torch.from_numpy(np.ascontiguousarray(np.transpose(img_HR, (3, 0, 1, 2)))).float()
# BGR to RGB,
if img_LR.shape[3] == 3:
img_LR = img_LR[:, :, :, [2, 1, 0]]
img_LR = torch.from_numpy(np.ascontiguousarray(np.transpose(img_LR, (3, 0, 1, 2)))).float()
if self.HR_hdf5:
return {'LR': img_LR, 'HR': img_HR, 'hz': self.LR_hdf5['hz'][index]}
else:
return {'LR': img_LR, 'hz': self.LR_hdf5['hz'][index]}
def __getitem__(self, index):
# get HR image
# load frequence as BGR, HWDC
img_HR = self.HR_hdf5['data'][index]
# get LR image
# load frequence as BGR, HWDC
if 'data' in self.LRx2_hdf5:
img_LRx2 = self.LRx2_hdf5['data'][index]
if 'data' in self.LRx4_hdf5:
img_LRx4 = self.LRx4_hdf5['data'][index]
if self.opt['phase'] == 'train':
# augmentation - flip, rotate
img_LRx4, img_LRx2, img_HR = util.augment(
[img_LRx4, img_LRx2, img_HR], self.opt['use_flip'], self.opt['use_rot']
)
# BGR to RGB,
if img_HR.shape[3] == 3:
img_HR = img_HR[:, :, :, [2, 1, 0]]
img_LRx2 = img_LRx2[:, :, :, [2, 1, 0]]
img_LRx4 = img_LRx4[:, :, :, [2, 1, 0]]
# HWC to CHW, numpy to tensor
img_HR = torch.from_numpy(np.ascontiguousarray(np.transpose(img_HR, (3, 0, 1, 2)))).float()
img_LRx2 = torch.from_numpy(np.ascontiguousarray(np.transpose(img_LRx2, (3, 0, 1, 2)))).float()
img_LRx4 = torch.from_numpy(np.ascontiguousarray(np.transpose(img_LRx4, (3, 0, 1, 2)))).float()
return {'LRx4': img_LRx4, 'LRx2': img_LRx2, 'HR': img_HR}
def __getitem__(self, index):
GT_path, LQ_path = None, None
GT_size = self.opt['GT_size']
# get GT audio
GT_path = self.paths_GT[index]
audio_GT = util.read_audio(GT_path)
# get LQ audio
LQ_path = self.paths_LQ[index]
resolution = None
audio_LQ = util.read_audio(LQ_path)
if self.opt['phase'] == 'train':
LQ_size = GT_size
# randomly trim
start = random.randint(0, audio_GT.shape[1] - GT_size)
stop = start + GT_size
audio_GT = audio_GT[:,start:stop,:]
audio_LQ = audio_LQ[:,start:stop,:]
# augmentation - reverse
audio_LQ, audio_GT = util.augment([audio_LQ, audio_GT], self.opt['use_reverse'])
audio_GT = torch.from_numpy(np.transpose(audio_GT)).float()
audio_LQ = torch.from_numpy(np.transpose(audio_LQ)).float()
return {'LQ': audio_LQ, 'GT': audio_GT, 'LQ_path': LQ_path, 'GT_path': GT_path}
def __getitem__(self, index):
GT_path, NOISY_path = None, None
GT_size = self.opt['GT_size']
GT_path = self.paths_GT[index]
specular_ref = util_disney.loadDisneyEXR_multi_ref_shading(
GT_path, self.opt["feature"] + ["specular"])
NOISY_path = self.paths_NOISY[index]
specular_in, features = util_disney.loadDisneyEXR_feature_shading(
NOISY_path, self.opt["feature"] + ["specular"])
if self.opt['phase'] == 'train':
# augmentation - flip, rotate
specular_ref, specular_in, features= util.augment([ specular_ref, specular_in, features], self.opt['use_flip'], \
self.opt['use_rot'])
features = torch.from_numpy(
np.ascontiguousarray(np.transpose(features, (2, 0, 1)))).float()
specular_in = torch.from_numpy(
np.ascontiguousarray(np.transpose(specular_in,
(2, 0, 1)))).float()
specular_ref = torch.from_numpy(
np.ascontiguousarray(np.transpose(specular_ref,
(2, 0, 1)))).float()
return {
"seg": features,
"specular_in": specular_in,
"specular_ref": specular_ref,
'NOISY_path': NOISY_path,
'GT_path': GT_path,
"x_offset": 128,
"y_offset": 128
}
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_channels) # return: Numpy float32, HWC, BGR, [0,1]
scale = self.scale
patch_size = self.patch_size
if self.opt.phase == 'train':
H_s, W_s, _ = img_HR.shape
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)
# using matlab imresize
img_LR = util.imresize_np(img_HR, 1 / scale, True)
H, W, C = img_LR.shape
if img_LR.ndim == 2:
img_LR = np.expand_dims(img_LR, axis=2)
if self.opt.phase == 'train':
patch_H, patch_L = util.paired_random_crop(img_HR, img_LR, patch_size, scale, gt_path)
# augmentation - flip, rotate
img_HR , img_LR = util.augment([patch_H, patch_L], self.opt.use_flip,self.opt.use_rot)
img_HR = util.img2tensor(img_HR)
img_LR = util.img2tensor(img_LR)
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.HQ_envs is None:
self._init_lmdb()
HQ_size = self.opt["HQ_size"]
env_idx, key = self.paths_HQ[index]
name_a, name_b = key.split("_")
target_frame_idx = int(name_b)
# determine the neighbor frames
# ensure not exceeding the borders
neighbor_list = [target_frame_idx]
name_b = "{:08d}".format(neighbor_list[0])
# get the HQ image (as the center frame)
img_HQ_l = []
for v in neighbor_list:
img_HQ = util.read_img(self.HQ_envs[env_idx],
"{}_{:08d}".format(name_a,
v), (3, 720, 1280))
img_HQ_l.append(img_HQ)
# get LQ images
img_LQ = util.read_img(self.LQ_envs[env_idx],
"{}_{:08d}".format(name_a, neighbor_list[-1]),
(3, 720, 1280))
if self.opt["phase"] == "train":
_, H, W = 3, 720, 1280 # LQ size
# randomly crop
rnd_h = random.randint(0, max(0, H - HQ_size))
rnd_w = random.randint(0, max(0, W - HQ_size))
img_LQ = img_LQ[rnd_h:rnd_h + HQ_size, rnd_w:rnd_w + HQ_size, :]
img_HQ_l = [
v[rnd_h:rnd_h + HQ_size, rnd_w:rnd_w + HQ_size, :]
for v in img_HQ_l
]
# augmentation - flip, rotate
img_HQ_l.append(img_LQ)
rlt = util.augment(img_HQ_l, self.opt["use_flip"],
self.opt["use_rot"])
img_HQ_l = rlt[0:-1]
img_LQ = rlt[-1]
# stack LQ images to NHWC, N is the frame number
img_HQs = np.stack(img_HQ_l, axis=0)
# BGR to RGB, HWC to CHW, numpy to tensor
img_LQ = img_LQ[:, :, [2, 1, 0]]
img_HQs = img_HQs[:, :, :, [2, 1, 0]]
img_LQ = torch.from_numpy(
np.ascontiguousarray(np.transpose(img_LQ, (2, 0, 1)))).float()
img_HQs = torch.from_numpy(
np.ascontiguousarray(np.transpose(img_HQs, (0, 3, 1, 2)))).float()
# print(img_LQ.shape, img_HQs.shape)
if self.use_identical and np.random.randint(0, 10) == 0:
img_LQ = img_HQs[-1, :, :, :]
return {"LQ": img_LQ, "HQs": img_HQs, "identical_w": 10}
return {"LQ": img_LQ, "HQs": img_HQs, "identical_w": 0}
def __getitem__(self, index):
if self.data_type == 'lmdb' and self.GT_env is None:
self._init_lmdb()
image_index, scale_index = index
scale = self.scales[scale_index]
LQ_size = self.LQ_sizes[scale_index]
key = self.paths_GT[image_index]
name_a, name_b = key.split('_')
center_frame_idx = int(name_b)
neighbor_list, name_b = self.get_neighbor_list(center_frame_idx)
key = name_a + '_' + name_b
imgs = self.read_imgs(self.GT_root, name_a, neighbor_list)
if self.opt['phase'] == 'train':
_, H, W, _ = imgs.shape
rnd_h = random.randint(0, max(0, H - self.GT_size))
rnd_w = random.randint(0, max(0, W - self.GT_size))
imgs = imgs[:, rnd_h:rnd_h + self.GT_size,
rnd_w:rnd_w + self.GT_size, :]
img_GT = imgs[self.half_N_frames] / 255.
is_PreGT = True if random.random() < self.pre_GT else False
img_PreGT = imgs[self.half_N_frames - 1] / 255.
img_LQs = [
np.array(
Image.fromarray(img).resize(
(LQ_size, LQ_size), Image.BICUBIC)) / 255.
for img in imgs
]
# augmentation - flip, rotate
img_LQs.extend([img_GT, img_PreGT])
rlt = util.augment(img_LQs, self.opt['use_flip'],
self.opt['use_rot'])
img_LQs = rlt[0:-2]
img_GT = rlt[-2]
img_PreGT = rlt[-1]
# stack LQ images to NHWC, N is the frame number
img_LQs = np.stack(img_LQs, axis=0)
# BGR => RGB
img_GT = img_GT[:, :, [2, 1, 0]]
img_PreGT = img_PreGT[:, :, [2, 1, 0]]
img_LQs = img_LQs[:, :, :, [2, 1, 0]]
# NHWC => NCHW
img_GT = torch.from_numpy(
np.ascontiguousarray(np.transpose(img_GT, (2, 0, 1)))).float()
img_PreGT = torch.from_numpy(
np.ascontiguousarray(np.transpose(img_PreGT, (2, 0, 1)))).float()
img_LQs = torch.from_numpy(
np.ascontiguousarray(np.transpose(img_LQs, (0, 3, 1, 2)))).float()
return {
'LQs': img_LQs,
'GT': img_GT,
'Pre': img_PreGT,
'if_pre': is_PreGT,
'key': key,
'scale': scale
}
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):
# 从文件中读出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):
HR_path, LR_path = None, None
scale = self.opt['scale']
# get HR image
HR_path = self.paths_HR[index]
img_HR = util.read_img(self.HR_env, HR_path)
# get LR image
if self.paths_LR:
LR_path = self.paths_LR[index]
D1_path = self.paths_D1[index]
D2_path = self.paths_D2[index]
D3_path = self.paths_D3[index]
img_LR = util.read_img(self.LR_env, LR_path)
img_D1 = util.read_img(self.D1_env, D1_path)
img_D2 = util.read_img(self.D2_env, D2_path)
img_D3 = util.read_img(self.D3_env, D3_path)
else: # down-sampling on-the-fly
if self.opt['phase'] == 'train':
# force to 3 channels
if img_HR.ndim == 2:
img_HR = cv2.cvtColor(img_HR, cv2.COLOR_GRAY2BGR)
H, W, _ = img_HR.shape
# using matlab imresize
img_LR = cv2.resize(img_HR, dsize=(int(H / scale), int(W / scale)), interpolation=cv2.INTER_CUBIC)
img_D1 = cv2.resize(img_HR, dsize=(int(H / scale * 2), int(W / scale * 2)), interpolation=cv2.INTER_CUBIC)
img_D2 = cv2.resize(img_HR, dsize=(int(H / scale * 4), int(W / scale * 4)), interpolation=cv2.INTER_CUBIC)
img_D3 = cv2.resize(img_HR, dsize=(int(H / scale * 8), int(W / scale * 8)), interpolation=cv2.INTER_CUBIC)
if img_LR.ndim == 2:
img_LR = np.expand_dims(img_LR, axis=2)
img_D1 = np.expand_dims(img_D1, axis=2)
img_D2 = np.expand_dims(img_D2, axis=2)
img_D3 = np.expand_dims(img_D3, axis=2)
if self.opt['phase'] == 'train':
# augmentation - flip, rotate
img_LR, img_HR, img_D1, img_D2, img_D3 = util.augment([img_LR, img_HR, img_D1, img_D2, img_D3],
self.opt['use_flip'], self.opt['use_rot'])
# BGR to RGB, HWC to CHW, numpy to tensor
if img_HR.shape[2] == 3:
img_HR = img_HR[:, :, [2, 1, 0]]
img_LR = img_LR[:, :, [2, 1, 0]]
img_D1 = img_D1[:, :, [2, 1, 0]]
img_D2 = img_D2[:, :, [2, 1, 0]]
img_D3 = img_D3[:, :, [2, 1, 0]]
img_HR = torch.from_numpy(np.ascontiguousarray(np.transpose(img_HR, (2, 0, 1)))).float()
img_LR = torch.from_numpy(np.ascontiguousarray(np.transpose(img_LR, (2, 0, 1)))).float()
img_D1 = torch.from_numpy(np.ascontiguousarray(np.transpose(img_D1, (2, 0, 1)))).float()
img_D2 = torch.from_numpy(np.ascontiguousarray(np.transpose(img_D2, (2, 0, 1)))).float()
img_D3 = torch.from_numpy(np.ascontiguousarray(np.transpose(img_D3, (2, 0, 1)))).float()
if LR_path is None:
LR_path = HR_path
return {'LR': img_LR, 'HR': img_HR, 'LR_path': LR_path, 'HR_path': HR_path,
'D1': img_D1, 'D2': img_D2, 'D3': img_D3, 'D1_path': D1_path, 'D2_path': D2_path, 'D3_path': D3_path}
def __getitem__(self, index):
HR_path, LR_path, MR_path = None, None, None
scale = self.opt['scale']
HR_size = self.opt['HR_size']
# get HR image
HR_path = self.paths_HR[index]
img_HR = util.read_img(self.HR_env, HR_path)
# modcrop in the validation / test phase
# if self.opt['phase'] != 'train':
# img_HR = util.modcrop(img_HR, scale)
# change color space if necessary
if self.opt['color']:
img_HR = util.channel_convert(img_HR.shape[2], self.opt['color'], [img_HR])[0]
LR_path = self.paths_LR[index]
img_LR = util.read_img(self.LR_env, LR_path)
MR_path = self.paths_MR[index]
img_MR = util.read_img(self.MR_env, MR_path)
if self.opt['noise_gt']:
img_MR = img_LR - img_MR
if self.opt['phase'] == 'train':
# if the image size is too small
H, W, C = img_LR.shape
LR_size = HR_size // scale
# randomly crop
rnd_h = random.randint(0, max(0, H - LR_size))
rnd_w = random.randint(0, max(0, W - LR_size))
img_MR = img_MR[rnd_h:rnd_h + LR_size, rnd_w:rnd_w + LR_size, :]
img_LR = img_LR[rnd_h:rnd_h + LR_size, rnd_w:rnd_w + LR_size, :]
rnd_h_HR, rnd_w_HR = int(rnd_h * scale), int(rnd_w * scale)
img_HR = img_HR[rnd_h_HR:rnd_h_HR + HR_size, rnd_w_HR:rnd_w_HR + HR_size, :]
# augmentation - flip, rotate
img_MR, img_LR, img_HR = util.augment([img_MR, img_LR, img_HR], self.opt['use_flip'], \
self.opt['use_rot'])
# channel conversion
if self.opt['color']:
# img_HR, img_LR, img_MR = util.channel_convert(C, self.opt['color'], [img_HR, img_LR, img_MR])
img_LR = util.channel_convert(C, self.opt['color'], [img_LR])[0]
img_MR = util.channel_convert(C, self.opt['color'], [img_MR])[0]
# BGR to RGB, HWC to CHW, numpy to tensor
if img_HR.shape[2] == 3:
img_HR = img_HR[:, :, [2, 1, 0]]
img_LR = img_LR[:, :, [2, 1, 0]]
img_MR = img_MR[:, :, [2, 1, 0]]
img_HR = torch.from_numpy(np.ascontiguousarray(np.transpose(img_HR, (2, 0, 1)))).float()
img_LR = torch.from_numpy(np.ascontiguousarray(np.transpose(img_LR, (2, 0, 1)))).float()
img_MR = torch.from_numpy(np.ascontiguousarray(np.transpose(img_MR, (2, 0, 1)))).float()
return {'HR': img_HR, 'LR': img_LR, 'MR': img_MR, 'HR_path': HR_path, 'MR_path': MR_path, 'LR_path': LR_path}
def __getitem__(self, index):
if self.data_type == 'lmdb':
if (self.GT_env is None) or (self.LQ_env is None):
self._init_lmdb()
GT_path, LQ_path = None, None
scale = self.opt['scale']
GT_size = self.opt['GT_size']
# get GT image
GT_path = self.paths_GT[index]
if self.data_type == 'lmdb':
resolution = [int(s) for s in self.sizes_GT[index].split('_')]
else:
resolution = None
img_GT = util.read_img(self.GT_env, GT_path, resolution)
# modcrop in the validation / test phase
if self.opt['phase'] != 'train':
img_GT = util.modcrop(img_GT, scale)
# print('val img_GT shape: ', img_GT.shape)
# change color space if necessary
if self.opt['color']:
img_GT = util.channel_convert(img_GT.shape[2], self.opt['color'],
[img_GT])[0]
# print('img_GT shape: ', img_GT.shape)
if img_GT.ndim == 2:
img_GT = np.expand_dims(img_GT, axis=2)
if self.opt['phase'] == 'train':
# if the image size is too small
H, W, C = img_GT.shape
if H < GT_size or W < GT_size:
img_GT = cv2.resize(np.copy(img_GT), (GT_size, GT_size),
interpolation=cv2.INTER_LINEAR)
# randomly crop - GT img shape: [H, W, 1]
rnd_h = random.randint(0, max(0, H - GT_size))
rnd_w = random.randint(0, max(0, W - GT_size))
img_GT = img_GT[rnd_h:rnd_h + GT_size, rnd_w:rnd_w + GT_size, :]
# augmentation - flip, rotate
img_GT = util.augment([img_GT], self.opt['use_flip'],
self.opt['use_rot'])[0]
# transform() - subsample k space - img_LF
# input img shape [H, W, 1] or [H, W, 3]
img_LF, img_GT = self.transform(img_GT, self.mask_func, self.seed)
if LQ_path is None:
LQ_path = GT_path
return {
'LQ': img_LF,
'GT': img_GT,
'LQ_path': LQ_path,
'GT_path': GT_path
}
def __getitem__(self, index):
if self.data_type == 'lmdb' and self.GT_env is None:
self._init_lmdb()
LQ_size = self.opt['LQ_size']
key = self.paths_GT[index]
name_a, name_b = key.split('_')
center_frame_idx = int(name_b)
neighbor_list, name_b = self.get_neighbor_list(center_frame_idx)
key = name_a + '_' + name_b
extra = self.valid_nf // 2
name_bs = neighbor_list[extra:self.nf - extra]
img_GTs = self.read_imgs(self.GT_root, name_a, name_bs, True)
img_LQs = self.read_imgs(self.LQ_root, name_a, neighbor_list)
if self.opt['phase'] == 'train':
H, W, _ = img_LQs[0].shape
rnd_h = random.randint(0, max(0, H - LQ_size))
rnd_w = random.randint(0, max(0, W - LQ_size))
img_LQs = [
v[rnd_h:rnd_h + LQ_size, rnd_w:rnd_w + LQ_size, :]
for v in img_LQs
]
rnd_h_HR, rnd_w_HR = int(rnd_h * self.scale), int(rnd_w *
self.scale)
GT_size = int(LQ_size * self.scale)
img_GTs = [
v[rnd_h_HR:rnd_h_HR + GT_size, rnd_w_HR:rnd_w_HR + GT_size, :]
for v in img_GTs
]
# augmentation - flip, rotate
num_gt = len(img_GTs)
img_LQs += img_GTs
img_LQs = util.augment(img_LQs, self.opt['use_flip'],
self.opt['use_rot'])
img_GTs = img_LQs[-num_gt:]
img_LQs = img_LQs[:-num_gt]
# stack LQ images to NHWC, N is the frame number
img_LQs = np.stack(img_LQs, axis=0)
img_GTs = np.stack(img_GTs, axis=0)
# BGR => RGB
img_GTs = img_GTs[:, :, :, [2, 1, 0]]
img_LQs = img_LQs[:, :, :, [2, 1, 0]]
# NHWC => NCHW
img_GTs = torch.from_numpy(
np.ascontiguousarray(np.transpose(img_GTs, (0, 3, 1, 2)))).float()
img_LQs = torch.from_numpy(
np.ascontiguousarray(np.transpose(img_LQs, (0, 3, 1, 2)))).float()
return {'LQs': img_LQs, 'GT': img_GTs, 'key': key, 'scale': self.scale}
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_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):
GT_path, NOISY_path = None, None
GT_size = self.opt['GT_size']
# get GT image
GT_path = self.paths_GT[index]
# img_GT = util.read_img(self.GT_env, GT_path)
# for HDRs
img_GT = util.load_reference_mat("", GT_path)
# get NOISY image
if self.paths_NOISY:
NOISY_path = self.paths_NOISY[index]
all_feature = util.load_feature_mat(
self.FEATURE_DIR,
NOISY_path) #util.load_feature_mat_complete_tungsten
img_NOISY = all_feature[:, :, :3]
features = all_feature[:, :, 3:]
# do the cropping
H, W, _ = img_GT.shape
# x = random.randint(0, np.maximum(0, W - GT_size))
# y = random.randint(0, np.maximum(0, H - GT_size))
x = 0
y = 0
img_GT = util._crop(img_GT, (y, x), GT_size)
img_NOISY = util._crop(img_NOISY, (y, x), GT_size)
features = util._crop(features, (y, x), GT_size)
if self.opt['phase'] == 'train':
# augmentation - flip, rotate
img_NOISY, img_GT, features = util.augment([img_NOISY, img_GT, features], self.opt['use_flip'], \
self.opt['use_rot'])
img_GT = torch.from_numpy(
np.ascontiguousarray(np.transpose(img_GT, (2, 0, 1)))).float()
img_NOISY = torch.from_numpy(
np.ascontiguousarray(np.transpose(img_NOISY, (2, 0, 1)))).float()
features = torch.from_numpy(
np.ascontiguousarray(np.transpose(features, (2, 0, 1)))).float()
return {
'diffuse_in': img_NOISY,
'diffuse_ref': img_GT,
"seg": features,
"category": 1,
'NOISY_path': NOISY_path,
'GT_path': GT_path,
"x_offset": 512,
"y_offset": 512
}
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):
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):
GT_path, NOISY_path = None, None
GT_size = self.opt['GT_size']
# get GT image
GT_path = self.paths_GT[index]
NOISY_path = self.paths_NOISY[index]
# print(GT_path)
# print(NOISY_path)
img_GT = util.read_img_GIGAN_version(GT_path)
# for HDRs
# img_GT = util.load_reference_mat_shading("", GT_path, NOISY_path)#图片那边还需要改！！！！
# get NOISY image
all_feature = util.load_feature_GIGAN_version(
NOISY_path,
self.DIFFUSE_DIR) #util.load_feature_mat_complete_tungsten
img_NOISY = all_feature[:, :, :3]
features = all_feature[:, :, 3:]
# do the cropping
# H, W, _ = img_GT.shape
# x = random.randint(0, np.maximum(0, W - GT_size))
# y = random.randint(0, np.maximum(0, H - GT_size))
# img_GT = util._crop(img_GT, (y,x), GT_size)
# img_NOISY = util._crop(img_NOISY, (y,x), GT_size)
# features = util._crop(features, (y,x), GT_size)
if self.opt['phase'] == 'train':
# augmentation - flip, rotate
img_NOISY, img_GT, features = util.augment([img_NOISY, img_GT, features], self.opt['use_flip'], \
self.opt['use_rot'])
img_GT = torch.from_numpy(
np.ascontiguousarray(np.transpose(img_GT, (2, 0, 1)))).float()
img_NOISY = torch.from_numpy(
np.ascontiguousarray(np.transpose(img_NOISY, (2, 0, 1)))).float()
features = torch.from_numpy(
np.ascontiguousarray(np.transpose(features, (2, 0, 1)))).float()
# all_feature = torch.from_numpy(np.ascontiguousarray(np.transpose(all_feature, (2, 0, 1)))).float()
return {
'NOISY': img_NOISY,
'GT': img_GT,
"seg": features,
"category": 1,
'NOISY_path': NOISY_path,
'GT_path': GT_path
}
def __getitem__(self, index):
# 从文件中读出HR图片，返回HR、LR图片对
HR_path = self.HR_paths[index]
LR_path = self.LR_paths[index]
img_HR = util.read_img_(HR_path, self.n_colors)
img_LR = util.read_img_(LR_path, self.n_colors)
if not self.opt.test_only == 'train':
patch_H, patch_L = util.paired_random_crop(img_HR, img_LR,
self.patch_size,
self.scale)
img_HR, img_LR = util.augment([patch_H, patch_L],
self.opt.use_flip, self.opt.use_rot)
img_HR = util.img2tensor(img_HR)
img_LR = util.img2tensor(img_LR)
return {'LR': img_LR, 'HR': img_HR}
def __getitem__(self, index):
lrs, gt, video_name, filename = self.read_data(index)
if self.train:
lrs, gt = self.random_crop(lrs, gt, self.patch_size, self.scale)
lrs.append(gt)
rlt = util.augment(lrs)
# rlt = [util.rgb2ycbcr(img, only_y=True) for img in rlt]
lrs, gt = rlt[:-1], rlt[-1]
else:
lrs.append(gt)
# rlt = [util.rgb2ycbcr(img, only_y=True) for img in lrs]
lrs, gt = rlt[:-1], rlt[-1]
lrs, gt = self.to_tensor(lrs, gt)
# return lrs, gt, video_name, filename
key = '{}_{}'.format(video_name, filename)
return {'LQs': lrs, 'GT': gt, 'key': key}
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):
image_index, scale_index = index
scale = self.scales[scale_index]
LQ_size = self.opt['LQ_size']
key = self.paths_GT[image_index]
name_a, name_b = key.split('_')
center_frame_idx = int(name_b)
neighbor_list, name_b = self.get_neighbor_list(center_frame_idx)
#### get the images
img_GT = self.read_imgs(self.GT_root, name_a, name_b)[0]
img_LQs = self.read_imgs(self.LQ_root, name_a, neighbor_list, scale)
if self.opt['phase'] == 'train':
H, W, _ = img_LQs[0].shape
rnd_h = random.randint(0, max(0, H - LQ_size))
rnd_w = random.randint(0, max(0, W - LQ_size))
img_LQs = [
v[rnd_h:rnd_h + LQ_size, rnd_w:rnd_w + LQ_size, :]
for v in img_LQs
]
rnd_h_HR, rnd_w_HR = int(rnd_h * scale), int(rnd_w * scale)
GT_size = int(LQ_size * scale)
img_GT = img_GT[rnd_h_HR:rnd_h_HR + GT_size,
rnd_w_HR:rnd_w_HR + GT_size, :]
# augmentation - flip, rotate
img_LQs.append(img_GT)
rlt = util.augment(img_LQs, self.opt['use_flip'],
self.opt['use_rot'])
img_LQs = rlt[0:-1]
img_GT = rlt[-1]
# stack LQ images to NHWC, N is the frame number
img_LQs = np.stack(img_LQs, axis=0)
# BGR => RGB
img_GT = img_GT[:, :, [2, 1, 0]]
img_LQs = img_LQs[:, :, :, [2, 1, 0]]
# NHWC => NCHW
img_GT = torch.from_numpy(
np.ascontiguousarray(np.transpose(img_GT, (2, 0, 1)))).float()
img_LQs = torch.from_numpy(
np.ascontiguousarray(np.transpose(img_LQs, (0, 3, 1, 2)))).float()
return {'LQs': img_LQs, 'GT': img_GT, 'key': key, 'scale': scale}
def __getitem__(self, index):
# get LQ and HQ image
LQ_path = self.paths_LQ[index]
HQ_path = self.paths_HQ[index]
img_LQ = util.read_img(None, LQ_path, None)
img_HQ = util.read_img(None, HQ_path, None)
if self.opt['crop_size']:
# randomly crop
LH, LW, _ = img_LQ.shape
rnd_h = random.randint(0, max(0, LH - self.crop_size))
rnd_w = random.randint(0, max(0, LW - self.crop_size))
rnd_hh = rnd_h * self.opt['scale']
rnd_wh = rnd_w * self.opt['scale']
patch_size = self.crop_size * self.opt['scale']
img_LQ = img_LQ[rnd_h:rnd_h + self.crop_size, rnd_w:rnd_w + self.crop_size, :]
img_HQ = img_HQ[rnd_hh:rnd_hh + patch_size, rnd_wh:rnd_wh + patch_size, :]
if self.opt['phase'] == 'train':
# augmenttation cutclur
if self.opt['use_cutblur']:
img_LQ,img_HQ=util.argment_cutblur(img_LQ, img_HQ,self.opt['scale'])
if self.opt['use_rgbpermute']:
img_LQ, img_HQ=util.argment_rgb(img_LQ, img_HQ)
# augmentation - flip, rotate
img_LQ, img_HQ = util.augment([img_LQ, img_HQ], self.opt['use_flip'],
self.opt['use_rot'])
# BGR to RGB, HWC to CHW, numpy to tensor
if img_LQ.shape[2] == 3:
img_LQ = img_LQ[:, :, [2, 1, 0]]
if img_HQ.shape[2] == 3:
img_HQ = img_HQ[:, :, [2, 1, 0]]
img_LQ = torch.from_numpy(np.ascontiguousarray(np.transpose(img_LQ, (2, 0, 1)))).float()
img_HQ = torch.from_numpy(np.ascontiguousarray(np.transpose(img_HQ, (2, 0, 1)))).float()
return {'LQ': img_LQ, 'LQ_path': LQ_path, 'HQ': img_HQ, 'HQ_path': HQ_path}
def __getitem__(self, index):
if self.data_type == 'lmdb' and self.GT_env is None:
self._init_lmdb()
image_index, scale_index = index
scale = self.scales[scale_index]
LQ_size = self.LQ_sizes[scale_index]
key = self.paths_GT[image_index]
imgs = self.read_imgs(key)
if self.opt['phase'] == 'train':
_, H, W, _ = imgs.shape
rnd_h = random.randint(0, max(0, H - self.GT_size))
rnd_w = random.randint(0, max(0, W - self.GT_size))
imgs = imgs[:, rnd_h:rnd_h + self.GT_size,
rnd_w:rnd_w + self.GT_size, :]
img_GT = imgs[self.half_N_frames] / 255.
img_LQs = [
np.array(
Image.fromarray(img).resize(
(LQ_size, LQ_size), Image.BICUBIC)) / 255.
for img in imgs
]
img_LQs.append(img_GT)
rlt = util.augment(img_LQs, self.opt['use_flip'],
self.opt['use_rot'])
img_LQs = rlt[0:-1]
img_GT = rlt[-1]
# stack LQ images to NHWC, N is the frame number
img_LQs = np.stack(img_LQs, axis=0)
# BGR => RGB, HWC to CHW, numpy to tensor
img_GT = img_GT[:, :, [2, 1, 0]]
img_LQs = img_LQs[:, :, :, [2, 1, 0]]
# NHWC => NCHW
img_GT = torch.from_numpy(
np.ascontiguousarray(np.transpose(img_GT, (2, 0, 1)))).float()
img_LQs = torch.from_numpy(
np.ascontiguousarray(np.transpose(img_LQs, (0, 3, 1, 2)))).float()
return {'LQs': img_LQs, 'GT': img_GT, 'key': key, 'scale': scale}
def __getitem__(self, index):
if self.data_type == 'lmdb' and self.GT_env is None:
self._init_lmdb()
key = self.paths_GT[index]
name_a, name_b = key.split('_')
center_frame_idx = int(name_b)
neighbors, name_b = self.get_neighbor_list(center_frame_idx)
key = name_a + "_" + name_b
keys = [name_a + "_{:08d}".format(nei) for nei in neighbors]
img_GTs = self.read_imgs(keys, True)
img_LQs = self.read_imgs(keys, False)
if self.opt['phase'] == 'train':
H, W, _ = img_LQs[0].shape
rnd_h = random.randint(0, max(0, H - self.LQ_size))
rnd_w = random.randint(0, max(0, W - self.LQ_size))
img_LQs = [v[rnd_h : rnd_h + self.LQ_size, rnd_w : rnd_w + \
self.LQ_size, :] for v in img_LQs]
rnd_h_HR, rnd_w_HR = int(rnd_h * self.scale), int(rnd_w *
self.scale)
img_GTs = [v[rnd_h_HR : rnd_h_HR + self.GT_size, rnd_w_HR : rnd_w_HR + \
self.GT_size, :] for v in img_GTs]
# augmentation - flip, rotate
img_LQs += img_GTs
rlts = util.augment(img_LQs, self.opt['use_flip'],
self.opt['use_rot'])
img_GTs = rlts[self.nframes:]
img_LQs = rlts[:self.nframes]
# stack images to NHWC; BGR -> RGB; NHWC -> CNHW
img_GTs = np.stack(img_GTs, axis=0)[:, :, :, (2, 1, 0)].transpose(
(3, 0, 1, 2))
img_LQs = np.stack(img_LQs, axis=0)[:, :, :, (2, 1, 0)].transpose(
(3, 0, 1, 2))
img_GTs = torch.from_numpy(img_GTs).float()
img_LQs = torch.from_numpy(img_LQs).float()
return {'LQs': img_LQs, 'GT': img_GTs, 'key': key, 'scale': self.scale}
def __getitem__(self, index):
if self.data_type == "mc":
self._ensure_memcached()
elif self.data_type == "lmdb" and (self.HQ_env is None
or self.LQ_env is None):
self._init_lmdb()
HQ_size = self.opt["HQ_size"]
key = self.paths_HQ[index]
name_a, name_b = key.split("_")
# get the HQ image
img_HQ = util.read_img(self.HQ_env, key, (3, 720, 1280))
# get the LQ image
img_LQ = util.read_img(self.LQ_env, key, (3, 720, 1280))
if self.opt["phase"] == "train":
_, H, W = 3, 720, 1280 # LQ size
# randomly crop
rnd_h = random.randint(0, max(0, H - HQ_size))
rnd_w = random.randint(0, max(0, W - HQ_size))
img_LQ = img_LQ[rnd_h:rnd_h + HQ_size, rnd_w:rnd_w + HQ_size, :]
img_HQ = img_HQ[rnd_h:rnd_h + HQ_size, rnd_w:rnd_w + HQ_size, :]
# augmentation - flip, rotate
imgs = [img_HQ, img_LQ]
rlt = util.augment(imgs, self.opt["use_flip"], self.opt["use_rot"])
img_HQ = rlt[0]
img_LQ = rlt[1]
# BGR to RGB, HWC to CHW, numpy to tensor
img_LQ = img_LQ[:, :, [2, 1, 0]]
img_HQ = img_HQ[:, :, [2, 1, 0]]
img_LQ = torch.from_numpy(
np.ascontiguousarray(np.transpose(img_LQ, (2, 0, 1)))).float()
img_HQ = torch.from_numpy(
np.ascontiguousarray(np.transpose(img_HQ, (2, 0, 1)))).float()
return {"LQ": img_LQ, "HQ": img_HQ}
def __getitem__(self, index):
if self.opt['data_type'] == 'lmdb':
if self.LR_env is None:
self._init_lmdb()
LR_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):
path, target = self.samples[index]
'''img_dis = self.loader(path[0])
img_ref = self.loader(path[1])'''
img_dis = read_img(env=None, path=path[0])
img_ref = read_img(env=None, path=path[1])
'''if self.transform is not None:
img_dis = self.transform(img_dis)
img_ref = self.transform(img_ref)'''
if self.patch_size < 288:
H, W, _ = img_ref.shape
crop_size = self.patch_size
rnd_h = random.randint(0, max(0, (H - crop_size)))
rnd_w = random.randint(0, max(0, (W - crop_size)))
img_dis = img_dis[rnd_h:rnd_h + crop_size,
rnd_w:rnd_w + crop_size, :]
img_ref = img_ref[rnd_h:rnd_h + crop_size,
rnd_w:rnd_w + crop_size, :]
# augmentation - flip, rotate
img_dis, img_ref = augment([img_dis, img_ref],
self.opt['use_flip'],
rot=False)
if img_ref.shape[2] == 3:
img_ref = img_ref[:, :, [2, 1, 0]]
img_dis = img_dis[:, :, [2, 1, 0]]
img_ref = torch.from_numpy(
np.ascontiguousarray(np.transpose(img_ref, (2, 0, 1)))).float()
img_dis = torch.from_numpy(
np.ascontiguousarray(np.transpose(img_dis, (2, 0, 1)))).float()
img_dis = self.transform(img_dis)
img_ref = self.transform(img_ref)
return {'Dis': img_dis, 'Ref': img_ref, 'Label': target}