def __init__(self, opt):
super(LRHRDataset, self).__init__()
self.opt = opt
self.paths_LR = None
self.paths_HR = None
self.LR_env = None # environment for lmdb
self.HR_env = None
# read image list from subset list txt
if opt['subset_file'] is not None and opt['phase'] == 'train':
with open(opt['subset_file']) as f:
self.paths_HR = sorted([os.path.join(opt['dataroot_HR'], line.rstrip('\n')) \
for line in f])
if opt['dataroot_LR'] is not None:
raise NotImplementedError(
'Now subset only supports generating LR on-the-fly.')
else: # read image list from lmdb or image files
self.HR_env, self.paths_HR = util.get_image_paths(
opt['data_type'], opt['dataroot_HR'])
self.LR_env, self.paths_LR = util.get_image_paths(
opt['data_type'], opt['dataroot_LR'])
assert self.paths_HR, 'Error: HR path is empty.'
if self.paths_LR and self.paths_HR:
assert len(self.paths_LR) == len(self.paths_HR), \
'HR and LR datasets have different number of images - {}, {}.'.format(\
len(self.paths_LR), len(self.paths_HR))
self.random_scale_list = [1]
def __init__(self, opt):
super(LRHRSeg_BG_Dataset, self).__init__()
self.opt = opt
self.paths_LR = None
self.paths_HR = None
self.paths_HR_bg = None # HR images for background scenes
self.LR_env = None # environment for lmdb
self.HR_env = None
self.HR_env_bg = None
# read image list from lmdb or image files
self.HR_env, self.paths_HR = util.get_image_paths(
opt['data_type'], opt['dataroot_HR'])
self.LR_env, self.paths_LR = util.get_image_paths(
opt['data_type'], opt['dataroot_LR'])
self.HR_env_bg, self.paths_HR_bg = util.get_image_paths(opt['data_type'], \
opt['dataroot_HR_bg'])
assert self.paths_HR, 'Error: HR path is empty.'
if self.paths_LR and self.paths_HR:
assert len(self.paths_LR) == len(self.paths_HR), \
'HR and LR datasets have different number of images - {}, {}.'.format(\
len(self.paths_LR), len(self.paths_HR))
self.random_scale_list = [1, 0.9, 0.8, 0.7, 0.6, 0.5]
self.ratio = 10 # 10 OST data samples and 1 DIV2K general data samples(background)
def __init__(self, opt):
super(LRHRSeg_BG_Dataset, self).__init__(opt, keys_ds=['LR', 'HR'])
self.znorm = opt.get(
'znorm',
False) # Alternative: images are z-normalized to the [-1,1] range
# self.opt = opt
# self.paths_LR = None
# self.paths_HR = None
self.paths_HR_bg = None # HR images for background scenes
self.LR_env = None
self.HR_env = None
self.HR_env_bg = None # environment for lmdb
# get images paths (and optional environments for lmdb) from dataroots
self.paths_LR, self.paths_HR = get_dataroots_paths(
opt, strict=False, keys_ds=self.keys_ds)
# read backgrounds image list from lmdb or image files
self.paths_HR_bg = util.get_image_paths(opt['data_type'],
opt['dataroot_HR_bg'])
if self.opt.get('data_type') == 'lmdb':
self.LR_env = util._init_lmdb(
opt.get('dataroot_' + self.keys_ds[0]))
self.HR_env = util._init_lmdb(
opt.get('dataroot_' + self.keys_ds[1]))
self.HR_env_bg = util._init_lmdb(opt.get('dataroot_HR_bg'))
assert len(self.paths_HR) == len(self.paths_HR_bg)
self.random_scale_list = [1, 0.9, 0.8, 0.7, 0.6, 0.5]
self.ratio = 10 # 10 OST data samples and 1 DIV2K general data samples(background)
def __init__(self, opt):
super(DVDIDataset, self).__init__()
self.opt = opt
self.paths_in = None
# read image list from lmdb or image files
self.paths_in = get_image_paths('img', opt['dataroot_in'])
assert self.paths_in, 'Error: Interlaced paths are empty.'
def __init__(self, opt):
super(DVDDataset, self).__init__()
self.opt = opt
self.debug = None # path to a directory to save debug files, None = Disable
self.paths_in = get_image_paths('img', opt['dataroot_in'])
self.paths_top = get_image_paths('img', opt['dataroot_top'])
self.paths_bot = get_image_paths('img', opt['dataroot_bottom'])
self.paths_progressive = get_image_paths('img', opt['dataroot_progressive'])
if self.paths_in and self.paths_top and self.paths_bot:
assert len(self.paths_top) >= len(self.paths_in), \
'Top dataset contains fewer images than interlaced dataset - {}, {}.'.format(
len(self.paths_top), len(self.paths_in))
assert len(self.paths_bot) >= len(self.paths_in), \
'Bottom dataset contains fewer images than interlaced dataset - {}, {}.'.format(
len(self.paths_bot), len(self.paths_in))
def __init__(self, opt):
super(LRDataset, self).__init__()
self.opt = opt
self.paths_LR = None
self.LR_env = None # environment for lmdb
self.znorm = opt.get(
'znorm',
False) # Alternative: images are z-normalized to the [-1,1] range
# read image list from lmdb or image files
self.LR_env, self.paths_LR = util.get_image_paths(
opt['data_type'], opt['dataroot_LR'])
assert self.paths_LR, 'Error: LR paths are empty.'
def paired_dataset_validation(A_images_paths,
B_images_paths,
data_type='img',
max_dataset_size=float("inf")):
if isinstance(A_images_paths, str) and isinstance(B_images_paths, str):
A_images_paths = [A_images_paths]
B_images_paths = [B_images_paths]
paths_A = []
paths_B = []
for paths in zip(A_images_paths, B_images_paths):
A_paths = get_image_paths(data_type, paths[0],
max_dataset_size) # get image paths
B_paths = get_image_paths(data_type, paths[1],
max_dataset_size) # get image paths
for imgs in zip(A_paths, B_paths):
_, A_filename = os.path.split(imgs[0])
_, B_filename = os.path.split(imgs[1])
assert A_filename == B_filename, f'Wrong pair of images {A_filename} and {B_filename}'
paths_A.append(imgs[0])
paths_B.append(imgs[1])
return paths_A, paths_B
def process_img_paths(images_paths=None,
data_type='img',
max_dataset_size=float("inf")):
if not images_paths:
return images_paths
# process images_paths
paths_list = []
for path in images_paths:
paths = get_image_paths(data_type, path, max_dataset_size)
for imgs in paths:
paths_list.append(imgs)
paths_list = sorted(paths_list)
return paths_list
def __init__(self, opt):
super(LRHRDataset, self).__init__()
self.opt = opt
self.paths_LR, self.paths_HR = None, None
self.LR_env, self.HR_env = None, None # environment for lmdb
self.output_sample_imgs = None
if opt.get('dataroot_kernels', None):
#TODO: note: use the model scale to get the right kernel
scale = opt.get('scale', 4)
self.ds_kernels = KernelDownscale(
scale=scale, kernel_paths=opt['dataroot_kernels'])
if opt['phase'] == 'train' and opt.get(
'lr_noise_types', 3) and "patches" in opt['lr_noise_types']:
assert opt['noise_data']
self.noise_patches = NoisePatches(
opt['noise_data'],
opt.get('HR_size', 128) / opt.get('scale', 4))
else:
self.noise_patches = None
# read image list from subset list txt
if opt['subset_file'] is not None and opt['phase'] == 'train':
with open(opt['subset_file']) as f:
self.paths_HR = sorted([os.path.join(opt['dataroot_HR'], line.rstrip('\n')) \
for line in f])
if opt['dataroot_LR'] is not None:
raise NotImplementedError(
'Subset only supports generating LR on-the-fly.')
else: # read image list from lmdb or image files
# Check if dataroot_HR is a list of directories or a single directory. Note: lmdb will not currently work with a list
HR_images_paths = opt['dataroot_HR']
# if receiving a single path in str format, convert to list
if type(HR_images_paths) is str:
HR_images_paths = [HR_images_paths]
# Check if dataroot_LR is a list of directories or a single directory. Note: lmdb will not currently work with a list
LR_images_paths = opt['dataroot_LR']
if not LR_images_paths:
LR_images_paths = []
# if receiving a single path in str format, convert to list
if type(LR_images_paths) is str:
LR_images_paths = [LR_images_paths]
self.paths_HR = []
self.paths_LR = []
# special case when dealing with duplicate HR_images_paths or LR_images_paths
# LMDB will not be supported with this option
if len(HR_images_paths) != len(set(HR_images_paths)) or \
len(LR_images_paths) != len(set(LR_images_paths)):
# only resolve when the two path lists coincide in the number of elements,
# they have to be ordered specifically as they will be used in the options file
assert len(HR_images_paths) == len(LR_images_paths), \
'Error: When using duplicate paths, dataroot_HR and dataroot_LR must contain the same number of elements.'
for paths in zip(HR_images_paths, LR_images_paths):
_, paths_HR = util.get_image_paths(opt['data_type'],
paths[0])
_, paths_LR = util.get_image_paths(opt['data_type'],
paths[1])
for imgs in zip(paths_HR, paths_LR):
_, HR_filename = os.path.split(imgs[0])
_, LR_filename = os.path.split(imgs[1])
assert HR_filename == LR_filename, 'Wrong pair of images {} and {}'.format(
HR_filename, LR_filename)
self.paths_HR.append(imgs[0])
self.paths_LR.append(imgs[1])
else: # for cases with extra HR directories for OTF images or original single directories
self.HR_env = []
self.LR_env = []
# process HR_images_paths
# if type(HR_images_paths) is list:
for path in HR_images_paths:
HR_env, paths_HR = util.get_image_paths(
opt['data_type'], path)
if type(HR_env) is list:
for imgs in HR_env:
self.HR_env.append(imgs)
for imgs in paths_HR:
self.paths_HR.append(imgs)
if self.HR_env.count(None) == len(self.HR_env):
self.HR_env = None
else:
self.HR_env = sorted(self.HR_env)
self.paths_HR = sorted(self.paths_HR)
# elif type(HR_images_paths) is str:
# self.HR_env, self.paths_HR = util.get_image_paths(opt['data_type'], HR_images_paths)
# process LR_images_paths
# if type(LR_images_paths) is list:
for path in LR_images_paths:
LR_env, paths_LR = util.get_image_paths(
opt['data_type'], path)
if type(LR_env) is list:
for imgs in LR_env:
self.LR_env.append(imgs)
for imgs in paths_LR:
self.paths_LR.append(imgs)
if self.LR_env.count(None) == len(self.LR_env):
self.LR_env = None
else:
self.LR_env = sorted(self.LR_env)
self.paths_LR = sorted(self.paths_LR)
# elif type(LR_images_paths) is str:
# self.LR_env, self.paths_LR = util.get_image_paths(opt['data_type'], LR_images_paths)
assert self.paths_HR, 'Error: HR path is empty.'
if self.paths_LR and self.paths_HR:
# Modified to allow using HR and LR folders with different amount of images
# - If an LR image pair is not found, downscale HR on the fly, else, use the LR
# - If all LR are provided and 'lr_downscale' is enabled, randomize use of provided LR and OTF LR for augmentation
"""
assert len(self.paths_LR) == len(self.paths_HR), \
'HR and LR datasets have different number of images - {}, {}.'.format(\
len(self.paths_LR), len(self.paths_HR))
"""
#"""
assert len(self.paths_HR) >= len(self.paths_LR), \
'HR dataset contains less images than LR dataset - {}, {}.'.format(\
len(self.paths_HR), len(self.paths_LR))
#"""
if len(self.paths_LR) < len(self.paths_HR):
print(
'LR contains less images than HR dataset - {}, {}. Will generate missing images on the fly.'
.format(len(self.paths_LR), len(self.paths_HR)))
i = 0
tmp = []
for idx in range(0, len(self.paths_HR)):
_, HRtail = os.path.split(self.paths_HR[idx])
if i < len(self.paths_LR):
LRhead, LRtail = os.path.split(self.paths_LR[i])
if LRtail == HRtail:
LRimg_path = os.path.join(LRhead, LRtail)
tmp.append(LRimg_path)
i += 1
else:
LRimg_path = None
tmp.append(LRimg_path)
else: #if the last image is missing
LRimg_path = None
tmp.append(LRimg_path)
self.paths_LR = tmp
def __getitem__(self, idx):
scale = self.opt.get('scale', 4)
HR_size = self.opt.get('HR_size', 128)
LR_size = HR_size // scale
idx_center = (self.num_frames - 1) // 2
ds_kernel = None
# Default case: tensor will result in the [0,1] range
# Alternative: tensor will be z-normalized to the [-1,1] range
znorm = self.opt.get('znorm', False)
if self.opt['phase'] == 'train':
if self.opt.get('lr_downscale', None) and self.opt.get('dataroot_kernels', None) and 999 in self.opt["lr_downscale_types"]:
ds_kernel = self.ds_kernels #KernelDownscale(scale, self.kernel_paths, self.num_kernel)
# get a random video directory
idx_video = random.randint(0, len(self.video_list)-1)
video_dir = self.video_list[idx_video]
# print(video_dir)
else:
# only one video and paths_LR/paths_HR is already the video dir
video_dir = ""
# list the frames in the directory
# hr_dir = self.trainset_dir + '/' + video_dir + '/hr'
paths_HR = util.get_image_paths(self.opt['data_type'], os.path.join(self.paths_HR, video_dir))
# print(paths_HR)
if self.opt['phase'] == 'train':
# random reverse augmentation
random_reverse = self.opt.get('random_reverse', False)
# skipping intermediate frames to learn from low FPS videos augmentation
# testing random frameskip up to 'max_frameskip' frames
max_frameskip = self.opt.get('max_frameskip', 0)
if max_frameskip > 0:
max_frameskip = min(max_frameskip, len(paths_HR)//(self.num_frames-1))
frameskip = random.randint(1, max_frameskip)
else:
frameskip = 1
# print("max_frameskip: ", max_frameskip)
assert ((self.num_frames-1)*frameskip) <= (len(paths_HR)-1), (
f'num_frame*frameskip must be smaller than the number of frames per video, check {video_dir}')
# if number of frames of training video is for example 31, "max index -num_frames" = 31-3=28
idx_frame = random.randint(0, (len(paths_HR)-1)-((self.num_frames-1)*frameskip))
# print('frameskip:', frameskip)
else:
frameskip = 1
idx_frame = idx
'''
List based frames loading
'''
if self.paths_LR:
paths_LR = util.get_image_paths(self.opt['data_type'], os.path.join(self.paths_LR, video_dir))
else:
paths_LR = paths_HR
ds_algo = 777 # default to matlab-like bicubic downscale
if self.opt.get('lr_downscale', None): # if manually set and scale algorithms are provided, then:
ds_algo = self.opt.get('lr_downscale_types', 777)
# get the video directory
HR_dir, _ = os.path.split(paths_HR[idx_frame])
LR_dir, _ = os.path.split(paths_HR[idx_frame])
# read HR & LR frames
HR_list = []
LR_list = []
resize_type = None
LR_bicubic = None
HR_center = None
# print('len(paths_HR)', len(paths_HR))
for i_frame in range(self.num_frames):
# print('frame path:', paths_HR[int(idx_frame)+(frameskip*i_frame)])
HR_img = util.read_img(None, paths_HR[int(idx_frame)+(frameskip*i_frame)], out_nc=self.image_channels)
HR_img = util.modcrop(HR_img, scale)
if self.opt['phase'] == 'train':
'''
If using individual image augmentations, get cropping parameters for reuse
'''
if self.otf_noise and i_frame == 0: #only need to calculate once, from the first frame
# reuse the cropping parameters for all LR and HR frames
hr_crop_params, lr_crop_params = get_crop_params(HR_img, LR_size, scale)
if self.opt.get('lr_noise', None):
# reuse the same noise type for all the frames
noise_option = get_noise(self.opt.get('lr_noise_types', None), self.noise_patches)
if self.opt.get('lr_blur', None):
# reuse the same blur type for all the frames
blur_option = get_blur(self.opt.get('lr_blur_types', None))
if self.paths_LR:
# LR images are provided at the correct scale
LR_img = util.read_img(None, paths_LR[int(idx_frame)+(frameskip*i_frame)], out_nc=self.image_channels)
if LR_img.shape == HR_img.shape:
LR_img, resize_type = Scale(img=HR_img, scale=scale, algo=ds_algo, ds_kernel=ds_kernel, resize_type=resize_type)
else:
# generate LR images on the fly
LR_img, resize_type = Scale(img=HR_img, scale=scale, algo=ds_algo, ds_kernel=ds_kernel, resize_type=resize_type)
# get the bicubic upscale of the center frame to concatenate for SR
if self.y_only and self.srcolors and i_frame == idx_center:
LR_bicubic, _ = Scale(img=LR_img, scale=1/scale, algo=777) # bicubic upscale
HR_center = HR_img
# tmp_vis(LR_bicubic, False)
# tmp_vis(HR_center, False)
if self.y_only:
# extract Y channel from frames
# normal path, only Y for both
HR_img = util.bgr2ycbcr(HR_img, only_y=True)
LR_img = util.bgr2ycbcr(LR_img, only_y=True)
# crop patches randomly if using otf noise
#TODO: make a BasicSR composable random_crop
#TODO: note the original crop should go here and crop after loading each image, but could also be much simpler
# to crop after concatenating. Check the speed difference.
if self.otf_noise and self.opt['phase'] == 'train':
HR_img, LR_img = apply_crop_params(HR_img, LR_img, hr_crop_params, lr_crop_params)
if self.y_only and self.srcolors and i_frame == idx_center:
LR_bicubic, _ = apply_crop_params(LR_bicubic, None, hr_crop_params, None)
HR_center, _ = apply_crop_params(HR_center, None, hr_crop_params, None)
# expand Y images to add the channel dimension
# normal path, only Y for both
if self.y_only:
HR_img = util.fix_img_channels(HR_img, 1)
LR_img = util.fix_img_channels(LR_img, 1)
if self.opt['phase'] == 'train':
# single frame augmentation (noise, blur, etc). Would only be efficient if patches are cropped in this loop
if self.opt.get('lr_blur', None):
if blur_option:
LR_img = blur_option(LR_img)
if self.opt.get('lr_noise', None):
if noise_option:
LR_img = noise_option(LR_img)
# expand LR images to add the channel dimension again if needed (blur removes the grayscale channel)
#TODO: add a if condition, can compare to the ndim before the augs, maybe move inside the aug condition
# if not fullimgchannels: #TODO: TMP, this should be when using srcolors for HR or when training with 3 channels tests, separatedly
if self.y_only:
LR_img = util.fix_img_channels(LR_img, 1)
# print("HR_img.shape: ", HR_img.shape)
# print("LR_img.shape", LR_img.shape)
HR_list.append(HR_img) # h, w, c
LR_list.append(LR_img) # h, w, c
# print(len(HR_list))
# print(len(LR_list))
if self.opt['phase'] == 'train':
# random reverse sequence augmentation
if random_reverse and random.random() < 0.5:
HR_list.reverse()
LR_list.reverse()
if not self.y_only:
t = self.num_frames
HR = [np.asarray(GT) for GT in HR_list] # list -> numpy # input: list (contatin numpy: [H,W,C])
HR = np.asarray(HR) # numpy, [T,H,W,C]
h_HR, w_HR, c = HR_img.shape #HR_center.shape #TODO: check, may be risky
HR = HR.transpose(1,2,3,0).reshape(h_HR, w_HR, -1) # numpy, [H',W',CT]
LR = [np.asarray(LT) for LT in LR_list] # list -> numpy # input: list (contatin numpy: [H,W,C])
LR = np.asarray(LR) # numpy, [T,H,W,C]
LR = LR.transpose(1,2,3,0).reshape(h_HR//scale, w_HR//scale, -1) # numpy, [Hl',Wl',CT]
else:
HR = np.concatenate((HR_list), axis=2) # h, w, t
LR = np.concatenate((LR_list), axis=2) # h, w, t
if self.opt['phase'] == 'train':
'''
# If not using individual image augmentations, this cropping should be faster, only once
'''
# crop patches randomly. If not using otf noise, crop all concatenated images
if not self.otf_noise:
HR, LR, hr_crop_params, _ = random_crop_mod(HR, LR, LR_size, scale)
if self.y_only and self.srcolors:
LR_bicubic, _, _, _ = random_crop_mod(LR_bicubic, _, LR_size, scale, hr_crop_params)
HR_center, _, _, _ = random_crop_mod(HR_center, _, LR_size, scale, hr_crop_params)
# tmp_vis(LR_bicubic, False)
# tmp_vis(HR_center, False)
# data augmentation
#TODO: use BasicSR augmentations
#TODO: use variables from config
LR, HR, LR_bicubic, HR_center = augmentation()([LR, HR, LR_bicubic, HR_center])
# tmp_vis(HR, False)
# tmp_vis(LR, False)
# tmp_vis(LR_bicubic, False)
# tmp_vis(HR_center, False)
if self.y_only:
HR = util.np2tensor(HR, normalize=znorm, bgr2rgb=False, add_batch=False) # Tensor, [CT',H',W'] or [T, H, W]
LR = util.np2tensor(LR, normalize=znorm, bgr2rgb=False, add_batch=False) # Tensor, [CT',H',W'] or [T, H, W]
else:
HR = util.np2tensor(HR, normalize=znorm, bgr2rgb=True, add_batch=False) # Tensor, [CT',H',W'] or [T, H, W]
LR = util.np2tensor(LR, normalize=znorm, bgr2rgb=True, add_batch=False) # Tensor, [CT',H',W'] or [T, H, W]
#TODO: TMP to test generating 3 channel images for SR loss
# HR = util.np2tensor(HR, normalize=znorm, bgr2rgb=False, add_batch=True) # Tensor, [CT',H',W'] or [T, H, W]
# LR = util.np2tensor(LR, normalize=znorm, bgr2rgb=False, add_batch=True) # Tensor, [CT',H',W'] or [T, H, W]
# if self.srcolors:
# HR = HR.view(c,t,HR_size,HR_size) # Tensor, [C,T,H,W]
if not self.y_only:
HR = HR.view(c,t,HR_size,HR_size) # Tensor, [C,T,H,W]
LR = LR.view(c,t,LR_size,LR_size) # Tensor, [C,T,H,W]
if self.shape == 'TCHW':
HR = HR.transpose(0,1) # Tensor, [T,C,H,W]
LR = LR.transpose(0,1) # Tensor, [T,C,H,W]
# generate Cr, Cb channels using bicubic interpolation
#TODO: check, it might be easier to return the whole image and separate later when needed
if self.y_only and self.srcolors:
LR_bicubic = util.bgr2ycbcr(LR_bicubic, only_y=False)
# HR_center = util.bgr2ycbcr(HR_center, only_y=False) #not needed, can directly use rgb image
## LR_bicubic = util.ycbcr2rgb(LR_bicubic, only_y=False) #test, looks ok
## HR_center = util.ycbcr2rgb(HR_center, only_y=False) #test, looks ok
## _, SR_cb, SR_cr = util.bgr2ycbcr(LR_bicubic, only_y=False, separate=True)
LR_bicubic = util.np2tensor(LR_bicubic, normalize=znorm, bgr2rgb=False, add_batch=False)
# HR_center = util.np2tensor(HR_center, normalize=znorm, bgr2rgb=False, add_batch=False) # will test using rgb image instead
HR_center = util.np2tensor(HR_center, normalize=znorm, bgr2rgb=True, add_batch=False)
#TODO: TMP to test generating 3 channel images for SR loss
# LR_bicubic = util.np2tensor(LR_bicubic, normalize=znorm, bgr2rgb=False, add_batch=True)
# HR_center = util.np2tensor(HR_center, normalize=znorm, bgr2rgb=False, add_batch=True)
elif self.y_only and not self.srcolors:
LR_bicubic = []
HR_center = []
else:
HR_center = HR[:,idx_center,:,:] if self.shape == 'CTHW' else HR[idx_center,:,:,:]
LR_bicubic = []
# return toTensor(LR), toTensor(HR)
return {'LR': LR, 'HR': HR, 'LR_path': LR_dir, 'HR_path': HR_dir, 'LR_bicubic': LR_bicubic, 'HR_center': HR_center}
def __getitem__(self, idx):
scale = self.opt.get('scale', 4)
idx_center = (self.num_frames - 1) // 2
h_LR = None
w_LR = None
# Default case: tensor will result in the [0,1] range
# Alternative: tensor will be z-normalized to the [-1,1] range
znorm = self.opt.get('znorm', False)
# only one video and paths_LR/paths_HR is already the video dir
video_dir = ""
# list the frames in the directory
# hr_dir = self.trainset_dir + '/' + video_dir + '/hr'
'''
List based frames loading
'''
paths_LR = util.get_image_paths(self.opt['data_type'], os.path.join(self.paths_LR, video_dir))
assert self.num_frames <= len(paths_LR), (
f'num_frame must be smaller than the number of frames per video, check {video_dir}')
idx_frame = idx
LR_name = paths_LR[idx_frame + 1] # center frame
# print(LR_name)
# print(len(self.video_list))
# read LR frames
# HR_list = []
LR_list = []
resize_type = None
LR_bicubic = None
for i_frame in range(self.num_frames):
if idx_frame == len(self.video_list)-2 and self.num_frames == 3:
# print("second to last frame:", i_frame)
if i_frame == 0:
LR_img = util.read_img(None, paths_LR[int(idx_frame)], out_nc=self.image_channels)
else:
LR_img = util.read_img(None, paths_LR[int(idx_frame)+1], out_nc=self.image_channels)
elif idx_frame == len(self.video_list)-1 and self.num_frames == 3:
# print("last frame:", i_frame)
LR_img = util.read_img(None, paths_LR[int(idx_frame)], out_nc=self.image_channels)
# every other internal frame
else:
# print("normal frame:", idx_frame)
LR_img = util.read_img(None, paths_LR[int(idx_frame)+(i_frame)], out_nc=self.image_channels)
#TODO: check if this is necessary
LR_img = util.modcrop(LR_img, scale)
# get the bicubic upscale of the center frame to concatenate for SR
if not self.y_only and self.srcolors and i_frame == idx_center:
if self.opt.get('denoise_LRbic', False):
LR_bicubic = transforms.RandomAverageBlur(p=1, kernel_size=3)(LR_img)
# LR_bicubic = transforms.RandomBoxBlur(p=1, kernel_size=3)(LR_img)
else:
LR_bicubic = LR_img
LR_bicubic, _ = Scale(img=LR_bicubic, scale=1/scale, algo=777) # bicubic upscale
# HR_center = HR_img
# tmp_vis(LR_bicubic, False)
# tmp_vis(HR_center, False)
if self.y_only:
# extract Y channel from frames
# normal path, only Y for both
LR_img = util.bgr2ycbcr(LR_img, only_y=True)
# expand Y images to add the channel dimension
# normal path, only Y for both
LR_img = util.fix_img_channels(LR_img, 1)
# print("HR_img.shape: ", HR_img.shape)
# print("LR_img.shape", LR_img.shape)
LR_list.append(LR_img) # h, w, c
if not self.y_only and (not h_LR or not w_LR):
h_LR, w_LR, c = LR_img.shape
if not self.y_only:
t = self.num_frames
LR = [np.asarray(LT) for LT in LR_list] # list -> numpy # input: list (contatin numpy: [H,W,C])
LR = np.asarray(LR) # numpy, [T,H,W,C]
LR = LR.transpose(1,2,3,0).reshape(h_LR, w_LR, -1) # numpy, [Hl',Wl',CT]
else:
LR = np.concatenate((LR_list), axis=2) # h, w, t
if self.y_only:
LR = util.np2tensor(LR, normalize=znorm, bgr2rgb=False, add_batch=False) # Tensor, [CT',H',W'] or [T, H, W]
else:
LR = util.np2tensor(LR, normalize=znorm, bgr2rgb=True, add_batch=False) # Tensor, [CT',H',W'] or [T, H, W]
LR = LR.view(c,t,h_LR,w_LR) # Tensor, [C,T,H,W]
if self.shape == 'TCHW':
LR = LR.transpose(0,1) # Tensor, [T,C,H,W]
if self.y_only and self.srcolors:
# generate Cr, Cb channels using bicubic interpolation
LR_bicubic = util.bgr2ycbcr(LR_bicubic, only_y=False)
LR_bicubic = util.np2tensor(LR_bicubic, normalize=znorm, bgr2rgb=False, add_batch=False)
HR_center = []
else:
LR_bicubic = []
HR_center = []
# return toTensor(LR), toTensor(HR)
return {'LR': LR, 'LR_path': LR_name, 'LR_bicubic': LR_bicubic, 'HR_center': HR_center}
