def __getitem__(self, index):
path_LQ = self.data_info['path_LQ'][index]
path_GT = self.data_info['path_GT'][index]
folder = self.data_info['folder'][index]
idx, max_idx = self.data_info['idx'][index].split('/')
idx, max_idx = int(idx), int(max_idx)
if self.cache_data:
# select_idx = util.index_generation(idx, max_idx, self.opt['N_frames'], padding=self.opt['padding'])
# select_idx = [i for i in range(max_idx)]
# imgs_LQ = self.imgs_LQ[folder].index_select(0, torch.LongTensor(select_idx))
imgs_LQ = self.imgs_LQ[folder][idx]
img_GT = self.imgs_GT[folder][idx]
else:
imgs_LQ = util.read_img_seq(path_LQ)
img_GT = util.read_img_seq(path_GT)
return {
'LQ': imgs_LQ,
'GT': img_GT,
'folder': folder,
'idx': self.data_info['idx'][index],
'LQ_path': path_LQ,
'GT_path': path_GT
}
def __init__(self, opt):
super(VideoTestDataset, self).__init__()
self.opt = opt
self.cache_data = opt['cache_data']
self.half_N_frames = opt['N_frames'] // 2
self.GT_root, self.LQ_root = opt['dataroot_GT'], opt['dataroot_LQ']
self.data_type = self.opt['data_type']
self.data_info = {
'path_LQ': [],
'path_GT': [],
'folder': [],
'idx': [],
'border': []
}
if self.data_type == 'lmdb':
raise ValueError('No need to use LMDB during validation/test.')
#### Generate data info and cache data
self.imgs_LQ, self.imgs_GT = {}, {}
if opt['name'].lower() in ['vid4', 'reds4']:
if opt['name'].lower() == 'vid4':
subfolders_LQ = sorted(glob.glob(self.LQ_root))
subfolders_GT = sorted(glob.glob(self.GT_root))
else:
subfolders_LQ = util.glob_file_list(self.LQ_root)
subfolders_GT = util.glob_file_list(self.GT_root)
for subfolder_LQ, subfolder_GT in zip(subfolders_LQ,
subfolders_GT):
subfolder_name = osp.basename(subfolder_GT)
if subfolder_name in ['blur4', 'truth']:
subfolder_name = osp.basename(osp.dirname(subfolder_GT))
img_paths_LQ = util.glob_file_list(subfolder_LQ)
img_paths_GT = util.glob_file_list(subfolder_GT)
max_idx = len(img_paths_LQ)
assert max_idx == len(
img_paths_GT
), 'Different number of images in LQ and GT folders'
self.data_info['path_LQ'].extend(img_paths_LQ)
self.data_info['path_GT'].extend(img_paths_GT)
self.data_info['folder'].extend([subfolder_name] * max_idx)
for i in range(max_idx):
self.data_info['idx'].append('{}/{}'.format(i, max_idx))
border_l = [0] * max_idx
for i in range(self.half_N_frames):
border_l[i] = 1
border_l[max_idx - i - 1] = 1
self.data_info['border'].extend(border_l)
if self.cache_data:
self.imgs_LQ[subfolder_name] = util.read_img_seq(
img_paths_LQ)
self.imgs_GT[subfolder_name] = util.read_img_seq(
img_paths_GT)
elif opt['name'].lower() in ['vimeo90k-test']:
pass # TODO
else:
raise ValueError(
'Not support video test dataset. Support Vid4, REDS4 and Vimeo90k-Test.'
)
def __getitem__(self, index):
# path_LQ = self.data_info['path_LQ'][index]
# path_GT = self.data_info['path_GT'][index]
folder = self.data_info['folder'][index]
idx, max_idx = self.data_info['idx'][index].split('/')
idx, max_idx = int(idx), int(max_idx)
border = self.data_info['border'][index]
select_idx = util.index_generation(idx,
max_idx,
self.opt['N_frames'],
padding=self.opt['padding'])
if self.data_type == 'lmdb':
if self.GT_env is None or self.LQ_env is None:
self._init_lmdb()
key = self.lmdb_paths_GT[index]
name_a, name_b = key.split('_')
center_frame_idx = int(name_b)
GT_size_tuple = self.opt['GT_shape']
LQ_size_tuple = self.opt['LQ_shape']
img_GT = util.read_img(self.GT_env, key, GT_size_tuple)
img_LQ_l = []
for v in select_idx:
img_LQ = util.read_img(self.LQ_env,
'{}_{:08d}'.format(name_a,
v), LQ_size_tuple)
img_LQ_l.append(img_LQ)
# stack LQ images to NHWC, N is the frame number
img_LQs = np.stack(img_LQ_l, axis=0)
# BGR to RGB, HWC to CHW, numpy to tensor
img_GT = img_GT[:, :, [2, 1, 0]]
img_LQs = img_LQs[:, :, :, [2, 1, 0]]
img_GT = torch.from_numpy(
np.ascontiguousarray(np.transpose(img_GT, (2, 0, 1)))).float()
imgs_LQ = torch.from_numpy(
np.ascontiguousarray(np.transpose(img_LQs,
(0, 3, 1, 2)))).float()
elif self.cache_data:
imgs_LQ = self.imgs_LQ[folder].index_select(
0, torch.LongTensor(select_idx))
img_GT = self.imgs_GT[folder][idx]
else:
imgs_LQ = util.read_img_seq(self.imgs_LQ[folder]).index_select(
0, torch.LongTensor(select_idx))
img_GT = util.read_img_seq(self.imgs_GT[folder])[idx]
return {
'LQs': imgs_LQ,
'GT': img_GT,
'folder': folder,
'idx': self.data_info['idx'][index],
'border': border
}
def __init__(self, opt):
super(AI4KTestdataset, self).__init__()
self.opt = opt
self.cache_data = opt['cache_data']
self.half_N_frames = opt['N_frames'] // 2
self.GT_root, self.LQ_root = opt['dataroot_GT'], opt['dataroot_LQ']
self.data_type = self.opt['data_type']
self.data_info = {
'path_LQ': [],
'path_GT': [],
'folder': [],
'idx': []
}
if self.cache_data == False:
raise ValueError('Currently only support cache_data=True.')
if self.data_type == 'lmdb':
raise ValueError(
'Currently not support LMDB during validation/test.')
# scene change index dictionary, key example: XXXXXXXX, value example: [0, 10, 51, 100]
if opt['scene_index']:
self.scene_dict = pickle.load(open(opt['scene_index'], 'rb'))
else:
raise ValueError(
'Need to supply scene change index dictionary by running [cache_keys.py]'
)
#### Generate data info and cache data
self.imgs_LQ, self.imgs_GT = {}, {}
subfolders_LQ = util.glob_file_list(self.LQ_root)
subfolders_GT = util.glob_file_list(self.GT_root)
for subfolder_LQ, subfolder_GT in zip(subfolders_LQ, subfolders_GT):
subfolder_name = osp.basename(
subfolder_GT) # subfolders_name: XXXXXXXX
img_paths_LQ = util.glob_file_list(subfolder_LQ)
img_paths_GT = util.glob_file_list(subfolder_GT)
max_idx = len(img_paths_LQ)
assert len(img_paths_LQ) == len(
img_paths_GT
), 'Different number of images in LQ and GT folders'
self.data_info['path_LQ'].extend(img_paths_LQ)
self.data_info['path_GT'].extend(img_paths_GT)
self.data_info['folder'].extend([subfolder_name] * max_idx)
for i in range(max_idx):
self.data_info['idx'].append('{}/{}'.format(i, max_idx))
if self.cache_data:
self.imgs_LQ[subfolder_name] = util.read_img_seq(img_paths_LQ)
self.imgs_GT[subfolder_name] = util.read_img_seq(img_paths_GT)
def img2HQimg(self):
select_idx = data_util.index_generation(0, 1, 5, padding=self.padding)
imgs = data_util.read_img_seq(os.path.dirname(self.input_path))
imgs_in = imgs.index_select(0, torch.LongTensor(select_idx)).unsqueeze(0).to(self.device)
output = util.single_forward(self.model, imgs_in)
output = util.tensor2img(output.squeeze(0))
util.save_img(img=output, img_path=self.output_path)
def __init__(self, opt):
super(VideoSameSizeDataset, self).__init__()
self.opt = opt
self.cache_data = opt['cache_data']
self.half_N_frames = opt['N_frames'] // 2
self.GT_root, self.LQ_root = opt['dataroot_GT'], opt['dataroot_LQ']
self.data_type = self.opt['data_type']
self.data_info = {
'path_LQ': [],
'path_GT': [],
'folder': [],
'idx': [],
'border': []
}
if self.data_type == 'lmdb':
raise ValueError('No need to use LMDB during validation/test.')
# Generate data info and cache data
self.imgs_LQ, self.imgs_GT = {}, {}
# read data:
subfolders_LQ = util.glob_file_list(self.LQ_root)
subfolders_GT = util.glob_file_list(self.GT_root)
for subfolder_LQ, subfolder_GT in zip(subfolders_LQ, subfolders_GT):
# for frames in each video:
subfolder_name = osp.basename(subfolder_GT)
img_paths_LQ = util.glob_file_list(subfolder_LQ)
img_paths_GT = util.glob_file_list(subfolder_GT)
max_idx = len(img_paths_LQ)
assert max_idx == len(
img_paths_GT
), 'Different number of images in LQ and GT folders'
self.data_info['path_LQ'].extend(
img_paths_LQ) # list of path str of images
self.data_info['path_GT'].extend(img_paths_GT)
self.data_info['folder'].extend([subfolder_name] * max_idx)
for i in range(max_idx):
self.data_info['idx'].append('{}/{}'.format(i, max_idx))
border_l = [0] * max_idx
for i in range(self.half_N_frames):
border_l[i] = 1
border_l[max_idx - i - 1] = 1
self.data_info['border'].extend(border_l)
if self.cache_data:
self.imgs_LQ[subfolder_name] = util.read_img_seq(img_paths_LQ)
self.imgs_GT[subfolder_name] = util.read_img_seq(img_paths_GT)
def __init__(self, opt):
super(MultiImgTestDataset, self).__init__()
self.opt = opt
self.cache_data = opt['cache_data']
self.GT_root, self.LQ_root = opt['dataroot_GT'], opt['dataroot_LQ']
self.data_info = {
'path_LQ': [],
'path_GT': [],
'folder': [],
'idx': []
}
#### Generate data info and cache data
self.scale = opt['scale'][-1] if isinstance(opt['scale'],
(tuple,
list)) else opt['scale']
self.imgs_LQ, self.imgs_GT = {}, {}
if opt['name'].lower() in ['vid4', 'reds4']:
subfolders_LQ = util.glob_file_list(self.LQ_root)
subfolders_GT = util.glob_file_list(self.GT_root)
for subfolder_LQ, subfolder_GT in zip(subfolders_LQ,
subfolders_GT):
subfolder_name = osp.basename(subfolder_GT)
img_paths_LQ = util.glob_file_list(subfolder_LQ)
img_paths_GT = util.glob_file_list(subfolder_GT)
max_idx = len(img_paths_LQ)
assert max_idx == len(
img_paths_GT
), 'Different number of images in LQ and GT folders'
self.data_info['path_LQ'].extend(img_paths_LQ)
self.data_info['path_GT'].extend(img_paths_GT)
self.data_info['folder'].extend([subfolder_name] * max_idx)
for i in range(max_idx):
self.data_info['idx'].append('{}/{}'.format(i, max_idx))
if self.cache_data:
self.imgs_LQ[subfolder_name] = util.read_img_seq(
img_paths_LQ)
self.imgs_GT[subfolder_name] = util.read_img_seq(
img_paths_GT)
elif opt['name'].lower() in ['vimeo90k-test']:
pass # TODO
else:
raise ValueError(
'Not support video test dataset. Support Vid4, REDS4 and Vimeo90k-Test.'
)
def __init__(self, opt):
super(ImageTestDataset, self).__init__()
self.opt = opt
self.cache_data = opt['cache_data']
self.GT_root, self.LQ_root = opt['dataroot_GT'], opt['dataroot_LQ']
self.data_type = self.opt['data_type']
self.data_info = {'path_LQ': [], 'path_GT': []}
if self.data_type == 'lmdb':
raise ValueError('No need to use LMDB during validation/test.')
#### Generate data info and cache data
self.imgs_LQ, self.imgs_GT = {}, {}
img_paths_LQ = util.glob_file_list(self.LQ_root)
img_paths_GT = util.glob_file_list(self.GT_root)
assert len(img_paths_LQ) == len(
img_paths_GT), 'Different number of images in LQ and GT folders'
self.data_info['path_LQ'].extend(img_paths_LQ)
self.data_info['path_GT'].extend(img_paths_GT)
if self.cache_data:
self.imgs_LQ = util.read_img_seq(img_paths_LQ)
self.imgs_GT = util.read_img_seq(img_paths_GT)
def __init__(self, opt, **kwargs):
super(Demo, self).__init__()
self.scale = kwargs['scale']
self.kernel_size = kwargs['kernel_size']
self.model_name = kwargs['model_name']
idx = kwargs['idx'] if 'idx' in kwargs else None
self.opt = opt
self.cache_data = opt['cache_data']
self.half_N_frames = opt['N_frames'] // 2
if idx is None:
self.name = opt['name']
self.root = opt['dataroot']
else:
self.name = opt['name'].split('+')[idx]
self.root = opt['dataroot'].split('+')[idx]
self.data_type = self.opt['data_type']
self.data_info = {'path': [], 'folder': [], 'idx': [], 'border': []}
if self.data_type == 'lmdb':
raise ValueError('No need to use LMDB during validation/test.')
#### Generate data info and cache data
self.imgs = {}
subfolder_name = osp.basename(self.root)
self.data_info['path'] = util.glob_file_list(self.root)
max_idx = len(self.data_info['path'])
self.data_info['folder'] = [subfolder_name] * max_idx
for i in range(max_idx):
self.data_info['idx'].append('{}/{}'.format(i, max_idx))
border_l = [0] * max_idx
for i in range(self.half_N_frames):
border_l[i] = 1
border_l[max_idx - i - 1] = 1
self.data_info['border'].extend(border_l)
self.imgs[subfolder_name] = util.read_img_seq(self.data_info['path'], 'img')
def val(model_name, current_step, arch='EDVR'):
#################
# configurations
#################
device = torch.device('cuda')
#os.environ['CUDA_VISIBLE_DEVICES'] = '1,2,3,4'
test_set = 'REDS4' # Vid4 | YouKu10 | REDS4 | AI4K_val
data_mode = 'sharp_bicubic' # sharp_bicubic | blur_bicubic
N_in = 5
# load test set
if test_set == 'Vid4':
test_dataset_folder = '../datasets/Vid4/BIx4'
GT_dataset_folder = '../datasets/Vid4/GT'
elif test_set == 'YouKu10':
test_dataset_folder = '../datasets/YouKu10/LR'
GT_dataset_folder = '../datasets/YouKu10/HR'
elif test_set == 'YouKu_val':
test_dataset_folder = '/data0/yhliu/DATA/YouKuVid/valid/valid_lr_bmp'
GT_dataset_folder = '/data0/yhliu/DATA/YouKuVid/valid/valid_hr_bmp'
elif test_set == 'REDS4':
test_dataset_folder = '../datasets/REDS4/{}'.format(data_mode)
GT_dataset_folder = '../datasets/REDS4/GT'
elif test_set == 'AI4K_val':
test_dataset_folder = '/data0/yhliu/AI4K/contest1/val1_LR_png/'
GT_dataset_folder = '/data0/yhliu/AI4K/contest1/val1_HR_png/'
elif test_set == 'AI4K_val_small':
test_dataset_folder = '/home/yhliu/AI4K/contest1/val1_LR_png_small/'
GT_dataset_folder = '/home/yhliu/AI4K/contest1/val1_HR_png_small/'
flip_test = False
#model_path = '../experiments/pretrained_models/EDVR_Vimeo90K_SR_L.pth'
model_path = os.path.join('../experiments/', model_name,
'models/{}_G.pth'.format(current_step))
predeblur, HR_in = False, False
back_RBs = 10
if data_mode == 'blur_bicubic':
predeblur = True
if data_mode == 'blur' or data_mode == 'blur_comp':
predeblur, HR_in = True, True
if arch == 'EDVR':
model = EDVR_arch.EDVR(64,
N_in,
8,
5,
back_RBs,
predeblur=predeblur,
HR_in=HR_in)
elif arch == 'MY_EDVR':
model = my_EDVR_arch.MYEDVR(64,
N_in,
8,
5,
back_RBs,
predeblur=predeblur,
HR_in=HR_in)
#### evaluation
crop_border = 0
border_frame = N_in // 2 # border frames when evaluate
# temporal padding mode
if data_mode == 'Vid4' or data_mode == 'sharp_bicubic':
padding = 'new_info'
else:
padding = 'replicate'
save_imgs = False
save_folder = '../validation/{}'.format(test_set)
util.mkdirs(save_folder)
util.setup_logger('base',
save_folder,
'test',
level=logging.INFO,
screen=True,
tofile=True)
logger = logging.getLogger('base')
#### log info
logger.info('Data: {} - {}'.format(data_mode, test_dataset_folder))
logger.info('Padding mode: {}'.format(padding))
logger.info('Model path: {}'.format(model_path))
logger.info('Save images: {}'.format(save_imgs))
logger.info('Flip test: {}'.format(flip_test))
#### set up the models
model.load_state_dict(torch.load(model_path), strict=True)
model.eval()
model = model.to(device)
model = nn.DataParallel(model, device_ids=[0, 1, 2, 3])
avg_psnr_l, avg_psnr_center_l, avg_psnr_border_l = [], [], []
subfolder_name_l = []
subfolder_l = sorted(glob.glob(osp.join(test_dataset_folder, '*')))
subfolder_GT_l = sorted(glob.glob(osp.join(GT_dataset_folder, '*')))
#print(subfolder_l)
#print(subfolder_GT_l)
#exit()
# for each subfolder
for subfolder, subfolder_GT in zip(subfolder_l, subfolder_GT_l):
subfolder_name = osp.basename(subfolder)
subfolder_name_l.append(subfolder_name)
save_subfolder = osp.join(save_folder, subfolder_name)
img_path_l = sorted(glob.glob(osp.join(subfolder, '*')))
#print(img_path_l)
max_idx = len(img_path_l)
if save_imgs:
util.mkdirs(save_subfolder)
#### read LQ and GT images
imgs_LQ = data_util.read_img_seq(subfolder)
img_GT_l = []
for img_GT_path in sorted(glob.glob(osp.join(subfolder_GT, '*'))):
#print(img_GT_path)
img_GT_l.append(data_util.read_img(None, img_GT_path))
#print(img_GT_l[0].shape)
avg_psnr, avg_psnr_border, avg_psnr_center, N_border, N_center = 0, 0, 0, 0, 0
# process each image
for img_idx, img_path in enumerate(img_path_l):
img_name = osp.splitext(osp.basename(img_path))[0]
select_idx = data_util.index_generation(img_idx,
max_idx,
N_in,
padding=padding)
imgs_in = imgs_LQ.index_select(
0, torch.LongTensor(select_idx)).unsqueeze(0).to(device)
#print(imgs_in.size())
if flip_test:
output = util.flipx4_forward(model, imgs_in)
else:
output = util.single_forward(model, imgs_in)
output = util.tensor2img(output.squeeze(0))
if save_imgs:
cv2.imwrite(
osp.join(save_subfolder, '{}.png'.format(img_name)),
output)
# calculate PSNR
output = output / 255.
GT = np.copy(img_GT_l[img_idx])
# For REDS, evaluate on RGB channels; for Vid4, evaluate on the Y channel
'''
if data_mode == 'Vid4': # bgr2y, [0, 1]
GT = data_util.bgr2ycbcr(GT, only_y=True)
output = data_util.bgr2ycbcr(output, only_y=True)
'''
output, GT = util.crop_border([output, GT], crop_border)
crt_psnr = util.calculate_psnr(output * 255, GT * 255)
#logger.info('{:3d} - {:25} \tPSNR: {:.6f} dB'.format(img_idx + 1, img_name, crt_psnr))
if img_idx >= border_frame and img_idx < max_idx - border_frame: # center frames
avg_psnr_center += crt_psnr
N_center += 1
else: # border frames
avg_psnr_border += crt_psnr
N_border += 1
avg_psnr = (avg_psnr_center + avg_psnr_border) / (N_center + N_border)
avg_psnr_center = avg_psnr_center / N_center
avg_psnr_border = 0 if N_border == 0 else avg_psnr_border / N_border
avg_psnr_l.append(avg_psnr)
avg_psnr_center_l.append(avg_psnr_center)
avg_psnr_border_l.append(avg_psnr_border)
logger.info('Folder {} - Average PSNR: {:.6f} dB for {} frames; '
'Center PSNR: {:.6f} dB for {} frames; '
'Border PSNR: {:.6f} dB for {} frames.'.format(
subfolder_name, avg_psnr, (N_center + N_border),
avg_psnr_center, N_center, avg_psnr_border, N_border))
logger.info('################ Tidy Outputs ################')
for subfolder_name, psnr, psnr_center, psnr_border in zip(
subfolder_name_l, avg_psnr_l, avg_psnr_center_l,
avg_psnr_border_l):
logger.info('Folder {} - Average PSNR: {:.6f} dB. '
'Center PSNR: {:.6f} dB. '
'Border PSNR: {:.6f} dB.'.format(subfolder_name, psnr,
psnr_center, psnr_border))
logger.info('################ Final Results ################')
logger.info('Data: {} - {}'.format(data_mode, test_dataset_folder))
logger.info('Padding mode: {}'.format(padding))
logger.info('Model path: {}'.format(model_path))
logger.info('Save images: {}'.format(save_imgs))
logger.info('Flip test: {}'.format(flip_test))
logger.info('Total Average PSNR: {:.6f} dB for {} clips. '
'Center PSNR: {:.6f} dB. Border PSNR: {:.6f} dB.'.format(
sum(avg_psnr_l) / len(avg_psnr_l), len(subfolder_l),
sum(avg_psnr_center_l) / len(avg_psnr_center_l),
sum(avg_psnr_border_l) / len(avg_psnr_border_l)))
return sum(avg_psnr_l) / len(avg_psnr_l)
def main():
####################
# arguments parser #
####################
# [format] dataset(vid4, REDS4) N(number of frames)
# data_mode = str(args.dataset)
# N_in = int(args.n_frames)
# metrics = str(args.metrics)
# output_format = str(args.output_format)
#################
# configurations
#################
device = torch.device('cuda')
os.environ['CUDA_VISIBLE_DEVICES'] = '0'
#data_mode = 'Vid4' # Vid4 | sharp_bicubic | blur_bicubic | blur | blur_comp
# Vid4: SR
# REDS4: sharp_bicubic (SR-clean), blur_bicubic (SR-blur);
# blur (deblur-clean), blur_comp (deblur-compression).
# STAGE Vid4
# Collecting results for Vid4
model_path = '../experiments/pretrained_models/EDVR_Vimeo90K_SR_L.pth'
stage = 1 # 1 or 2, use two stage strategy for REDS dataset.
flip_test = False
predeblur, HR_in = False, False
back_RBs = 40
N_model_default = 7
data_mode = 'Vid4'
# vid4_dir_map = {"calendar": 0, "city": 1, "foliage": 2, "walk": 3}
vid4_results = {"calendar": {}, "city": {}, "foliage": {}, "walk": {}}
#vid4_results = 4 * [[]]
for N_in in range(1, N_model_default + 1):
raw_model = EDVR_arch.EDVR(128, N_model_default, 8, 5, back_RBs, predeblur=predeblur, HR_in=HR_in)
model = EDVR_arch.EDVR(128, N_in, 8, 5, back_RBs, predeblur=predeblur, HR_in=HR_in)
test_dataset_folder = '../datasets/Vid4/BIx4'
GT_dataset_folder = '../datasets/Vid4/GT'
aposterior_GT_dataset_folder = '../datasets/Vid4/GT_7'
crop_border = 0
border_frame = N_in // 2 # border frames when evaluate
padding = 'new_info'
save_imgs = False
raw_model.load_state_dict(torch.load(model_path), strict=True)
model.nf = raw_model.nf
model.center = N_in // 2 # if center is None else center
model.is_predeblur = raw_model.is_predeblur
model.HR_in = raw_model.HR_in
model.w_TSA = raw_model.w_TSA
if model.is_predeblur:
model.pre_deblur = raw_model.pre_deblur # Predeblur_ResNet_Pyramid(nf=nf, HR_in=self.HR_in)
model.conv_1x1 = raw_model.conv_1x1 # nn.Conv2d(nf, nf, 1, 1, bias=True)
else:
if model.HR_in:
model.conv_first_1 = raw_model.conv_first_1 # nn.Conv2d(3, nf, 3, 1, 1, bias=True)
model.conv_first_2 = raw_model.conv_first_2 # nn.Conv2d(nf, nf, 3, 2, 1, bias=True)
model.conv_first_3 = raw_model.conv_first_3 # nn.Conv2d(nf, nf, 3, 2, 1, bias=True)
else:
model.conv_first = raw_model.conv_first # nn.Conv2d(3, nf, 3, 1, 1, bias=True)
model.feature_extraction = raw_model.feature_extraction # arch_util.make_layer(ResidualBlock_noBN_f, front_RBs)
model.fea_L2_conv1 = raw_model.fea_L2_conv1 # nn.Conv2d(nf, nf, 3, 2, 1, bias=True)
model.fea_L2_conv2 = raw_model.fea_L2_conv2 # nn.Conv2d(nf, nf, 3, 1, 1, bias=True)
model.fea_L3_conv1 = raw_model.fea_L3_conv1 # nn.Conv2d(nf, nf, 3, 2, 1, bias=True)
model.fea_L3_conv2 = raw_model.fea_L3_conv2 # nn.Conv2d(nf, nf, 3, 1, 1, bias=True)
model.pcd_align = raw_model.pcd_align # PCD_Align(nf=nf, groups=groups)
model.tsa_fusion.center = model.center
model.tsa_fusion.tAtt_1 = raw_model.tsa_fusion.tAtt_1
model.tsa_fusion.tAtt_2 = raw_model.tsa_fusion.tAtt_2
model.tsa_fusion.fea_fusion = copy.deepcopy(raw_model.tsa_fusion.fea_fusion)
model.tsa_fusion.fea_fusion.weight = copy.deepcopy(torch.nn.Parameter(raw_model.tsa_fusion.fea_fusion.weight[:, 0:N_in * 128, :, :]))
model.tsa_fusion.sAtt_1 = copy.deepcopy(raw_model.tsa_fusion.sAtt_1)
model.tsa_fusion.sAtt_1.weight = copy.deepcopy(torch.nn.Parameter(raw_model.tsa_fusion.sAtt_1.weight[:, 0:N_in * 128, :, :]))
model.tsa_fusion.maxpool = raw_model.tsa_fusion.maxpool
model.tsa_fusion.avgpool = raw_model.tsa_fusion.avgpool
model.tsa_fusion.sAtt_2 = raw_model.tsa_fusion.sAtt_2
model.tsa_fusion.sAtt_3 = raw_model.tsa_fusion.sAtt_3
model.tsa_fusion.sAtt_4 = raw_model.tsa_fusion.sAtt_4
model.tsa_fusion.sAtt_5 = raw_model.tsa_fusion.sAtt_5
model.tsa_fusion.sAtt_L1 = raw_model.tsa_fusion.sAtt_L1
model.tsa_fusion.sAtt_L2 = raw_model.tsa_fusion.sAtt_L2
model.tsa_fusion.sAtt_L3 = raw_model.tsa_fusion.sAtt_L3
model.tsa_fusion.sAtt_add_1 = raw_model.tsa_fusion.sAtt_add_1
model.tsa_fusion.sAtt_add_2 = raw_model.tsa_fusion.sAtt_add_2
model.tsa_fusion.lrelu = raw_model.tsa_fusion.lrelu
model.recon_trunk = raw_model.recon_trunk
model.upconv1 = raw_model.upconv1
model.upconv2 = raw_model.upconv2
model.pixel_shuffle = raw_model.pixel_shuffle
model.HRconv = raw_model.HRconv
model.conv_last = raw_model.conv_last
model.lrelu = raw_model.lrelu
#####################################################
model.eval()
model = model.to(device)
#avg_psnr_l, avg_psnr_center_l, avg_psnr_border_l = [], [], []
subfolder_name_l = []
subfolder_l = sorted(glob.glob(osp.join(test_dataset_folder, '*')))
subfolder_GT_l = sorted(glob.glob(osp.join(GT_dataset_folder, '*')))
subfolder_GT_a_l = sorted(glob.glob(osp.join(aposterior_GT_dataset_folder, "*")))
# for each subfolder
for subfolder, subfolder_GT, subfolder_GT_a in zip(subfolder_l, subfolder_GT_l, subfolder_GT_a_l):
subfolder_name = osp.basename(subfolder)
subfolder_name_l.append(subfolder_name)
img_path_l = sorted(glob.glob(osp.join(subfolder, '*')))
max_idx = len(img_path_l)
print("MAX_IDX: ", max_idx)
#### read LQ and GT images
imgs_LQ = data_util.read_img_seq(subfolder)
img_GT_l = []
for img_GT_path in sorted(glob.glob(osp.join(subfolder_GT, '*'))):
img_GT_l.append(data_util.read_img(None, img_GT_path))
img_GT_a = []
for img_GT_a_path in sorted(glob.glob(osp.join(subfolder_GT_a, '*'))):
img_GT_a.append(data_util.read_img(None, img_GT_a_path))
#avg_psnr, avg_psnr_border, avg_psnr_center, N_border, N_center = 0, 0, 0, 0, 0
# process each image
for img_idx, img_path in enumerate(img_path_l):
img_name = osp.splitext(osp.basename(img_path))[0]
select_idx = data_util.index_generation(img_idx, max_idx, N_in, padding=padding)
imgs_in = imgs_LQ.index_select(0, torch.LongTensor(select_idx)).unsqueeze(0).to(device)
if flip_test:
output = util.flipx4_forward(model, imgs_in)
else:
print("IMGS_IN SHAPE: ", imgs_in.shape)
output = util.single_forward(model, imgs_in)
output = util.tensor2img(output.squeeze(0))
# calculate PSNR
output = output / 255.
GT = np.copy(img_GT_l[img_idx])
# For REDS, evaluate on RGB channels; for Vid4, evaluate on the Y channel
#if data_mode == 'Vid4': # bgr2y, [0, 1]
GT = data_util.bgr2ycbcr(GT, only_y=True)
output = data_util.bgr2ycbcr(output, only_y=True)
GT_a = np.copy(img_GT_a[img_idx])
GT_a = data_util.bgr2ycbcr(GT_a, only_y=True)
output_a = copy.deepcopy(output)
output, GT = util.crop_border([output, GT], crop_border)
crt_psnr = util.calculate_psnr(output * 255, GT * 255)
crt_ssim = util.calculate_ssim(output * 255, GT * 255)
output_a, GT_a = util.crop_border([output_a, GT_a], crop_border)
crt_aposterior = util.calculate_ssim(output_a * 255, GT_a * 255) # CHANGE
t = vid4_results[subfolder_name].get(str(img_name))
if t != None:
vid4_results[subfolder_name][img_name].add_psnr(crt_psnr)
vid4_results[subfolder_name][img_name].add_gt_ssim(crt_ssim)
vid4_results[subfolder_name][img_name].add_aposterior_ssim(crt_aposterior)
else:
vid4_results[subfolder_name].update({img_name: metrics_file(img_name)})
vid4_results[subfolder_name][img_name].add_psnr(crt_psnr)
vid4_results[subfolder_name][img_name].add_gt_ssim(crt_ssim)
vid4_results[subfolder_name][img_name].add_aposterior_ssim(crt_aposterior)
############################################################################
#### model
#### writing vid4 results
util.mkdirs('../results/calendar')
util.mkdirs('../results/city')
util.mkdirs('../results/foliage')
util.mkdirs('../results/walk')
save_folder = '../results/'
for i, dir_name in enumerate(["calendar", "city", "foliage", "walk"]):
save_subfolder = osp.join(save_folder, dir_name)
for j, value in vid4_results[dir_name].items():
# cur_result = json.dumps(_)
with open(osp.join(save_subfolder, '{}.json'.format(value.name)), 'w') as outfile:
json.dump(value.__dict__, outfile, ensure_ascii=False, indent=4)
#json.dump(cur_result, outfile)
#cv2.imwrite(osp.join(save_subfolder, '{}.png'.format(img_name)), output)
###################################################################################
# STAGE REDS
reds4_results = {"000": {}, "011": {}, "015": {}, "020": {}}
data_mode = 'sharp_bicubic'
N_model_default = 5
for N_in in range(1, N_model_default + 1):
for stage in range(1,3):
flip_test = False
if data_mode == 'sharp_bicubic':
if stage == 1:
model_path = '../experiments/pretrained_models/EDVR_REDS_SR_L.pth'
else:
model_path = '../experiments/pretrained_models/EDVR_REDS_SR_Stage2.pth'
elif data_mode == 'blur_bicubic':
if stage == 1:
model_path = '../experiments/pretrained_models/EDVR_REDS_SRblur_L.pth'
else:
model_path = '../experiments/pretrained_models/EDVR_REDS_SRblur_Stage2.pth'
elif data_mode == 'blur':
if stage == 1:
model_path = '../experiments/pretrained_models/EDVR_REDS_deblur_L.pth'
else:
model_path = '../experiments/pretrained_models/EDVR_REDS_deblur_Stage2.pth'
elif data_mode == 'blur_comp':
if stage == 1:
model_path = '../experiments/pretrained_models/EDVR_REDS_deblurcomp_L.pth'
else:
model_path = '../experiments/pretrained_models/EDVR_REDS_deblurcomp_Stage2.pth'
else:
raise NotImplementedError
predeblur, HR_in = False, False
back_RBs = 40
if data_mode == 'blur_bicubic':
predeblur = True
if data_mode == 'blur' or data_mode == 'blur_comp':
predeblur, HR_in = True, True
if stage == 2:
HR_in = True
back_RBs = 20
if stage == 1:
test_dataset_folder = '../datasets/REDS4/{}'.format(data_mode)
else:
test_dataset_folder = '../results/REDS-EDVR_REDS_SR_L_flipx4'
print('You should modify the test_dataset_folder path for stage 2')
GT_dataset_folder = '../datasets/REDS4/GT'
raw_model = EDVR_arch.EDVR(128, N_model_default, 8, 5, back_RBs, predeblur=predeblur, HR_in=HR_in)
model = EDVR_arch.EDVR(128, N_in, 8, 5, back_RBs, predeblur=predeblur, HR_in=HR_in)
crop_border = 0
border_frame = N_in // 2 # border frames when evaluate
# temporal padding mode
if data_mode == 'Vid4' or data_mode == 'sharp_bicubic':
padding = 'new_info'
else:
padding = 'replicate'
save_imgs = True
data_mode_t = copy.deepcopy(data_mode)
if stage == 1 and data_mode_t != 'Vid4':
data_mode = 'REDS-EDVR_REDS_SR_L_flipx4'
save_folder = '../results/{}'.format(data_mode)
data_mode = copy.deepcopy(data_mode_t)
util.mkdirs(save_folder)
util.setup_logger('base', save_folder, 'test', level=logging.INFO, screen=True, tofile=True)
aposterior_GT_dataset_folder = '../datasets/REDS4/GT_5'
crop_border = 0
border_frame = N_in // 2 # border frames when evaluate
raw_model.load_state_dict(torch.load(model_path), strict=True)
model.nf = raw_model.nf
model.center = N_in // 2 # if center is None else center
model.is_predeblur = raw_model.is_predeblur
model.HR_in = raw_model.HR_in
model.w_TSA = raw_model.w_TSA
if model.is_predeblur:
model.pre_deblur = raw_model.pre_deblur # Predeblur_ResNet_Pyramid(nf=nf, HR_in=self.HR_in)
model.conv_1x1 = raw_model.conv_1x1 # nn.Conv2d(nf, nf, 1, 1, bias=True)
else:
if model.HR_in:
model.conv_first_1 = raw_model.conv_first_1 # nn.Conv2d(3, nf, 3, 1, 1, bias=True)
model.conv_first_2 = raw_model.conv_first_2 # nn.Conv2d(nf, nf, 3, 2, 1, bias=True)
model.conv_first_3 = raw_model.conv_first_3 # nn.Conv2d(nf, nf, 3, 2, 1, bias=True)
else:
model.conv_first = raw_model.conv_first # nn.Conv2d(3, nf, 3, 1, 1, bias=True)
model.feature_extraction = raw_model.feature_extraction # arch_util.make_layer(ResidualBlock_noBN_f, front_RBs)
model.fea_L2_conv1 = raw_model.fea_L2_conv1 # nn.Conv2d(nf, nf, 3, 2, 1, bias=True)
model.fea_L2_conv2 = raw_model.fea_L2_conv2 # nn.Conv2d(nf, nf, 3, 1, 1, bias=True)
model.fea_L3_conv1 = raw_model.fea_L3_conv1 # nn.Conv2d(nf, nf, 3, 2, 1, bias=True)
model.fea_L3_conv2 = raw_model.fea_L3_conv2 # nn.Conv2d(nf, nf, 3, 1, 1, bias=True)
model.pcd_align = raw_model.pcd_align # PCD_Align(nf=nf, groups=groups)
model.tsa_fusion.center = model.center
model.tsa_fusion.tAtt_1 = raw_model.tsa_fusion.tAtt_1
model.tsa_fusion.tAtt_2 = raw_model.tsa_fusion.tAtt_2
model.tsa_fusion.fea_fusion = copy.deepcopy(raw_model.tsa_fusion.fea_fusion)
model.tsa_fusion.fea_fusion.weight = copy.deepcopy(torch.nn.Parameter(raw_model.tsa_fusion.fea_fusion.weight[:, 0:N_in * 128, :, :]))
model.tsa_fusion.sAtt_1 = copy.deepcopy(raw_model.tsa_fusion.sAtt_1)
model.tsa_fusion.sAtt_1.weight = copy.deepcopy(torch.nn.Parameter(raw_model.tsa_fusion.sAtt_1.weight[:, 0:N_in * 128, :, :]))
model.tsa_fusion.maxpool = raw_model.tsa_fusion.maxpool
model.tsa_fusion.avgpool = raw_model.tsa_fusion.avgpool
model.tsa_fusion.sAtt_2 = raw_model.tsa_fusion.sAtt_2
model.tsa_fusion.sAtt_3 = raw_model.tsa_fusion.sAtt_3
model.tsa_fusion.sAtt_4 = raw_model.tsa_fusion.sAtt_4
model.tsa_fusion.sAtt_5 = raw_model.tsa_fusion.sAtt_5
model.tsa_fusion.sAtt_L1 = raw_model.tsa_fusion.sAtt_L1
model.tsa_fusion.sAtt_L2 = raw_model.tsa_fusion.sAtt_L2
model.tsa_fusion.sAtt_L3 = raw_model.tsa_fusion.sAtt_L3
model.tsa_fusion.sAtt_add_1 = raw_model.tsa_fusion.sAtt_add_1
model.tsa_fusion.sAtt_add_2 = raw_model.tsa_fusion.sAtt_add_2
model.tsa_fusion.lrelu = raw_model.tsa_fusion.lrelu
model.recon_trunk = raw_model.recon_trunk
model.upconv1 = raw_model.upconv1
model.upconv2 = raw_model.upconv2
model.pixel_shuffle = raw_model.pixel_shuffle
model.HRconv = raw_model.HRconv
model.conv_last = raw_model.conv_last
model.lrelu = raw_model.lrelu
#####################################################
model.eval()
model = model.to(device)
#avg_psnr_l, avg_psnr_center_l, avg_psnr_border_l = [], [], []
subfolder_name_l = []
subfolder_l = sorted(glob.glob(osp.join(test_dataset_folder, '*')))
subfolder_GT_l = sorted(glob.glob(osp.join(GT_dataset_folder, '*')))
subfolder_GT_a_l = sorted(glob.glob(osp.join(aposterior_GT_dataset_folder, "*")))
# for each subfolder
for subfolder, subfolder_GT, subfolder_GT_a in zip(subfolder_l, subfolder_GT_l, subfolder_GT_a_l):
subfolder_name = osp.basename(subfolder)
subfolder_name_l.append(subfolder_name)
save_subfolder = osp.join(save_folder, subfolder_name)
img_path_l = sorted(glob.glob(osp.join(subfolder, '*')))
max_idx = len(img_path_l)
print("MAX_IDX: ", max_idx)
print("SAVE FOLDER::::::", save_folder)
if save_imgs:
util.mkdirs(save_subfolder)
#### read LQ and GT images
imgs_LQ = data_util.read_img_seq(subfolder)
img_GT_l = []
for img_GT_path in sorted(glob.glob(osp.join(subfolder_GT, '*'))):
img_GT_l.append(data_util.read_img(None, img_GT_path))
img_GT_a = []
for img_GT_a_path in sorted(glob.glob(osp.join(subfolder_GT_a, '*'))):
img_GT_a.append(data_util.read_img(None, img_GT_a_path))
#avg_psnr, avg_psnr_border, avg_psnr_center, N_border, N_center = 0, 0, 0, 0, 0
# process each image
for img_idx, img_path in enumerate(img_path_l):
img_name = osp.splitext(osp.basename(img_path))[0]
select_idx = data_util.index_generation(img_idx, max_idx, N_in, padding=padding)
imgs_in = imgs_LQ.index_select(0, torch.LongTensor(select_idx)).unsqueeze(0).to(device)
if flip_test:
output = util.flipx4_forward(model, imgs_in)
else:
print("IMGS_IN SHAPE: ", imgs_in.shape)
output = util.single_forward(model, imgs_in)
output = util.tensor2img(output.squeeze(0))
if save_imgs and stage == 1:
cv2.imwrite(osp.join(save_subfolder, '{}.png'.format(img_name)), output)
# calculate PSNR
if stage == 2:
output = output / 255.
GT = np.copy(img_GT_l[img_idx])
# For REDS, evaluate on RGB channels; for Vid4, evaluate on the Y channel
#if data_mode == 'Vid4': # bgr2y, [0, 1]
GT_a = np.copy(img_GT_a[img_idx])
output_a = copy.deepcopy(output)
output, GT = util.crop_border([output, GT], crop_border)
crt_psnr = util.calculate_psnr(output * 255, GT * 255)
crt_ssim = util.calculate_ssim(output * 255, GT * 255)
output_a, GT_a = util.crop_border([output_a, GT_a], crop_border)
crt_aposterior = util.calculate_ssim(output_a * 255, GT_a * 255) # CHANGE
t = reds4_results[subfolder_name].get(str(img_name))
if t != None:
reds4_results[subfolder_name][img_name].add_psnr(crt_psnr)
reds4_results[subfolder_name][img_name].add_gt_ssim(crt_ssim)
reds4_results[subfolder_name][img_name].add_aposterior_ssim(crt_aposterior)
else:
reds4_results[subfolder_name].update({img_name: metrics_file(img_name)})
reds4_results[subfolder_name][img_name].add_psnr(crt_psnr)
reds4_results[subfolder_name][img_name].add_gt_ssim(crt_ssim)
reds4_results[subfolder_name][img_name].add_aposterior_ssim(crt_aposterior)
############################################################################
#### model
#### writing reds4 results
util.mkdirs('../results/000')
util.mkdirs('../results/011')
util.mkdirs('../results/015')
util.mkdirs('../results/020')
save_folder = '../results/'
for i, dir_name in enumerate(["000", "011", "015", "020"]): # +
save_subfolder = osp.join(save_folder, dir_name)
for j, value in reds4_results[dir_name].items():
# cur_result = json.dumps(value.__dict__)
with open(osp.join(save_subfolder, '{}.json'.format(value.name)), 'w') as outfile:
json.dump(value.__dict__, outfile, ensure_ascii=False, indent=4)
def main():
#################
# configurations
#################
#torch.backends.cudnn.benchmark = True
#torch.backends.cudnn.enabled = True
device = torch.device('cuda')
os.environ['CUDA_VISIBLE_DEVICES'] = '5'
test_set = 'AI4K_test' # Vid4 | YouKu10 | REDS4 | AI4K_test
data_mode = 'sharp_bicubic' # sharp_bicubic | blur_bicubic
test_name = 'Contest2_Test18_A38_color_EDVR_35_220000_A01_5in_64f_10b_128_pretrain_A01xxx_900000_fix_before_pcd_165000' #'AI4K_TEST_Denoise_A02_265000' | AI4K_test_A01b_145000
N_in = 5
# load test set
if test_set == 'AI4K_test':
#test_dataset_folder = '/data1/yhliu/AI4K/Corrected_TestA_Contest2_001_ResNet_alpha_beta_gaussian_65000/' #'/data1/yhliu/AI4K/testA_LR_png/'
test_dataset_folder = '/home/yhliu/AI4K/contest2/testA_LR_png/'
flip_test = False #False
#model_path = '../experiments/pretrained_models/EDVR_Vimeo90K_SR_L.pth'
#model_path = '../experiments/002_EDVR_EDVRwoTSAIni_lr4e-4_600k_REDS_LrCAR4S_fixTSA50k_new/models/latest_G.pth'
#model_path = '../experiments/A02_predenoise/models/415000_G.pth'
model_path = '../experiments/A38_color_EDVR_35_220000_A01_5in_64f_10b_128_pretrain_A01xxx_900000_fix_before_pcd/models/165000_G.pth'
color_model_path = '/home/yhliu/BasicSR/experiments/35_ResNet_alpha_beta_decoder_3x3_IN_encoder_8HW_re_100k/models/220000_G.pth'
predeblur, HR_in = False, False
back_RBs = 10
if data_mode == 'blur_bicubic':
predeblur = True
if data_mode == 'blur' or data_mode == 'blur_comp':
predeblur, HR_in = True, True
model = EDVR_arch.EDVR(64,
N_in,
8,
5,
back_RBs,
predeblur=predeblur,
HR_in=HR_in)
#model = my_EDVR_arch.MYEDVR_FusionDenoise(64, N_in, 8, 5, back_RBs, predeblur=predeblur, HR_in=HR_in, deconv=False)
color_model = SRResNet_arch.ResNet_alpha_beta_multi_in(
structure='ResNet_alpha_beta_decoder_3x3_IN_encoder_8HW')
#### evaluation
crop_border = 0
border_frame = N_in // 2 # border frames when evaluate
# temporal padding mode
if data_mode == 'Vid4' or data_mode == 'sharp_bicubic':
padding = 'new_info'
else:
padding = 'replicate'
save_imgs = True
save_folder = '../results/{}'.format(test_name)
util.mkdirs(save_folder)
util.setup_logger('base',
save_folder,
'test',
level=logging.INFO,
screen=True,
tofile=True)
logger = logging.getLogger('base')
#### log info
logger.info('Data: {} - {}'.format(data_mode, test_dataset_folder))
logger.info('Padding mode: {}'.format(padding))
logger.info('Model path: {}'.format(model_path))
logger.info('Save images: {}'.format(save_imgs))
logger.info('Flip test: {}'.format(flip_test))
#### set up the models
model.load_state_dict(torch.load(model_path), strict=True)
model.eval()
model = model.to(device)
model = nn.DataParallel(model)
#### set up the models
load_net = torch.load(color_model_path)
load_net_clean = OrderedDict() # add prefix 'color_net.'
for k, v in load_net.items():
k = 'color_net.' + k
load_net_clean[k] = v
color_model.load_state_dict(load_net_clean, strict=True)
color_model.eval()
color_model = color_model.to(device)
color_model = nn.DataParallel(color_model)
avg_psnr_l, avg_psnr_center_l, avg_psnr_border_l = [], [], []
subfolder_name_l = []
subfolder_l = sorted(glob.glob(osp.join(test_dataset_folder, '*')))
#print(subfolder_l)
#print(subfolder_GT_l)
#exit()
# for each subfolder
for subfolder in subfolder_l:
subfolder_name = osp.basename(subfolder)
subfolder_name_l.append(subfolder_name)
save_subfolder = osp.join(save_folder, subfolder_name)
img_path_l = sorted(glob.glob(osp.join(subfolder, '*')))
#print(img_path_l)
max_idx = len(img_path_l)
if save_imgs:
util.mkdirs(save_subfolder)
#### read LQ and GT images
imgs_LQ = data_util.read_img_seq(subfolder)
# process each image
for img_idx, img_path in enumerate(img_path_l):
img_name = osp.splitext(osp.basename(img_path))[0]
select_idx = data_util.index_generation(img_idx,
max_idx,
N_in,
padding=padding)
imgs_in = imgs_LQ.index_select(
0, torch.LongTensor(select_idx)).unsqueeze(0).cpu()
print(imgs_in.size())
if flip_test:
imgs_in = util.single_forward(color_model, imgs_in)
output = util.flipx4_forward(model, imgs_in)
else:
start_time = time.time()
imgs_in = util.single_forward(color_model, imgs_in)
output = util.single_forward(model, imgs_in)
end_time = time.time()
print('Forward One image:', end_time - start_time)
output = util.tensor2img(output.squeeze(0))
if save_imgs:
cv2.imwrite(
osp.join(save_subfolder, '{}.png'.format(img_name)),
output)
logger.info('{:3d} - {:25}'.format(img_idx + 1, img_name))
logger.info('################ Tidy Outputs ################')
logger.info('################ Final Results ################')
logger.info('Data: {} - {}'.format(data_mode, test_dataset_folder))
logger.info('Padding mode: {}'.format(padding))
logger.info('Model path: {}'.format(model_path))
logger.info('Save images: {}'.format(save_imgs))
logger.info('Flip test: {}'.format(flip_test))
def main():
#################
# configurations
#################
device = torch.device('cuda')
os.environ['CUDA_VISIBLE_DEVICES'] = '0'
data_mode = 'HDR' # Vid4 | sharp_bicubic | blur_bicubic | blur | blur_comp
# Vid4: SR
# REDS4: sharp_bicubic (SR-clean), blur_bicubic (SR-blur);
# blur (deblur-clean), blur_comp (deblur-compression).
stage = 1 # 1 or 2, use two stage strategy for REDS dataset.
flip_test = True
############################################################################
#### model
if data_mode == 'Vid4':
if stage == 1:
model_path = '../experiments/pretrained_models/EDVR_Vimeo90K_SR_L.pth'
else:
raise ValueError('Vid4 does not support stage 2.')
elif data_mode == 'sharp_bicubic':
if stage == 1:
model_path = '../experiments/pretrained_models/EDVR_REDS_SR_L.pth'
else:
model_path = '../experiments/pretrained_models/EDVR_REDS_SR_Stage2.pth'
elif data_mode == 'blur_bicubic':
if stage == 1:
model_path = '../experiments/pretrained_models/EDVR_REDS_SRblur_L.pth'
else:
model_path = '../experiments/pretrained_models/EDVR_REDS_SRblur_Stage2.pth'
elif data_mode == 'blur':
if stage == 1:
model_path = '../experiments/pretrained_models/EDVR_REDS_deblur_L.pth'
else:
model_path = '../experiments/pretrained_models/EDVR_REDS_deblur_Stage2.pth'
elif data_mode == 'blur_comp':
if stage == 1:
model_path = '../experiments/pretrained_models/EDVR_REDS_deblurcomp_L.pth'
else:
model_path = '../experiments/pretrained_models/EDVR_REDS_deblurcomp_Stage2.pth'
elif data_mode == 'HDR':
# model_path = '../experiments/001_EDVR_scratch_lr4e-4_600k_HDR_LrCAR4S/models/20000_G.pth'
# model_path = '../experiments/005_EDVRwoTSA_scratch_lr4e-4_600k_HDR_LrCAR4S/models/490000_G.pth'
model_path = '../experiments/pretrained_models/50000_G.pth'
else:
raise NotImplementedError
if data_mode == 'Vid4':
N_in = 7 # use N_in images to restore one HR image
else:
N_in = 5
predeblur, HR_in = False, False
back_RBs = 10
if data_mode == 'blur_bicubic':
predeblur = True
if data_mode == 'blur' or data_mode == 'blur_comp':
predeblur, HR_in = True, True
if data_mode == 'HDR':
predeblur = True
if stage == 2:
HR_in = True
back_RBs = 20
model = EDVR_arch.EDVR(64,
N_in,
8,
5,
back_RBs,
predeblur=predeblur,
HR_in=HR_in,
w_TSA=True)
#### dataset
GT_dataset_folder = None
if data_mode == 'Vid4':
test_dataset_folder = '../datasets/Vid4/BIx4'
GT_dataset_folder = '../datasets/Vid4/GT'
elif data_mode == 'HDR':
test_dataset_folder = '../datasets/HDR/valid/new_method/540p'
GT_dataset_folder = '../datasets/HDR/valid/new_method/4k'
# test_dataset_folder = '../datasets/HDR/valid/sequences_540'
# GT_dataset_folder = '../datasets/HDR/valid/sequences_4k'
else:
if stage == 1:
test_dataset_folder = '../datasets/REDS4/{}'.format(data_mode)
else:
test_dataset_folder = '../results/REDS-EDVR_REDS_SR_L_flipx4'
print('You should modify the test_dataset_folder path for stage 2')
GT_dataset_folder = '../datasets/REDS4/GT'
#### evaluation
crop_border = 0
border_frame = N_in // 2 # border frames when evaluate
# temporal padding mode
if data_mode == 'Vid4' or data_mode == 'sharp_bicubic':
padding = 'new_info'
else:
padding = 'replicate'
save_imgs = True
save_folder = '../results/{}_50000'.format(data_mode)
util.mkdirs(save_folder)
util.setup_logger('base',
save_folder,
'test',
level=logging.INFO,
screen=True,
tofile=True)
logger = logging.getLogger('base')
#### log info
logger.info('Data: {} - {}'.format(data_mode, test_dataset_folder))
logger.info('Padding mode: {}'.format(padding))
logger.info('Model path: {}'.format(model_path))
logger.info('Save images: {}'.format(save_imgs))
logger.info('Flip test: {}'.format(flip_test))
#### set up the models
model.load_state_dict(torch.load(model_path), strict=True)
model.eval()
model = model.to(device)
avg_psnr_l, avg_psnr_center_l, avg_psnr_border_l = [], [], []
subfolder_name_l = []
subfolder_l = sorted(glob.glob(osp.join(test_dataset_folder, '*')))
subfolder_GT_l = sorted(glob.glob(osp.join(GT_dataset_folder, '*')))
# for each subfolder
for subfolder, subfolder_GT in zip(subfolder_l, subfolder_GT_l):
# print(subfolder, subfolde
if '10675978' not in subfolder:
print('pass')
continue
subfolder_name = osp.basename(subfolder)
subfolder_name_l.append(subfolder_name)
save_subfolder = osp.join(save_folder, subfolder_name)
img_path_l = sorted(glob.glob(osp.join(subfolder, '*')))
max_idx = len(img_path_l)
if save_imgs:
util.mkdirs(save_subfolder)
#### read LQ and GT images
imgs_LQ = data_util.read_img_seq(subfolder)
img_GT_l = []
for img_GT_path in sorted(glob.glob(osp.join(subfolder_GT, '*'))):
img_GT_l.append(data_util.read_img(None, img_GT_path))
avg_psnr, avg_psnr_border, avg_psnr_center, N_border, N_center = 0, 0, 0, 0, 0
# process each image
for img_idx, img_path in enumerate(img_path_l):
img_name = osp.splitext(osp.basename(img_path))[0]
select_idx = data_util.index_generation(img_idx,
max_idx,
N_in,
padding=padding)
imgs_in = imgs_LQ.index_select(
0, torch.LongTensor(select_idx)).unsqueeze(0).to(device)
if flip_test:
output = util.flipx4_forward(model, imgs_in)
else:
output = util.single_forward(model, imgs_in)
output = util.tensor2img(output.squeeze(0))
if save_imgs:
cv2.imwrite(
osp.join(save_subfolder, '{}.png'.format(img_name)),
output)
# calculate PSNR
# output = output / 255.
# GT = np.copy(img_GT_l[img_idx])
# # For REDS, evaluate on RGB channels; for Vid4, evaluate on the Y channel
# if data_mode == 'Vid4': # bgr2y, [0, 1]
# GT = data_util.bgr2ycbcr(GT, only_y=True)
# output = data_util.bgr2ycbcr(output, only_y=True)
# output, GT = util.crop_border([output, GT], crop_border)
# crt_psnr = util.calculate_psnr(output * 255, GT * 255)
# # crt_ssim = util.calculate_ssim(output * 255, GT * 255)
logger.info('{:3d} - {:25} '.format(img_idx + 1, img_name))
def main(gpu_id, start_id, step):
#################
# configurations
#################
device = torch.device('cuda')
os.environ['CUDA_VISIBLE_DEVICES'] = gpu_id
data_mode = 'AI4K'
stage = 1 # 1 or 2
flip_test = True
#### model
if data_mode == 'AI4K':
if stage == 1:
model_path = '/home/zenghui/projects/4KHDR/experiments/pretrained_models/EDVR_L_G300k.pth' # TODO: change path
else:
model_path = '../experiments/pretrained_models/EDVR_REDS_SR_Stage2.pth' # TODO: change path
else:
raise NotImplementedError()
N_in = 5 # use N_in images to restore one HR image
predeblur, HR_in = False, False
back_RBs = 40
if stage == 2:
HR_in = True
back_RBs = 20
model = EDVR_arch.EDVR(64,
N_in,
8,
5,
back_RBs,
predeblur=predeblur,
HR_in=HR_in)
#### dataset
if data_mode == 'AI4K':
test_dataset_folder = '../datasets/SDR_540p_PNG_test' # TODO: change path
else:
raise NotImplementedError()
#### scene information
scene_index_path = '../keys/test_scene_idx.pkl' # TODO: change path
scene_dict = pickle.load(open(scene_index_path, 'rb'))
#### evaluation
crop_border = 0
border_frame = N_in // 2 # border frames when evaluate
padding = 'replicate' # temporal padding mode
save_imgs = True
save_folder = '../results_edvr_l_tsa_300k_2/{}'.format(
data_mode) # TODO: change path
util.mkdirs(save_folder)
util.setup_logger('base',
save_folder,
'test',
level=logging.INFO,
screen=True,
tofile=True)
logger = logging.getLogger('base')
#### log info
logger.info('Data: {} - {}'.format(data_mode, test_dataset_folder))
logger.info('Padding mode: {}'.format(padding))
logger.info('Model path: {}'.format(model_path))
logger.info('Save images: {}'.format(save_imgs))
logger.info('Flip test: {}'.format(flip_test))
#### set up the models
model.load_state_dict(torch.load(model_path), strict=True)
model.eval()
model = model.to(device)
subfolder_name_l = []
subfolder_l = sorted(glob.glob(osp.join(test_dataset_folder, '*')))
seq_id = start_id
for subfolder in subfolder_l[start_id::step]:
subfolder_name = osp.basename(subfolder)
subfolder_name_l.append(subfolder_name)
save_subfolder = osp.join(save_folder, subfolder_name)
logger.info(
'Processing sequence: {}, seq_id = {}, crop_edge = {}'.format(
subfolder_name, seq_id, crop_edge[seq_id]))
hr_crop_edge = crop_edge[seq_id] * 4
img_path_l = sorted(glob.glob(osp.join(subfolder, '*')))
max_idx = len(img_path_l)
if save_imgs:
util.mkdirs(save_subfolder)
#### read LQ images
imgs_LQ = data_util.read_img_seq(subfolder)
# process each image
for img_idx, img_path in enumerate(img_path_l):
img_name = osp.splitext(osp.basename(img_path))[0]
select_idx = data_util.index_generation_with_scene_list(
img_idx,
max_idx,
N_in,
scene_dict[subfolder_name],
padding=padding)
imgs_in = imgs_LQ.index_select(
0, torch.LongTensor(select_idx)).unsqueeze(0).to(device)
if flip_test:
output = util.flipx4_forward(model, imgs_in)
else:
output = util.single_forward(model, imgs_in)
#if crop_edge[seq_id]>0:
# output[:,:,:hr_crop_edge, :] = 0
# output[:,:,-hr_crop_edge:, :] = 0
output = util.tensor2img(output.squeeze(0))
if save_imgs:
cv2.imwrite(
osp.join(save_subfolder, '{}.png'.format(img_name)),
output)
seq_id += step
def __init__(self, opt, **kwargs):
super(VideoTestDataset, self).__init__()
self.scale = kwargs['scale']
self.kernel_size = kwargs['kernel_size']
self.model_name = kwargs['model_name']
idx = kwargs['idx'] if 'idx' in kwargs else None
self.opt = opt
self.cache_data = opt['cache_data']
self.half_N_frames = opt['N_frames'] // 2
if idx is None:
self.name = opt['name']
self.GT_root, self.LQ_root = opt['dataroot_GT'], opt['dataroot_LQ']
degradation_type = opt['degradation_type']
opt_sigma_x = opt['sigma_x']
opt_sigma_y = opt['sigma_y']
opt_theta = opt['theta']
else:
self.name = opt['name'].split('+')[idx]
self.GT_root, self.LQ_root = opt['dataroot_GT'].split(
'+')[idx], opt['dataroot_LQ'].split('+')[idx]
if '+' in opt['degradation_type']:
degradation_type = opt['degradation_type'].split('+')[idx]
if '+' in str(opt['sigma_x']):
opt_sigma_x = float(opt['sigma_x'].split('+')[idx])
opt_sigma_y = float(opt['sigma_y'].split('+')[idx])
opt_theta = float(opt['theta'].split('+')[idx])
else:
opt_sigma_x = opt['sigma_x']
opt_sigma_y = opt['sigma_y']
opt_theta = opt['theta']
else:
degradation_type = opt['degradation_type']
opt_sigma_x = opt['sigma_x']
opt_sigma_y = opt['sigma_y']
opt_theta = opt['theta']
self.data_type = self.opt['data_type']
self.data_info = {
'path_SLQ': [],
'path_LQ': [],
'path_GT': [],
'folder': [],
'idx': [],
'border': []
}
if self.data_type == 'lmdb':
raise ValueError('No need to use LMDB during validation/test.')
#### Generate data info and cache data
self.imgs_SLQ, self.imgs_LQ, self.imgs_GT = {}, {}, {}
if opt['degradation_mode'] == 'preset':
self.LQ_root = self.LQ_root + '_preset'
else:
if isinstance(opt_sigma_x, list):
assert len(opt_sigma_x) == len(opt_sigma_y)
assert len(opt_sigma_x) == len(opt_theta)
LQ_root_list = []
for i, (sigma_x, sigma_y, theta) in enumerate(
zip(opt_sigma_x, opt_sigma_y, opt_theta)):
LQ_root_list.append(self.LQ_root + '_' + degradation_type + '_' + str('{:.1f}'.format(opt_sigma_x[i]))\
+ '_' + str('{:.1f}'.format(opt_sigma_y[i])) + '_' + str('{:.1f}'.format(opt_theta[i])))
self.LQ_root = LQ_root_list
else:
self.LQ_root = self.LQ_root + '_' + degradation_type + '_' + str('{:.1f}'.format(opt_sigma_x))\
+ '_' + str('{:.1f}'.format(opt_sigma_y)) + '_' + str('{:.1f}'.format(opt_theta))
slr_name = '' if opt['slr_mode'] is None else '_{}'.format(
opt['slr_mode'])
print(self.LQ_root)
if self.name.lower() in ['vid4', 'reds', 'mm522']:
if self.name.lower() == 'vid4':
img_type = 'img'
subfolders_GT = util.glob_file_list(self.GT_root)
if isinstance(self.LQ_root, list):
num_settings = len(self.LQ_root)
subfolders_LQ_list = [
util.glob_file_list(
osp.join(LQ_root, 'X{}'.format(self.scale)))
for LQ_root in self.LQ_root
]
subfolders_SLQ_list = [
util.glob_file_list(
osp.join(
LQ_root,
'X{}{}'.format(self.scale * self.scale,
slr_name)))
for LQ_root in self.LQ_root
]
subfolders_LQ = []
subfolders_SLQ = []
for i in range(len(subfolders_LQ_list[0])):
subfolders_LQ.append([
subfolders_LQ_list[j][i]
for j in range(len(subfolders_LQ_list))
])
subfolders_SLQ.append([
subfolders_SLQ_list[j][i]
for j in range(len(subfolders_SLQ_list))
])
else:
subfolders_LQ = util.glob_file_list(
osp.join(self.LQ_root, 'X{}'.format(self.scale)))
subfolders_SLQ = util.glob_file_list(
osp.join(
self.LQ_root,
'X{}{}'.format(self.scale * self.scale, slr_name)))
elif self.name.lower() == 'reds':
img_type = 'img'
list_hr_seq = util.glob_file_list(self.GT_root)
subfolders_GT = [
k for k in list_hr_seq
if k.find('000') >= 0 or k.find('011') >= 0
or k.find('015') >= 0 or k.find('020') >= 0
]
if isinstance(self.LQ_root, list):
num_settings = len(self.LQ_root)
subfolders_LQ_list = []
subfolders_SLQ_list = []
for i in range(num_settings):
list_lr_seq = util.glob_file_list(
osp.join(self.LQ_root[i],
'X{}'.format(self.scale)))
list_slr_seq = util.glob_file_list(
osp.join(
self.LQ_root[i],
'X{}{}'.format(self.scale * self.scale,
slr_name)))
subfolder_LQ = [
k for k in list_lr_seq
if k.find('000') >= 0 or k.find('011') >= 0
or k.find('015') >= 0 or k.find('020') >= 0
]
subfolder_SLQ = [
k for k in list_slr_seq
if k.find('000') >= 0 or k.find('011') >= 0
or k.find('015') >= 0 or k.find('020') >= 0
]
subfolders_LQ_list.append(subfolder_LQ)
subfolders_SLQ_list.append(subfolder_SLQ)
subfolders_LQ = []
subfolders_SLQ = []
for i in range(len(subfolders_LQ_list[0])):
subfolders_LQ.append([
subfolders_LQ_list[j][i]
for j in range(len(subfolders_LQ_list))
])
subfolders_SLQ.append([
subfolders_SLQ_list[j][i]
for j in range(len(subfolders_SLQ_list))
])
else:
list_lr_seq = util.glob_file_list(
osp.join(self.LQ_root, 'X{}'.format(self.scale)))
list_slr_seq = util.glob_file_list(
osp.join(
self.LQ_root,
'X{}{}'.format(self.scale * self.scale, slr_name)))
#subfolders_GT = [k for k in list_hr_seq if
# k.find('000') >= 0 or k.find('011') >= 0 or k.find('015') >= 0 or k.find('020') >= 0]
subfolders_LQ = [
k for k in list_lr_seq
if k.find('000') >= 0 or k.find('011') >= 0
or k.find('015') >= 0 or k.find('020') >= 0
]
subfolders_SLQ = [
k for k in list_slr_seq
if k.find('000') >= 0 or k.find('011') >= 0
or k.find('015') >= 0 or k.find('020') >= 0
]
else:
img_type = 'img'
list_hr_seq = util.glob_file_list(self.GT_root)
list_lr_seq = util.glob_file_list(
osp.join(self.LQ_root, 'X{}'.format(self.scale)))
list_slr_seq = util.glob_file_list(
osp.join(self.LQ_root,
'X{}'.format(self.scale * self.scale)))
subfolders_GT = [
k for k in list_hr_seq
if k.find('001') >= 0 or k.find('005') >= 0
or k.find('008') >= 0 or k.find('009') >= 0
]
subfolders_LQ = [
k for k in list_lr_seq
if k.find('001') >= 0 or k.find('005') >= 0
or k.find('008') >= 0 or k.find('009') >= 0
]
subfolders_SLQ = [
k for k in list_slr_seq
if k.find('001') >= 0 or k.find('005') >= 0
or k.find('008') >= 0 or k.find('009') >= 0
]
print(subfolders_GT[0], '\n', subfolders_LQ[0], '\n',
subfolders_SLQ[0])
for subfolder_SLQ, subfolder_LQ, subfolder_GT in zip(
subfolders_SLQ, subfolders_LQ, subfolders_GT):
subfolder_name = osp.basename(subfolder_GT)
img_paths_GT = util.glob_file_list(subfolder_GT)
if isinstance(subfolder_LQ, list):
img_paths_LQ_list = [
util.glob_file_list(subf_LQ)
for subf_LQ in subfolder_LQ
]
img_paths_SLQ_list = [
util.glob_file_list(subf_SLQ)
for subf_SLQ in subfolder_SLQ
]
img_paths_LQ = []
img_paths_SLQ = []
for i in range(len(img_paths_GT)):
img_paths_LQ.append(img_paths_LQ_list[i %
num_settings][i])
img_paths_SLQ.append(
img_paths_SLQ_list[i % num_settings][i])
else:
img_paths_LQ = util.glob_file_list(subfolder_LQ)
img_paths_SLQ = util.glob_file_list(subfolder_SLQ)
max_idx = len(img_paths_GT)
self.data_info['path_SLQ'].extend(img_paths_SLQ)
self.data_info['path_LQ'].extend(img_paths_LQ)
self.data_info['path_GT'].extend(img_paths_GT)
self.data_info['folder'].extend([subfolder_name] * max_idx)
for i in range(max_idx):
self.data_info['idx'].append('{}/{}'.format(i, max_idx))
border_l = [0] * max_idx
for i in range(self.half_N_frames):
border_l[i] = 1
border_l[max_idx - i - 1] = 1
self.data_info['border'].extend(border_l)
self.imgs_GT[subfolder_name] = util.read_img_seq(
img_paths_GT, img_type)
if opt['degradation_mode'] == 'preset':
self.imgs_LQ[subfolder_name] = torch.stack([
util.read_img_seq(util.glob_file_list(paths_LQ),
img_type)
for paths_LQ in img_paths_LQ
],
dim=0)
self.imgs_SLQ[subfolder_name] = torch.stack([
util.read_img_seq(util.glob_file_list(paths_SLQ),
img_type)
for paths_SLQ in img_paths_SLQ
],
dim=0)
else:
self.imgs_LQ[subfolder_name] = util.read_img_seq(
img_paths_LQ, img_type)
self.imgs_SLQ[subfolder_name] = util.read_img_seq(
img_paths_SLQ, img_type)
h, w = self.imgs_SLQ[subfolder_name].shape[-2:]
if h % 4 != 0 or w % 4 != 0:
self.imgs_SLQ[subfolder_name] = self.imgs_SLQ[
subfolder_name][..., :h - (h % 4), :w - (w % 4)]
self.imgs_LQ[subfolder_name] = self.imgs_LQ[
subfolder_name][..., :self.scale *
(h - (h % 4)), :self.scale * (w -
(w % 4))]
self.imgs_GT[subfolder_name] = self.imgs_GT[
subfolder_name][..., :self.scale * self.scale *
(h - (h % 4)), :self.scale *
self.scale * (w - (w % 4))]
else:
raise ValueError(
'Not support video test dataset. Support Vid4, REDS4 and Vimeo90k-Test.'
)
'''
if opt['degradation_mode'] == 'set':
sigma_x = float(opt['sigma_x'])
sigma_y = float(opt['sigma_y'])
theta = float(opt['theta'])
gen_kwargs = preprocessing.set_kernel_params(sigma_x=sigma_x, sigma_y=sigma_y, theta=theta)
self.kernel_gen = rkg.Degradation(self.kernel_size, self.scale, **gen_kwargs)
self.gen_kwargs_l = [gen_kwargs['sigma'][0], gen_kwargs['sigma'][1], gen_kwargs['theta']]
'''
if opt['degradation_mode'] == 'preset':
self.kernel_gen = rkg.Degradation(self.kernel_size, self.scale)
if self.name.lower() == 'vid4':
self.kernel_dict = np.load(
'../pretrained_models/Mixed/Vid4.npy')
elif self.name.lower() == 'reds':
self.kernel_dict = np.load(
'../pretrained_models/Mixed/REDS.npy')
else:
raise NotImplementedError()
def read_batch(self, batch_num : int):
"""load 100 imgs from cache"""
self.imgs_paths = sorted(glob.glob(os.path.join(f'{self.imgs_cache}{batch_num}/', '*.png')))
self.imgs = data_util.read_img_seq(self.imgs_paths)
def main():
#################
# configurations
#################
device = torch.device('cuda')
os.environ['CUDA_VISIBLE_DEVICES'] = '4'
data_mode = 'youku'
split_test = True
if split_test:
forward_fn = PatchedForward(3, 8, 4, 4)
else:
forward_fn = PaddingForward(4, 4)
############################################################################
#### model
model_path = '../experiments/pretrained_models/EDVR_YOUKU_M_woTSA.pth'
N_in = 5
predeblur, HR_in = False, False
back_RBs = 10
model = EDVR_arch.EDVR(64,
N_in,
8,
5,
10,
predeblur=predeblur,
HR_in=HR_in,
w_TSA=False,
block_type='rcab',
non_local=False)
ssim_fn = SSIM()
#### dataset
test_dataset_folder = '/media/tclwh2/public/youku/val/lq'
GT_dataset_folder = '/media/tclwh2/public/youku/val/gt'
#### evaluation
border_frame = N_in // 2 # border frames when evaluate
# temporal padding mode
padding = 'new_info'
save_imgs = False
save_folder = '../results/{}'.format(data_mode)
util.mkdirs(save_folder)
util.setup_logger('base',
save_folder,
'test',
level=logging.INFO,
screen=True,
tofile=True)
logger = logging.getLogger('base')
#### log info
logger.info('Data: {} - {}'.format(data_mode, test_dataset_folder))
logger.info('Padding mode: {}'.format(padding))
logger.info('Model path: {}'.format(model_path))
logger.info('Save images: {}'.format(save_imgs))
logger.info('Split test: {}'.format(split_test))
#### set up the models
model.load_state_dict(torch.load(model_path), strict=True)
model.eval()
model = model.to(device)
ssim_fn = ssim_fn.to(device)
ssim_fn.eval()
avg_psnr_l, avg_ssim_l = [], []
subfolder_name_l = []
subfolder_l = sorted(glob.glob(osp.join(test_dataset_folder, '*')))
subfolder_GT_l = sorted(glob.glob(osp.join(GT_dataset_folder, '*')))
# for each subfolder
for subfolder, subfolder_GT in zip(subfolder_l, subfolder_GT_l):
subfolder_name = osp.basename(subfolder)
subfolder_name_l.append(subfolder_name)
save_subfolder = osp.join(save_folder, subfolder_name)
img_path_l = sorted(glob.glob(osp.join(subfolder, '*')))
max_idx = len(img_path_l)
if save_imgs:
util.mkdirs(save_subfolder)
#### read LQ and GT images
imgs_LQ = data_util.read_img_seq(subfolder)
img_GT_l = []
for img_GT_path in sorted(glob.glob(osp.join(subfolder_GT, '*'))):
img_GT_l.append(data_util.read_img(None, img_GT_path))
avg_psnr = 0
avg_ssim = 0
# process each image
for img_idx, img_path in enumerate(img_path_l):
img_name = osp.splitext(osp.basename(img_path))[0]
select_idx = data_util.index_generation(img_idx,
max_idx,
N_in,
padding=padding)
imgs_in = imgs_LQ.index_select(
0, torch.LongTensor(select_idx)).unsqueeze(0).to(device)
with torch.no_grad():
model_output = forward_fn(model, imgs_in)
GT = np.copy(img_GT_l[img_idx][:, :, [2, 1, 0]])
GT = torch.from_numpy(GT.transpose(2, 0,
1)).unsqueeze_(0).to(device)
crt_ssim = ssim_fn(model_output, GT).data.cpu().item()
output = model_output.data.float().cpu()
output = util.tensor2img(output.squeeze(0))
if save_imgs:
cv2.imwrite(
osp.join(save_subfolder, '{}.png'.format(img_name)),
output)
# calculate PSNR
output = output / 255.
GT = np.copy(img_GT_l[img_idx])
crt_psnr = util.calculate_psnr(output * 255, GT * 255)
logger.info('{:3d} - {:25} \tPSNR/SSIM: {:.6f}/{:.6f}'.format(
img_idx + 1, img_name, crt_psnr, crt_ssim))
avg_psnr += crt_psnr
avg_ssim += crt_ssim
avg_psnr /= max_idx
avg_ssim /= max_idx
avg_psnr_l.append(avg_psnr)
avg_ssim_l.append(avg_ssim)
logger.info(
'Folder {} - Average PSNR/SSIM: {:.6f}/{:.6f} for {} frames. '.
format(subfolder_name, avg_psnr, avg_ssim, max_idx))
logger.info('################ Tidy Outputs ################')
for subfolder_name, psnr, ssim in zip(subfolder_name_l, avg_psnr_l,
avg_ssim_l):
logger.info('Folder {} - Average PSNR/SSIM: {:.6f}/{:.6f}. '.format(
subfolder_name, psnr, ssim))
logger.info('################ Final Results ################')
logger.info('Data: {} - {}'.format(data_mode, test_dataset_folder))
logger.info('Padding mode: {}'.format(padding))
logger.info('Model path: {}'.format(model_path))
logger.info('Save images: {}'.format(save_imgs))
logger.info('Total Average PSNR/SSIM: {:.6f}/{:.6f} for {} clips. '.format(
sum(avg_psnr_l) / len(avg_psnr_l),
sum(avg_ssim_l) / len(avg_ssim_l), len(subfolder_l)))
def main():
#################
# configurations
#################
device = torch.device('cuda')
os.environ['CUDA_VISIBLE_DEVICES'] = '0'
data_mode = 'Vid4' # Vid4 | sharp_bicubic | blur_bicubic | blur | blur_comp
# Vid4: SR
# REDS4: sharp_bicubic (SR-clean), blur_bicubic (SR-blur);
# blur (deblur-clean), blur_comp (deblur-compression).
stage = 1 # 1 or 2, use two stage strategy for REDS dataset.
flip_test = False
############################################################################
#### model
if data_mode == 'Vid4':
if stage == 1:
model_path = '../experiments/pretrained_models/cinepak_small2.pth'
else:
raise ValueError('Vid4 does not support stage 2.')
elif data_mode == 'sharp_bicubic':
if stage == 1:
model_path = '../experiments/pretrained_models/EDVR_REDS_SR_L.pth'
else:
model_path = '../experiments/pretrained_models/EDVR_REDS_SR_Stage2.pth'
elif data_mode == 'blur_bicubic':
if stage == 1:
model_path = '../experiments/pretrained_models/EDVR_REDS_SRblur_L.pth'
else:
model_path = '../experiments/pretrained_models/EDVR_REDS_SRblur_Stage2.pth'
elif data_mode == 'blur':
if stage == 1:
model_path = '../experiments/pretrained_models/EDVR_REDS_deblur_L.pth'
else:
model_path = '../experiments/pretrained_models/EDVR_REDS_deblur_Stage2.pth'
elif data_mode == 'blur_comp':
if stage == 1:
model_path = '../experiments/pretrained_models/EDVR_REDS_deblurcomp_L.pth'
else:
model_path = '../experiments/pretrained_models/EDVR_REDS_deblurcomp_Stage2.pth'
else:
raise NotImplementedError
if data_mode == 'Vid4':
N_in = 7 # use N_in images to restore one HR image
else:
N_in = 5
predeblur, HR_in = False, False
back_RBs = 10
if data_mode == 'blur_bicubic':
predeblur = True
if data_mode == 'blur' or data_mode == 'blur_comp':
predeblur, HR_in = True, True
if stage == 2:
HR_in = True
back_RBs = 20
model = EDVR_arch.EDVR(64,
N_in,
8,
5,
back_RBs,
predeblur=predeblur,
HR_in=HR_in)
#### dataset
if data_mode == 'Vid4':
test_dataset_folder = '../datasets/Vid4/BIx4'
GT_dataset_folder = '../datasets/Vid4/GT'
else:
if stage == 1:
test_dataset_folder = '../datasets/REDS4/{}'.format(data_mode)
else:
test_dataset_folder = '../results/REDS-EDVR_REDS_SR_L_flipx4'
print('You should modify the test_dataset_folder path for stage 2')
GT_dataset_folder = '../datasets/REDS4/GT'
#### evaluation
crop_border = 0
border_frame = N_in // 2 # border frames when evaluate
# temporal padding mode
if data_mode == 'Vid4' or data_mode == 'sharp_bicubic':
padding = 'new_info'
else:
padding = 'replicate'
save_imgs = True
save_folder = '../results/{}'.format(data_mode)
util.mkdirs(save_folder)
util.setup_logger('base',
save_folder,
'test',
level=logging.INFO,
screen=True,
tofile=True)
logger = logging.getLogger('base')
#### log info
logger.info('Data: {} - {}'.format(data_mode, test_dataset_folder))
logger.info('Padding mode: {}'.format(padding))
logger.info('Model path: {}'.format(model_path))
logger.info('Save images: {}'.format(save_imgs))
logger.info('Flip test: {}'.format(flip_test))
#### set up the models
model.load_state_dict(torch.load(model_path), strict=True)
model.eval()
model = model.to(device)
avg_psnr_l, avg_psnr_center_l, avg_psnr_border_l = [], [], []
subfolder_name_l = []
subfolder_l = sorted(glob.glob(osp.join(test_dataset_folder, '*')))
# for each subfolder
for subfolder in subfolder_l:
print('Processing video {:s}'.format(subfolder))
subfolder_name = osp.basename(subfolder)
subfolder_name_l.append(subfolder_name)
save_subfolder = osp.join(save_folder, subfolder_name)
img_path_l = sorted(glob.glob(osp.join(subfolder, '*')))
max_idx = len(img_path_l)
if save_imgs:
util.mkdirs(save_subfolder)
#### read LQ and GT images
imgs_LQ = data_util.read_img_seq(subfolder)
# process each image
for img_idx, img_path in enumerate(img_path_l):
print('\tProcessing frame {:s}')
img_name = osp.splitext(osp.basename(img_path))[0]
select_idx = data_util.index_generation(img_idx,
max_idx,
N_in,
padding=padding)
imgs_in = imgs_LQ.index_select(
0, torch.LongTensor(select_idx)).unsqueeze(0).to(device)
output = util.single_forward(model, imgs_in)
output = util.tensor2img(output.squeeze(0))
cv2.imwrite(osp.join(save_subfolder, '{}.png'.format(img_name)),
output)
def main():
# Create object for parsing command-line options
parser = argparse.ArgumentParser(description="Test with EDVR, requre path to test dataset folder.")
# Add argument which takes path to a bag file as an input
parser.add_argument("-i", "--input", type=str, help="Path to test folder")
# Parse the command line arguments to an object
args = parser.parse_args()
# Safety if no parameter have been given
if not args.input:
print("No input paramater have been given.")
print("For help type --help")
exit()
folder_name = args.input.split("/")[-1]
if folder_name == '':
index = len(args.input.split("/")) - 2
folder_name = args.input.split("/")[index]
#################
# configurations
#################
device = torch.device('cuda')
os.environ['CUDA_VISIBLE_DEVICES'] = '0'
data_mode = 'Vid4' # Vid4 | sharp_bicubic | blur_bicubic | blur | blur_comp
# Vid4: SR
# REDS4: sharp_bicubic (SR-clean), blur_bicubic (SR-blur);
# blur (deblur-clean), blur_comp (deblur-compression).
stage = 1 # 1 or 2, use two stage strategy for REDS dataset.
flip_test = False
############################################################################
#### model
if data_mode == 'Vid4':
if stage == 1:
model_path = '../experiments/pretrained_models/EDVR_Vimeo90K_SR_L.pth'
else:
raise ValueError('Vid4 does not support stage 2.')
else:
raise NotImplementedError
if data_mode == 'Vid4':
N_in = 7 # use N_in images to restore one HR image
else:
N_in = 5
predeblur, HR_in = False, False
back_RBs = 40
model = EDVR_arch.EDVR(128, N_in, 8, 5, back_RBs, predeblur=predeblur, HR_in=HR_in)
#### dataset
if data_mode == 'Vid4':
# debug
test_dataset_folder = os.path.join(args.input, 'BIx4')
GT_dataset_folder = os.path.join(args.input, 'GT')
else:
if stage == 1:
test_dataset_folder = '../datasets/REDS4/{}'.format(data_mode)
else:
test_dataset_folder = '../results/REDS-EDVR_REDS_SR_L_flipx4'
print('You should modify the test_dataset_folder path for stage 2')
GT_dataset_folder = '../datasets/REDS4/GT'
#### evaluation
crop_border = 0
border_frame = N_in // 2 # border frames when evaluate
# temporal padding mode
if data_mode == 'Vid4' or data_mode == 'sharp_bicubic':
padding = 'new_info'
else:
padding = 'replicate'
save_imgs = True
save_folder = '../results/{}'.format(folder_name)
util.mkdirs(save_folder)
util.setup_logger('base', save_folder, 'test', level=logging.INFO, screen=True, tofile=True)
logger = logging.getLogger('base')
#### log info
logger.info('Data: {} - {}'.format(folder_name, test_dataset_folder))
logger.info('Padding mode: {}'.format(padding))
logger.info('Model path: {}'.format(model_path))
logger.info('Save images: {}'.format(save_imgs))
logger.info('Flip test: {}'.format(flip_test))
#### set up the models
model.load_state_dict(torch.load(model_path), strict=True)
model.eval()
model = model.to(device)
avg_psnr_l, avg_psnr_center_l, avg_psnr_border_l = [], [], []
subfolder_name_l = []
subfolder_l = sorted(glob.glob(osp.join(test_dataset_folder, '*')))
subfolder_GT_l = sorted(glob.glob(osp.join(GT_dataset_folder, '*')))
# for each subfolder
for subfolder, subfolder_GT in zip(subfolder_l, subfolder_GT_l):
subfolder_name = osp.basename(subfolder)
subfolder_name_l.append(subfolder_name)
save_subfolder = osp.join(save_folder, subfolder_name)
img_path_l = sorted(glob.glob(osp.join(subfolder, '*')))
max_idx = len(img_path_l)
if save_imgs:
util.mkdirs(save_subfolder)
#### read LQ and GT images
imgs_LQ = data_util.read_img_seq(subfolder)
img_GT_l = []
for img_GT_path in sorted(glob.glob(osp.join(subfolder_GT, '*'))):
img_GT_l.append(data_util.read_img(None, img_GT_path))
avg_psnr, avg_psnr_border, avg_psnr_center, N_border, N_center = 0, 0, 0, 0, 0
# process each image
for img_idx, img_path in enumerate(img_path_l):
img_name = osp.splitext(osp.basename(img_path))[0]
select_idx = data_util.index_generation(img_idx, max_idx, N_in, padding=padding)
imgs_in = imgs_LQ.index_select(0, torch.LongTensor(select_idx)).unsqueeze(0).to(device)
if flip_test:
output = util.flipx4_forward(model, imgs_in)
else:
output = util.single_forward(model, imgs_in)
output = util.tensor2img(output.squeeze(0))
if save_imgs:
cv2.imwrite(osp.join(save_subfolder, '{}.png'.format(img_name)), output)
# calculate PSNR
output = output / 255.
GT = np.copy(img_GT_l[img_idx])
# For REDS, evaluate on RGB channels; for Vid4, evaluate on the Y channel
if data_mode == 'Vid4': # bgr2y, [0, 1]
GT = data_util.bgr2ycbcr(GT, only_y=True)
output = data_util.bgr2ycbcr(output, only_y=True)
output, GT = util.crop_border([output, GT], crop_border)
crt_psnr = util.calculate_psnr(output * 255, GT * 255)
logger.info('{:3d} - {:25} \tPSNR: {:.6f} dB'.format(img_idx + 1, img_name, crt_psnr))
if img_idx >= border_frame and img_idx < max_idx - border_frame: # center frames
avg_psnr_center += crt_psnr
N_center += 1
else: # border frames
avg_psnr_border += crt_psnr
N_border += 1
avg_psnr = (avg_psnr_center + avg_psnr_border) / (N_center + N_border)
avg_psnr_center = avg_psnr_center / N_center
avg_psnr_border = 0 if N_border == 0 else avg_psnr_border / N_border
avg_psnr_l.append(avg_psnr)
avg_psnr_center_l.append(avg_psnr_center)
avg_psnr_border_l.append(avg_psnr_border)
logger.info('Folder {} - Average PSNR: {:.6f} dB for {} frames; '
'Center PSNR: {:.6f} dB for {} frames; '
'Border PSNR: {:.6f} dB for {} frames.'.format(subfolder_name, avg_psnr,
(N_center + N_border),
avg_psnr_center, N_center,
avg_psnr_border, N_border))
logger.info('################ Tidy Outputs ################')
for subfolder_name, psnr, psnr_center, psnr_border in zip(subfolder_name_l, avg_psnr_l,
avg_psnr_center_l, avg_psnr_border_l):
logger.info('Folder {} - Average PSNR: {:.6f} dB. '
'Center PSNR: {:.6f} dB. '
'Border PSNR: {:.6f} dB.'.format(subfolder_name, psnr, psnr_center,
psnr_border))
logger.info('################ Final Results ################')
logger.info('Data: {} - {}'.format(folder_name, test_dataset_folder))
logger.info('Padding mode: {}'.format(padding))
logger.info('Model path: {}'.format(model_path))
logger.info('Save images: {}'.format(save_imgs))
logger.info('Flip test: {}'.format(flip_test))
logger.info('Total Average PSNR: {:.6f} dB for {} clips. '
'Center PSNR: {:.6f} dB. Border PSNR: {:.6f} dB.'.format(
sum(avg_psnr_l) / len(avg_psnr_l), len(subfolder_l),
sum(avg_psnr_center_l) / len(avg_psnr_center_l),
sum(avg_psnr_border_l) / len(avg_psnr_border_l)))
def main():
#################
# configurations
#################
device = torch.device('cuda')
os.environ['CUDA_VISIBLE_DEVICES'] = '1'
data_mode = 'ai4khdr_test'
flip_test = False
############################################################################
#### model
#################
if data_mode == 'ai4khdr_test':
model_path = '../experiments/002_EDVR_lr4e-4_600k_AI4KHDR/models/4000_G.pth'
else:
raise NotImplementedError
N_in = 5
front_RBs = 5
back_RBs = 10
predeblur, HR_in = False, False
model = EDVR_arch.EDVR(64, N_in, 8, front_RBs, back_RBs, predeblur=predeblur, HR_in=HR_in)
############################################################################
#### dataset
#################
if data_mode == 'ai4khdr_test':
test_dataset_folder = '/workspace/nas_mengdongwei/dataset/AI4KHDR/test/540p_frames'
else:
raise NotImplementedError
############################################################################
#### evaluation
crop_border = 0
border_frame = N_in // 2 # border frames when evaluate
# temporal padding mode
if data_mode == 'ai4khdr_test':
padding = 'new_info'
else:
padding = 'replicate'
save_imgs = True
save_folder = '../results/{}_{}'.format(data_mode, util.get_timestamp())
util.mkdirs(save_folder)
util.setup_logger('base', save_folder, 'test', level=logging.INFO, screen=True, tofile=True)
logger = logging.getLogger('base')
#### log info
logger.info('Data: {} - {}'.format(data_mode, test_dataset_folder))
logger.info('Padding mode: {}'.format(padding))
logger.info('Model path: {}'.format(model_path))
logger.info('Save images: {}'.format(save_imgs))
logger.info('Flip test: {}'.format(flip_test))
#### set up the models
model.load_state_dict(torch.load(model_path), strict=False)
model.eval()
model = model.to(device)
subfolder_name_l = []
subfolder_l = sorted(glob.glob(osp.join(test_dataset_folder, '*')))
# for each subfolder
for subfolder in subfolder_l:
subfolder_name = osp.basename(subfolder)
subfolder_name_l.append(subfolder_name)
save_subfolder = osp.join(save_folder, subfolder_name)
img_path_l = sorted(glob.glob(osp.join(subfolder, '*')))
max_idx = len(img_path_l)
if save_imgs:
util.mkdirs(save_subfolder)
#### read LQ and GT images
imgs_LQ = data_util.read_img_seq(subfolder)
# process each image
for img_idx, img_path in enumerate(img_path_l):
img_name = osp.splitext(osp.basename(img_path))[0]
select_idx = data_util.index_generation(img_idx, max_idx, N_in, padding=padding)
imgs_in = imgs_LQ.index_select(0, torch.LongTensor(select_idx)).unsqueeze(0).to(device)
if flip_test:
output = util.flipx4_forward(model, imgs_in)
else:
output = util.single_forward(model, imgs_in)
output = util.tensor2img(output.squeeze(0))
if save_imgs:
cv2.imwrite(osp.join(save_subfolder, '{}.png'.format(img_name)), output)
logger.info('Folder {}'.format(subfolder_name))
logger.info('################ Final Results ################')
logger.info('Data: {} - {}'.format(data_mode, test_dataset_folder))
logger.info('Padding mode: {}'.format(padding))
logger.info('Model path: {}'.format(model_path))
logger.info('Save images: {}'.format(save_imgs))
logger.info('Flip test: {}'.format(flip_test))
def main():
#################
# configurations
#################
device = torch.device('cuda')
os.environ['CUDA_VISIBLE_DEVICES'] = '6'
test_set = 'AI4K_val' # Vid4 | YouKu10 | REDS4 | AI4K_val | zhibo | AI4K_val_bic
test_name = 'PCD_Vis_Test_35_ResNet_alpha_beta_decoder_3x3_IN_encoder_8HW_A01xxx_900000_AI4K_5000' # 'AI4K_val_Denoise_A02_420000'
data_mode = 'sharp_bicubic' # sharp_bicubic | blur_bicubic
N_in = 5
# load test set
if test_set == 'Vid4':
test_dataset_folder = '../datasets/Vid4/BIx4'
GT_dataset_folder = '../datasets/Vid4/GT'
elif test_set == 'YouKu10':
test_dataset_folder = '../datasets/YouKu10/LR'
GT_dataset_folder = '../datasets/YouKu10/HR'
elif test_set == 'YouKu_val':
test_dataset_folder = '/data0/yhliu/DATA/YouKuVid/valid/valid_lr_bmp'
GT_dataset_folder = '/data0/yhliu/DATA/YouKuVid/valid/valid_hr_bmp'
elif test_set == 'REDS4':
test_dataset_folder = '../datasets/REDS4/{}'.format(data_mode)
GT_dataset_folder = '../datasets/REDS4/GT'
elif test_set == 'AI4K_val':
test_dataset_folder = '/home/yhliu/AI4K/contest2/val2_LR_png/'
GT_dataset_folder = '/home/yhliu/AI4K/contest1/val1_HR_png/'
elif test_set == 'AI4K_val_bic':
test_dataset_folder = '/home/yhliu/AI4K/contest1/val1_LR_png_bic/'
GT_dataset_folder = '/home/yhliu/AI4K/contest1/val1_HR_png_bic/'
elif test_set == 'zhibo':
test_dataset_folder = '/data1/yhliu/SR_ZHIBO_VIDEO/Test_video_LR/'
GT_dataset_folder = '/data1/yhliu/SR_ZHIBO_VIDEO/Test_video_HR/'
flip_test = False
#model_path = '../experiments/pretrained_models/EDVR_Vimeo90K_SR_L.pth'
#model_path = '../experiments/A01b/models/250000_G.pth'
#model_path = '../experiments/A02_predenoise/models/415000_G.pth'
model_path = '../experiments/A37_color_EDVR_35_220000_A01_5in_64f_10b_128_pretrain_A01xxx_900000_fix_before_pcd/models/5000_G.pth'
predeblur, HR_in = False, False
back_RBs = 10
if data_mode == 'blur_bicubic':
predeblur = True
if data_mode == 'blur' or data_mode == 'blur_comp':
predeblur, HR_in = True, True
model = EDVR_arch.EDVR(64,
N_in,
8,
5,
back_RBs,
predeblur=predeblur,
HR_in=HR_in)
#model = my_EDVR_arch.MYEDVR(64, N_in, 8, 5, back_RBs, predeblur=predeblur, HR_in=HR_in)
#model = my_EDVR_arch.MYEDVR_RES(64, N_in, 8, 5, back_RBs, predeblur=predeblur, HR_in=HR_in)
#### evaluation
crop_border = 0
border_frame = N_in // 2 # border frames when evaluate
# temporal padding mode
if data_mode == 'Vid4' or data_mode == 'sharp_bicubic':
padding = 'new_info'
else:
padding = 'replicate'
save_imgs = True #True | False
save_folder = '../results/{}'.format(test_name)
if test_set == 'zhibo':
save_folder = '/data1/yhliu/SR_ZHIBO_VIDEO/SR_png_sample_150'
util.mkdirs(save_folder)
util.setup_logger('base',
save_folder,
'test',
level=logging.INFO,
screen=True,
tofile=True)
logger = logging.getLogger('base')
#### log info
logger.info('Data: {} - {}'.format(data_mode, test_dataset_folder))
logger.info('Padding mode: {}'.format(padding))
logger.info('Model path: {}'.format(model_path))
logger.info('Save images: {}'.format(save_imgs))
logger.info('Flip test: {}'.format(flip_test))
#### set up the models
model.load_state_dict(torch.load(model_path), strict=True)
model.eval()
model = model.to(device)
model = nn.DataParallel(model)
avg_psnr_l, avg_psnr_center_l, avg_psnr_border_l = [], [], []
subfolder_name_l = []
subfolder_l = sorted(glob.glob(osp.join(test_dataset_folder, '*')))
subfolder_GT_l = sorted(glob.glob(osp.join(GT_dataset_folder, '*')))
print(subfolder_l)
print(subfolder_GT_l)
#exit()
# for each subfolder
for subfolder, subfolder_GT in zip(subfolder_l, subfolder_GT_l):
subfolder_name = osp.basename(subfolder)
subfolder_name_l.append(subfolder_name)
save_subfolder = osp.join(save_folder, subfolder_name)
img_path_l = sorted(glob.glob(osp.join(subfolder, '*')))
print(img_path_l)
max_idx = len(img_path_l)
if save_imgs:
util.mkdirs(save_subfolder)
#### read LQ and GT images
imgs_LQ = data_util.read_img_seq(subfolder)
img_GT_l = []
for img_GT_path in sorted(glob.glob(osp.join(subfolder_GT, '*'))):
#print(img_GT_path)
img_GT_l.append(data_util.read_img(None, img_GT_path))
#print(img_GT_l[0].shape)
avg_psnr, avg_psnr_border, avg_psnr_center, N_border, N_center = 0, 0, 0, 0, 0
# process each image
for img_idx, img_path in enumerate(img_path_l):
img_name = osp.splitext(osp.basename(img_path))[0]
select_idx = data_util.index_generation(img_idx,
max_idx,
N_in,
padding=padding)
imgs_in = imgs_LQ.index_select(
0,
torch.LongTensor(select_idx)).unsqueeze(0).cpu() #to(device)
print(imgs_in.size())
if flip_test:
output = util.flipx4_forward(model, imgs_in)
else:
start_time = time.time()
output = util.single_forward(model, imgs_in)
end_time = time.time()
print('Forward One image:', end_time - start_time)
output = util.tensor2img(output.squeeze(0))
if save_imgs:
cv2.imwrite(
osp.join(save_subfolder, '{}.png'.format(img_name)),
output)
# calculate PSNR
output = output / 255.
GT = np.copy(img_GT_l[img_idx])
# For REDS, evaluate on RGB channels; for Vid4, evaluate on the Y channel
'''
if data_mode == 'Vid4': # bgr2y, [0, 1]
GT = data_util.bgr2ycbcr(GT, only_y=True)
output = data_util.bgr2ycbcr(output, only_y=True)
'''
output, GT = util.crop_border([output, GT], crop_border)
crt_psnr = util.calculate_psnr(output * 255, GT * 255)
logger.info('{:3d} - {:25} \tPSNR: {:.6f} dB'.format(
img_idx + 1, img_name, crt_psnr))
if img_idx >= border_frame and img_idx < max_idx - border_frame: # center frames
avg_psnr_center += crt_psnr
N_center += 1
else: # border frames
avg_psnr_border += crt_psnr
N_border += 1
avg_psnr = (avg_psnr_center + avg_psnr_border) / (N_center + N_border)
avg_psnr_center = avg_psnr_center / N_center
avg_psnr_border = 0 if N_border == 0 else avg_psnr_border / N_border
avg_psnr_l.append(avg_psnr)
avg_psnr_center_l.append(avg_psnr_center)
avg_psnr_border_l.append(avg_psnr_border)
logger.info('Folder {} - Average PSNR: {:.6f} dB for {} frames; '
'Center PSNR: {:.6f} dB for {} frames; '
'Border PSNR: {:.6f} dB for {} frames.'.format(
subfolder_name, avg_psnr, (N_center + N_border),
avg_psnr_center, N_center, avg_psnr_border, N_border))
logger.info('################ Tidy Outputs ################')
for subfolder_name, psnr, psnr_center, psnr_border in zip(
subfolder_name_l, avg_psnr_l, avg_psnr_center_l,
avg_psnr_border_l):
logger.info('Folder {} - Average PSNR: {:.6f} dB. '
'Center PSNR: {:.6f} dB. '
'Border PSNR: {:.6f} dB.'.format(subfolder_name, psnr,
psnr_center, psnr_border))
logger.info('################ Final Results ################')
logger.info('Data: {} - {}'.format(data_mode, test_dataset_folder))
logger.info('Padding mode: {}'.format(padding))
logger.info('Model path: {}'.format(model_path))
logger.info('Save images: {}'.format(save_imgs))
logger.info('Flip test: {}'.format(flip_test))
logger.info('Total Average PSNR: {:.6f} dB for {} clips. '
'Center PSNR: {:.6f} dB. Border PSNR: {:.6f} dB.'.format(
sum(avg_psnr_l) / len(avg_psnr_l), len(subfolder_l),
sum(avg_psnr_center_l) / len(avg_psnr_center_l),
sum(avg_psnr_border_l) / len(avg_psnr_border_l)))
def main():
#################
# configurations
#################
device = torch.device('cuda')
os.environ['CUDA_VISIBLE_DEVICES'] = '0'
data_mode = 'SDR_4bit'
stage = 1 # 1 or 2, use two stage strategy for REDS dataset.
flip_test = False
############################################################################
#### model
if data_mode == 'SDR_4bit':
if stage == 1:
model_path = '../experiments/pretrained_models/EDVR_REDS_deblurcomp_L.pth'
else:
model_path = '../experiments/pretrained_models/EDVR_REDS_deblurcomp_Stage2.pth'
else:
raise NotImplementedError
# use N_in images to restore one high bitdepth image
N_in = 5
# predeblur: predeblur for blurry input
# HR_in: downsample high resolution input
predeblur, HR_in = False, False
back_RBs = 40
predeblur = True
HR_in = True
if data_mode == 'SDR_4bit':
# predeblur, HR_in = False, True
pass
if stage == 2:
HR_in = True
back_RBs = 20
# EDVR(num_feature_map, num_input_frames, deformable_groups?, front_RBs,
# back_RBs, predeblur, HR_in)
model = EDVR_arch.EDVR(128,
N_in,
8,
5,
back_RBs,
predeblur=predeblur,
HR_in=HR_in)
#### dataset
if stage == 1:
test_dataset_folder = '../datasets/{}'.format(data_mode)
else:
test_dataset_folder = '../'
print('You should modify the test_dataset_folder path for stage 2')
GT_dataset_folder = '../datasets/SDR_10bit/'
#### evaluation
crop_border = 0
border_frame = N_in // 2 # border frames when evaluate
# temporal padding mode
padding = 'replicate'
save_imgs = True
save_folder = '../results/{}'.format(data_mode)
util.mkdirs(save_folder)
util.setup_logger('base',
save_folder,
'test',
level=logging.INFO,
screen=True,
tofile=True)
logger = logging.getLogger('base')
#### log info
logger.info('Data: {} - {}'.format(data_mode, test_dataset_folder))
logger.info('Padding mode: {}'.format(padding))
logger.info('Model path: {}'.format(model_path))
logger.info('Save images: {}'.format(save_imgs))
logger.info('Flip test: {}'.format(flip_test))
#### set up the models
model.load_state_dict(torch.load(model_path), strict=True)
model.eval()
model = model.to(device)
avg_psnr_l, avg_psnr_center_l, avg_psnr_border_l = [], [], []
subfolder_name_l = []
subfolder_l = sorted(glob.glob(osp.join(test_dataset_folder, '*')))
subfolder_GT_l = sorted(glob.glob(osp.join(GT_dataset_folder, '*')))
# for each subfolder
for subfolder, subfolder_GT in zip(subfolder_l, subfolder_GT_l):
subfolder_name = osp.basename(subfolder)
subfolder_name_l.append(subfolder_name)
save_subfolder = osp.join(save_folder, subfolder_name)
img_path_l = sorted(glob.glob(osp.join(subfolder, '*')))
max_idx = len(img_path_l)
if save_imgs:
util.mkdirs(save_subfolder)
#### read LBD and GT images
#### resize to avoid cuda out of memory, 2160x3840->720x1280
imgs_LBD = data_util.read_img_seq(subfolder,
scale=65535.,
zoomout=(1280, 720))
img_GT_l = []
for img_GT_path in sorted(glob.glob(osp.join(subfolder_GT, '*'))):
img_GT_l.append(
data_util.read_img(None,
img_GT_path,
scale=65535.,
zoomout=True))
avg_psnr, avg_psnr_border, avg_psnr_center, N_border, N_center = 0, 0, 0, 0, 0
# process each image
for img_idx, img_path in enumerate(img_path_l):
img_name = osp.splitext(osp.basename(img_path))[0]
# generate frame index
select_idx = data_util.index_generation(img_idx,
max_idx,
N_in,
padding=padding)
imgs_in = imgs_LBD.index_select(
0, torch.LongTensor(select_idx)).unsqueeze(0).to(device)
if flip_test:
# self ensemble with fipping input at four different directions
output = util.flipx4_forward(model, imgs_in)
else:
output = util.single_forward(model, imgs_in)
output = util.tensor2img(output.squeeze(0), out_type=np.uint16)
if save_imgs:
cv2.imwrite(
osp.join(save_subfolder, '{}.png'.format(img_name)),
output)
# calculate PSNR
# output = output / 255.
output = output / 65535.
GT = np.copy(img_GT_l[img_idx])
# For REDS, evaluate on RGB channels; for Vid4, evaluate on the Y channel
if data_mode == 'Vid4': # bgr2y, [0, 1]
GT = data_util.bgr2ycbcr(GT, only_y=True)
output = data_util.bgr2ycbcr(output, only_y=True)
output, GT = util.crop_border([output, GT], crop_border)
crt_psnr = util.calculate_psnr(output * 65535, GT * 65535)
logger.info('{:3d} - {:25} \tPSNR: {:.6f} dB'.format(
img_idx + 1, img_name, crt_psnr))
if img_idx >= border_frame and img_idx < max_idx - border_frame: # center frames
avg_psnr_center += crt_psnr
N_center += 1
else: # border frames
avg_psnr_border += crt_psnr
N_border += 1
avg_psnr = (avg_psnr_center + avg_psnr_border) / (N_center + N_border)
avg_psnr_center = avg_psnr_center / N_center
avg_psnr_border = 0 if N_border == 0 else avg_psnr_border / N_border
avg_psnr_l.append(avg_psnr)
avg_psnr_center_l.append(avg_psnr_center)
avg_psnr_border_l.append(avg_psnr_border)
logger.info('Folder {} - Average PSNR: {:.6f} dB for {} frames; '
'Center PSNR: {:.6f} dB for {} frames; '
'Border PSNR: {:.6f} dB for {} frames.'.format(
subfolder_name, avg_psnr, (N_center + N_border),
avg_psnr_center, N_center, avg_psnr_border, N_border))
logger.info('################ Tidy Outputs ################')
for subfolder_name, psnr, psnr_center, psnr_border in zip(
subfolder_name_l, avg_psnr_l, avg_psnr_center_l,
avg_psnr_border_l):
logger.info('Folder {} - Average PSNR: {:.6f} dB. '
'Center PSNR: {:.6f} dB. '
'Border PSNR: {:.6f} dB.'.format(subfolder_name, psnr,
psnr_center, psnr_border))
logger.info('################ Final Results ################')
logger.info('Data: {} - {}'.format(data_mode, test_dataset_folder))
logger.info('Padding mode: {}'.format(padding))
logger.info('Model path: {}'.format(model_path))
logger.info('Save images: {}'.format(save_imgs))
logger.info('Flip test: {}'.format(flip_test))
logger.info('Total Average PSNR: {:.6f} dB for {} clips. '
'Center PSNR: {:.6f} dB. Border PSNR: {:.6f} dB.'.format(
sum(avg_psnr_l) / len(avg_psnr_l), len(subfolder_l),
sum(avg_psnr_center_l) / len(avg_psnr_center_l),
sum(avg_psnr_border_l) / len(avg_psnr_border_l)))
def main():
#################
# configurations
#################
device = torch.device('cuda')
os.environ['CUDA_VISIBLE_DEVICES'] = '1'
stage = 1 # 1 or 2, use two stage strategy for REDS dataset.
flip_test = False
#### model
data_mode = 'sharp'
if stage == 1:
model_path = '../experiments/001_EDVRwoTSA_scratch_lr4e-4_600k_SR4K_LrCAR4S_64_20_5/models/600000_G.pth'
else:
model_path = '../experiments/pretrained_models/EDVR_REDS_SR_Stage2.pth'
N_in = 5 # use N_in images to restore one HR image
predeblur, HR_in = False, False
back_RBs = 20
if stage == 2:
HR_in = True
back_RBs = 20
model = EDVR_arch.EDVR(64, N_in, 8, 5, back_RBs, predeblur=predeblur, HR_in=HR_in, w_TSA=True)
#### dataset
if stage == 1:
test_dataset_folder = '/home/mcc/4khdr/image/540p_test'
else:
test_dataset_folder = '../results/REDS-EDVR_REDS_SR_L_flipx4'
print('You should modify the test_dataset_folder path for stage 2')
#### evaluation
crop_border = 0
border_frame = N_in // 2 # border frames when evaluate
# temporal padding mode
if data_mode == 'sharp':
padding = 'new_info'
else:
padding = 'replicate'
save_imgs = True
save_folder = '../results/{}'.format(data_mode)
util.mkdirs(save_folder)
util.setup_logger('base', save_folder, 'test', level=logging.INFO, screen=True, tofile=True)
logger = logging.getLogger('base')
#### log info
logger.info('Data: {} - {}'.format(data_mode, test_dataset_folder))
logger.info('Padding mode: {}'.format(padding))
logger.info('Model path: {}'.format(model_path))
logger.info('Save images: {}'.format(save_imgs))
logger.info('Flip test: {}'.format(flip_test))
#### set up the models
model.load_state_dict(torch.load(model_path), strict=True)
model.eval()
model = model.to(device)
subfolder_name_l = []
subfolder_l = sorted(glob.glob(osp.join(test_dataset_folder, '*')))
# for each subfolder
for subfolder in subfolder_l:
subfolder_name = osp.basename(subfolder)
subfolder_name_l.append(subfolder_name)
save_subfolder = osp.join(save_folder, subfolder_name)
img_path_l = sorted(glob.glob(osp.join(subfolder, '*')))
max_idx = len(img_path_l)
if save_imgs:
util.mkdirs(save_subfolder)
#### read LQ and GT images
imgs_LQ = data_util.read_img_seq(subfolder)
avg_psnr, avg_psnr_border, avg_psnr_center, N_border, N_center = 0, 0, 0, 0, 0
# process each image
for img_idx, img_path in enumerate(img_path_l):
img_name = osp.splitext(osp.basename(img_path))[0]
select_idx = data_util.index_generation(img_idx, max_idx, N_in, padding=padding)
imgs_in = imgs_LQ.index_select(0, torch.LongTensor(select_idx)).unsqueeze(0).to(device)
if flip_test:
output = util.flipx4_forward(model, imgs_in)
else:
output = util.single_forward(model, imgs_in)
output = util.tensor2img(output.squeeze(0))
if save_imgs:
cv2.imwrite(osp.join(save_subfolder, '{}.png'.format(img_name)), output,
[int(cv2.IMWRITE_PNG_COMPRESSION), 1])
logger.info('{:3d} - {:25}'.format(img_idx + 1, img_name))
def main():
#################
# configurations
#################
device = torch.device('cuda')
os.environ['CUDA_VISIBLE_DEVICES'] = '1'
data_mode = 'Vid4' # Vid4 | sharp_bicubic | blur_bicubic | blur | blur_comp
# Vid4: SR
# REDS4: sharp_bicubic (SR-clean), blur_bicubic (SR-blur);
# blur (deblur-clean), blur_comp (deblur-compression).
stage = 1 # 1 or 2, use two stage strategy for REDS dataset.
flip_test = False
############################################################################
#### model
if data_mode == 'Vid4':
if stage == 1:
#model_path = '../experiments/pretrained_models/EDVR_REDS_SR_M.pth'
model_path = '../experiments/002_EDVR_lr4e-4_600k_AI4KHDR/models/4000_G.pth'
else:
raise ValueError('Vid4 does not support stage 2.')
elif data_mode == 'sharp_bicubic':
if stage == 1:
model_path = '../experiments/pretrained_models/EDVR_REDS_SR_L.pth'
else:
model_path = '../experiments/pretrained_models/EDVR_REDS_SR_Stage2.pth'
elif data_mode == 'blur_bicubic':
if stage == 1:
model_path = '../experiments/pretrained_models/EDVR_REDS_SRblur_L.pth'
else:
model_path = '../experiments/pretrained_models/EDVR_REDS_SRblur_Stage2.pth'
elif data_mode == 'blur':
if stage == 1:
model_path = '../experiments/pretrained_models/EDVR_REDS_deblur_L.pth'
else:
model_path = '../experiments/pretrained_models/EDVR_REDS_deblur_Stage2.pth'
elif data_mode == 'blur_comp':
if stage == 1:
model_path = '../experiments/pretrained_models/EDVR_REDS_deblurcomp_L.pth'
else:
model_path = '../experiments/pretrained_models/EDVR_REDS_deblurcomp_Stage2.pth'
else:
raise NotImplementedError
if data_mode == 'Vid4':
N_in = 5 # use N_in images to restore one HR image
else:
N_in = 5
predeblur, HR_in = False, False
back_RBs = 10
if data_mode == 'blur_bicubic':
predeblur = True
if data_mode == 'blur' or data_mode == 'blur_comp':
predeblur, HR_in = True, True
if stage == 2:
HR_in = True
back_RBs = 20
model = EDVR_arch.EDVR(64, 5, 8, 5, 10, predeblur=predeblur, HR_in=HR_in)
#### dataset
if data_mode == 'Vid4':
test_dataset_folder = '/workspace/nas_mengdongwei/dataset/AI4KHDR/valid/540p_frames'
GT_dataset_folder = '/workspace/nas_mengdongwei/dataset/AI4KHDR/valid/4k_frames'
#test_dataset_folder = '../datasets/Vid4/BIx4'
#GT_dataset_folder = '../datasets/Vid4/GT'
else:
if stage == 1:
test_dataset_folder = '../datasets/REDS4/{}'.format(data_mode)
else:
test_dataset_folder = '../results/REDS-EDVR_REDS_SR_L_flipx4'
print('You should modify the test_dataset_folder path for stage 2')
GT_dataset_folder = '../datasets/REDS4/GT'
#### evaluation
crop_border = 0
border_frame = N_in // 2 # border frames when evaluate
# temporal padding mode
if data_mode == 'Vid4' or data_mode == 'sharp_bicubic':
padding = 'new_info'
else:
padding = 'replicate'
save_imgs = True
save_folder = '../results/{}'.format(data_mode)
util.mkdirs(save_folder)
util.setup_logger('base',
save_folder,
'test',
level=logging.INFO,
screen=True,
tofile=True)
logger = logging.getLogger('base')
#### log info
logger.info('Data: {} - {}'.format(data_mode, test_dataset_folder))
logger.info('Padding mode: {}'.format(padding))
logger.info('Model path: {}'.format(model_path))
logger.info('Save images: {}'.format(save_imgs))
logger.info('Flip test: {}'.format(flip_test))
#### set up the models
model.load_state_dict(torch.load(model_path), strict=False)
model.eval()
model = model.to(device)
avg_psnr_l, avg_psnr_center_l, avg_psnr_border_l = [], [], []
subfolder_name_l = []
subfolder_l = sorted(glob.glob(osp.join(test_dataset_folder, '*')))
subfolder_GT_l = sorted(glob.glob(osp.join(GT_dataset_folder, '*')))
# for each subfolder
for subfolder, subfolder_GT in zip(subfolder_l, subfolder_GT_l):
subfolder_name = osp.basename(subfolder)
subfolder_name_l.append(subfolder_name)
save_subfolder = osp.join(save_folder, subfolder_name)
img_path_l = sorted(glob.glob(osp.join(subfolder, '*')))
max_idx = len(img_path_l)
if save_imgs:
util.mkdirs(save_subfolder)
#### read LQ and GT images
imgs_LQ = data_util.read_img_seq(subfolder)
img_GT_l = []
for img_GT_path in sorted(glob.glob(osp.join(subfolder_GT, '*'))):
img_GT_l.append(data_util.read_img(None, img_GT_path))
avg_psnr, avg_psnr_border, avg_psnr_center, N_border, N_center = 0, 0, 0, 0, 0
# process each image
for img_idx, img_path in enumerate(img_path_l):
img_name = osp.splitext(osp.basename(img_path))[0]
select_idx = data_util.index_generation(img_idx,
max_idx,
N_in,
padding=padding)
imgs_in = imgs_LQ.index_select(
0, torch.LongTensor(select_idx)).unsqueeze(0).to(device)
if flip_test:
output = util.flipx4_forward(model, imgs_in)
else:
output = util.single_forward(model, imgs_in)
output = util.tensor2img(output.squeeze(0))
if save_imgs:
cv2.imwrite(
osp.join(save_subfolder, '{}.png'.format(img_name)),
output)
# calculate PSNR
output = output / 255.
GT = np.copy(img_GT_l[img_idx])
# For REDS, evaluate on RGB channels; for Vid4, evaluate on the Y channel
if data_mode == 'Vid4': # bgr2y, [0, 1]
GT = data_util.bgr2ycbcr(GT, only_y=True)
output = data_util.bgr2ycbcr(output, only_y=True)
output, GT = util.crop_border([output, GT], crop_border)
crt_psnr = util.calculate_psnr(output * 255, GT * 255)
#logger.info('{:3d} - {:25} \tPSNR: {:.6f} dB'.format(img_idx + 1, img_name, crt_psnr))
if img_idx >= border_frame and img_idx < max_idx - border_frame: # center frames
avg_psnr_center += crt_psnr
N_center += 1
else: # border frames
avg_psnr_border += crt_psnr
N_border += 1
avg_psnr = (avg_psnr_center + avg_psnr_border) / (N_center + N_border)
avg_psnr_center = avg_psnr_center / N_center
avg_psnr_border = 0 if N_border == 0 else avg_psnr_border / N_border
avg_psnr_l.append(avg_psnr)
avg_psnr_center_l.append(avg_psnr_center)
avg_psnr_border_l.append(avg_psnr_border)
logger.info('Folder {} - Average PSNR: {:.6f} dB for {} frames; '
'Center PSNR: {:.6f} dB for {} frames; '
'Border PSNR: {:.6f} dB for {} frames.'.format(
subfolder_name, avg_psnr, (N_center + N_border),
avg_psnr_center, N_center, avg_psnr_border, N_border))
logger.info('################ Tidy Outputs ################')
for subfolder_name, psnr, psnr_center, psnr_border in zip(
subfolder_name_l, avg_psnr_l, avg_psnr_center_l,
avg_psnr_border_l):
logger.info('Folder {} - Average PSNR: {:.6f} dB. '
'Center PSNR: {:.6f} dB. '
'Border PSNR: {:.6f} dB.'.format(subfolder_name, psnr,
psnr_center, psnr_border))
logger.info('################ Final Results ################')
logger.info('Data: {} - {}'.format(data_mode, test_dataset_folder))
logger.info('Padding mode: {}'.format(padding))
logger.info('Model path: {}'.format(model_path))
logger.info('Save images: {}'.format(save_imgs))
logger.info('Flip test: {}'.format(flip_test))
logger.info('Total Average PSNR: {:.6f} dB for {} clips. '
'Center PSNR: {:.6f} dB. Border PSNR: {:.6f} dB.'.format(
sum(avg_psnr_l) / len(avg_psnr_l), len(subfolder_l),
sum(avg_psnr_center_l) / len(avg_psnr_center_l),
sum(avg_psnr_border_l) / len(avg_psnr_border_l)))
def main():
#################
# configurations
#################
device = torch.device('cuda')
# os.environ['CUDA_VISIBLE_DEVICES'] = '0'
prog = argparse.ArgumentParser()
prog.add_argument('--dataset_mode',
'-m',
type=str,
default='Vid4+REDS',
help='data_mode')
prog.add_argument('--degradation_mode',
'-d',
type=str,
default='impulse',
choices=('impulse', 'bicubic', 'preset'),
help='path to image output directory.')
prog.add_argument('--sigma_x',
'-sx',
type=float,
default=1,
help='sigma_x')
prog.add_argument('--sigma_y',
'-sy',
type=float,
default=0,
help='sigma_y')
prog.add_argument('--theta', '-t', type=float, default=0, help='theta')
prog.add_argument('--scale',
'-sc',
type=int,
default=2,
choices=(2, 4),
help='scale factor')
args = prog.parse_args()
data_modes = args.dataset_mode
degradation_mode = args.degradation_mode # impulse | bicubic
sig_x, sig_y, the = args.sigma_x, args.sigma_y, args.theta * math.pi / 180
if sig_y == 0:
sig_y = sig_x
scale = args.scale
kernel_size = 21
N_frames = 7
padding = 'new_info'
data_mode_l = data_modes.split('+')
for i in range(len(data_mode_l)):
data_mode = data_mode_l[i]
#### dataset
if data_mode == 'Vid4':
kernel_folder = '../pretrained_models/Mixed/Vid4.npy'
dataset_folder = '../dataset/Vid4'
elif data_mode == 'REDS':
kernel_folder = '../pretrained_models/Mixed/REDS.npy'
dataset_folder = '../dataset/REDS/train'
elif data_mode == 'Vimeo':
if degradation_mode == 'preset':
raise NotImplementedError(
'We do not support preset mode in Vimeo dataset')
dataset_folder = '../dataset/vimeo_septuplet'
else:
raise NotImplementedError()
save_folder_name = 'preset' if degradation_mode == 'preset' else degradation_mode + '_' + str(
'{:.1f}'.format(sig_x)) + '_' + str('{:.1f}'.format(
sig_y)) + '_' + str('{:.1f}'.format(args.theta))
save_folder = osp.join(dataset_folder, 'LR_' + save_folder_name, 'X2')
if not osp.exists(save_folder):
os.makedirs(save_folder)
save_folder2 = osp.join(dataset_folder, 'LR_' + save_folder_name, 'X4')
if not osp.exists(save_folder2):
os.makedirs(save_folder2)
if scale == 4:
save_folder3 = osp.join(dataset_folder, 'LR_' + save_folder_name,
'X16')
if not osp.exists(save_folder3):
os.makedirs(save_folder3)
if data_mode == 'Vimeo':
GT_dataset_folder = osp.join(dataset_folder, 'sequences')
meta = osp.join(dataset_folder, 'sep_testlist.txt')
with open(meta, 'r') as f:
seq_list = sorted(f.read().splitlines())
subfolder_GT_l = [
osp.join(GT_dataset_folder, seq_ind) for seq_ind in seq_list
]
else:
GT_dataset_folder = osp.join(dataset_folder, 'HR')
subfolder_GT_l = sorted(glob.glob(osp.join(GT_dataset_folder,
'*')))
if data_mode == 'REDS':
subfolder_GT_l = [
k for k in subfolder_GT_l
if k.find('000') >= 0 or k.find('011') >= 0
or k.find('015') >= 0 or k.find('020') >= 0
]
sig_x, sig_y, the = float(sig_x), float(sig_y), float(the)
for subfolder_GT in subfolder_GT_l:
print(subfolder_GT)
gen_kwargs = set_kernel_params(sigma_x=sig_x,
sigma_y=sig_y,
theta=the)
if degradation_mode == 'impulse' or degradation_mode == 'preset':
kernel_gen = rkg.Degradation(kernel_size, 2, **gen_kwargs)
if degradation_mode == 'preset':
kernel_preset = np.load(kernel_folder)
else:
kernel_gen = oldkg.Degradation(kernel_size,
2,
type=0.7,
**gen_kwargs)
if data_mode == 'Vimeo':
sub1 = osp.basename(osp.dirname(subfolder_GT))
sub2 = osp.basename(subfolder_GT)
subfolder_name = osp.join(sub1, sub2)
else:
subfolder_name = osp.basename(subfolder_GT)
save_subfolder = osp.join(save_folder, subfolder_name)
if not osp.exists(save_subfolder):
os.makedirs(save_subfolder)
save_subfolder2 = osp.join(save_folder2, subfolder_name)
if not osp.exists(save_subfolder2):
os.makedirs(save_subfolder2)
if scale == 4:
save_subfolder3 = osp.join(save_folder3, subfolder_name)
if not osp.exists(save_subfolder3):
os.makedirs(save_subfolder3)
img_GT_path_l = sorted(glob.glob(osp.join(subfolder_GT, '*')))
seq_length = len(img_GT_path_l)
imgs_GT = data_util.read_img_seq(subfolder_GT) # T C H W
if degradation_mode == 'preset':
for index in range(seq_length):
save_subsubfolder = osp.join(
save_subfolder,
osp.splitext(osp.basename(img_GT_path_l[index]))[0])
save_subsubfolder2 = osp.join(
save_subfolder2,
osp.splitext(osp.basename(img_GT_path_l[index]))[0])
if not osp.exists(save_subsubfolder):
os.mkdir(save_subsubfolder)
if not osp.exists(save_subsubfolder2):
os.mkdir(save_subsubfolder2)
if scale == 4 and not osp.exists(save_subsubfolder3):
os.mkdir(save_subsubfolder3)
kernel_gen.set_kernel_directly(kernel_preset[index])
imgs_HR = imgs_GT[data_util.index_generation(
index, seq_length, N_frames, padding)]
imgs_LR = kernel_gen.apply(imgs_HR)
imgs_LR = imgs_LR.mul(255).clamp(0, 255).round().div(255)
imgs_LR_np = imgs_LR.permute(0, 2, 3, 1).cpu().numpy()
imgs_LR_np = (imgs_LR_np * 255).astype('uint8')
for i, img_LR in enumerate(imgs_LR_np):
imageio.imwrite(
osp.join(save_subsubfolder, 'img{}.png'.format(i)),
img_LR)
imgs_SuperLR = kernel_gen.apply(imgs_LR)
imgs_SuperLR = imgs_SuperLR.mul(255).clamp(
0, 255).round().div(255)
imgs_SuperLR_np = imgs_SuperLR.permute(0, 2, 3,
1).cpu().numpy()
imgs_SuperLR_np = (imgs_SuperLR_np * 255).astype('uint8')
for i, img_SuperLR in enumerate(imgs_SuperLR_np):
imageio.imwrite(
osp.join(save_subsubfolder2,
'img{}.png'.format(i)), img_SuperLR)
if scale == 4:
imgs_SuperLR = kernel_gen.apply(imgs_LR)
imgs_SuperLR = imgs_SuperLR.mul(255).clamp(
0, 255).round().div(255)
imgs_SuperLR = kernel_gen.apply(imgs_LR)
imgs_SuperLR = imgs_SuperLR.mul(255).clamp(
0, 255).round().div(255)
imgs_SuperLR_np = imgs_SuperLR.permute(
0, 2, 3, 1).cpu().numpy()
imgs_SuperLR_np = (imgs_SuperLR_np *
255).astype('uint8')
for i, img_SuperLR in enumerate(imgs_SuperLR_np):
imageio.imwrite(
osp.join(save_subsubfolder3,
'img{}.png'.format(i)), img_SuperLR)
else:
count = 0
imgs_GT_l = imgs_GT.split(32)
for img_batch in imgs_GT_l:
img_lr_batch = kernel_gen.apply(img_batch)
img_lr_batch = img_lr_batch.permute(0, 2, 3,
1).cpu().numpy()
img_lr_batch = (img_lr_batch.clip(0, 1) * 255).round()
img_lr_batch = img_lr_batch.astype('uint8')
count_temp = count
for img_lr in img_lr_batch:
filename = osp.basename(img_GT_path_l[count])
imageio.imwrite(osp.join(save_subfolder, filename),
img_lr)
count += 1
img_lr_batch = img_lr_batch.astype('float32') / 255
img_lr_batch = torch.from_numpy(img_lr_batch).permute(
0, 3, 1, 2)
img_superlr_batch = kernel_gen.apply(img_lr_batch)
img_superlr_batch = img_superlr_batch.permute(
0, 2, 3, 1).cpu().numpy()
img_superlr_batch = (img_superlr_batch.clip(0, 1) *
255).round()
img_superlr_batch = img_superlr_batch.astype('uint8')
count = count_temp
for img_superlr in img_superlr_batch:
filename = osp.basename(img_GT_path_l[count])
imageio.imwrite(osp.join(save_subfolder2, filename),
img_superlr)
count += 1
if scale == 4:
img_superlr_batch = img_superlr_batch.astype(
'float32') / 255
img_superlr_batch = torch.from_numpy(
img_superlr_batch).permute(0, 3, 1, 2)
img_superlr_batch = kernel_gen.apply(img_superlr_batch)
img_superlr_batch = img_superlr_batch.permute(
0, 2, 3, 1).cpu().numpy()
img_superlr_batch = (img_superlr_batch.clip(0, 1) *
255).round()
img_superlr_batch = img_superlr_batch.astype(
'float32') / 255
img_superlr_batch = torch.from_numpy(
img_superlr_batch).permute(0, 3, 1, 2)
img_superlr_batch = kernel_gen.apply(img_superlr_batch)
img_superlr_batch = img_superlr_batch.permute(
0, 2, 3, 1).cpu().numpy()
img_superlr_batch = (img_superlr_batch.clip(0, 1) *
255).round()
img_superlr_batch = img_superlr_batch.astype('uint8')
count = count_temp
for img_superlr in img_superlr_batch:
filename = osp.basename(img_GT_path_l[count])
imageio.imwrite(
osp.join(save_subfolder2, filename),
img_superlr)
count += 1
def __init__(self, opt):
super(VideoTestDataset, self).__init__()
self.opt = opt
self.cache_data = opt['cache_data']
self.half_N_frames = opt['N_frames'] // 2
self.GT_root, self.LQ_root = opt['dataroot_GT'], opt['dataroot_LQ']
self.data_type = self.opt['data_type']
self.data_info = {
'path_LQ': [],
'path_GT': [],
'folder': [],
'idx': [],
'border': []
}
if self.data_type == 'lmdb':
raise ValueError('No need to use LMDB during validation/test.')
#### Generate data info and cache data
self.imgs_LQ, self.imgs_GT = {}, {}
if opt['name'].lower() in ['vid4', 'reds4']:
subfolders_LQ = util.glob_file_list(self.LQ_root)
subfolders_GT = util.glob_file_list(self.GT_root)
for subfolder_LQ, subfolder_GT in zip(subfolders_LQ,
subfolders_GT):
subfolder_name = osp.basename(subfolder_GT)
img_paths_LQ = util.glob_file_list(subfolder_LQ)
img_paths_GT = util.glob_file_list(subfolder_GT)
max_idx = len(img_paths_LQ)
assert max_idx == len(
img_paths_GT
), 'Different number of images in LQ and GT folders'
if opt['pred_interval'] == 0:
self.data_info['path_LQ'].extend(img_paths_LQ)
self.data_info['path_GT'].extend(img_paths_GT)
self.data_info['folder'].extend([subfolder_name] * max_idx)
for i in range(max_idx):
self.data_info['idx'].append('{}/{}'.format(
i, max_idx))
elif opt['pred_interval'] < 0:
LQs = []
GTs = []
LQs.append(img_paths_LQ)
GTs.append(img_paths_GT)
self.data_info['path_LQ'].extend(LQs)
self.data_info['path_GT'].extend(GTs)
self.data_info['folder'].extend([subfolder_name])
self.data_info['idx'].append('{}/{}'.format(1, 1))
else:
self.pred_interval = opt['pred_interval']
LQs = []
GTs = []
if max_idx % self.pred_interval == 1 or max_idx % self.pred_interval == 0:
num_clip = max_idx // self.pred_interval
else:
num_clip = max_idx // self.pred_interval + 1
for i in range(num_clip):
if i != max_idx // self.pred_interval:
LQs.append(
img_paths_LQ[i * self.pred_interval:(i + 1) *
self.pred_interval + 1])
GTs.append(
img_paths_GT[i * self.pred_interval:(i + 1) *
self.pred_interval + 1])
else:
LQs.append(img_paths_LQ[i * self.pred_interval:])
GTs.append(img_paths_GT[i * self.pred_interval:])
self.data_info['path_LQ'].extend(LQs)
self.data_info['path_GT'].extend(GTs)
self.data_info['folder'].extend([subfolder_name] *
num_clip)
for i in range(max_idx // self.pred_interval + 1):
self.data_info['idx'].append('{}/{}'.format(
i, num_clip))
if self.cache_data:
self.imgs_LQ[subfolder_name] = util.read_img_seq(
img_paths_LQ)
self.imgs_GT[subfolder_name] = util.read_img_seq(
img_paths_GT)
elif opt['name'].lower() in ['vimeo90k-test']:
pass # TODO
else:
raise ValueError(
'Not support video test dataset. Support Vid4, REDS4 and Vimeo90k-Test.'
)
def main():
####################
# arguments parser #
####################
# [format] dataset(vid4, REDS4) N(number of frames)
parser = argparse.ArgumentParser()
parser.add_argument('dataset')
parser.add_argument('n_frames')
parser.add_argument('stage')
args = parser.parse_args()
data_mode = str(args.dataset)
N_in = int(args.n_frames)
stage = int(args.stage)
#if args.command == 'start':
# start(int(args.params[0]))
#elif args.command == 'stop':
# stop(args.params[0], int(args.params[1]))
#elif args.command == 'stop_all':
# stop_all(args.params[0])
#################
# configurations
#################
device = torch.device('cuda')
os.environ['CUDA_VISIBLE_DEVICES'] = '0'
#data_mode = 'Vid4' # Vid4 | sharp_bicubic | blur_bicubic | blur | blur_comp
# Vid4: SR
# REDS4: sharp_bicubic (SR-clean), blur_bicubic (SR-blur);
# blur (deblur-clean), blur_comp (deblur-compression).
#stage = 1 # 1 or 2, use two stage strategy for REDS dataset.
flip_test = False
############################################################################
#### model
if data_mode == 'Vid4':
if stage == 1:
model_path = '../experiments/pretrained_models/EDVR_Vimeo90K_SR_L.pth'
else:
raise ValueError('Vid4 does not support stage 2.')
elif data_mode == 'sharp_bicubic':
if stage == 1:
model_path = '../experiments/pretrained_models/EDVR_REDS_SR_L.pth'
else:
model_path = '../experiments/pretrained_models/EDVR_REDS_SR_Stage2.pth'
elif data_mode == 'blur_bicubic':
if stage == 1:
model_path = '../experiments/pretrained_models/EDVR_REDS_SRblur_L.pth'
else:
model_path = '../experiments/pretrained_models/EDVR_REDS_SRblur_Stage2.pth'
elif data_mode == 'blur':
if stage == 1:
model_path = '../experiments/pretrained_models/EDVR_REDS_deblur_L.pth'
else:
model_path = '../experiments/pretrained_models/EDVR_REDS_deblur_Stage2.pth'
elif data_mode == 'blur_comp':
if stage == 1:
model_path = '../experiments/pretrained_models/EDVR_REDS_deblurcomp_L.pth'
else:
model_path = '../experiments/pretrained_models/EDVR_REDS_deblurcomp_Stage2.pth'
else:
raise NotImplementedError
predeblur, HR_in = False, False
back_RBs = 40
if data_mode == 'blur_bicubic':
predeblur = True
if data_mode == 'blur' or data_mode == 'blur_comp':
predeblur, HR_in = True, True
if stage == 2:
HR_in = True
back_RBs = 20
#### dataset
if data_mode == 'Vid4':
N_model_default = 7
test_dataset_folder = '../datasets/Vid4/BIx4'
GT_dataset_folder = '../datasets/Vid4/GT'
else:
N_model_default = 5
if stage == 1:
test_dataset_folder = '../datasets/REDS4/{}'.format(data_mode)
else:
test_dataset_folder = '../results/REDS-EDVR_REDS_SR_L_flipx4'
print('You should modify the test_dataset_folder path for stage 2')
GT_dataset_folder = '../datasets/REDS4/GT'
raw_model = EDVR_arch.EDVR(128,
N_model_default,
8,
5,
back_RBs,
predeblur=predeblur,
HR_in=HR_in)
model = EDVR_arch.EDVR(128,
N_in,
8,
5,
back_RBs,
predeblur=predeblur,
HR_in=HR_in)
#### evaluation
crop_border = 0
border_frame = N_in // 2 # border frames when evaluate
# temporal padding mode
if data_mode == 'Vid4' or data_mode == 'sharp_bicubic':
padding = 'new_info'
else:
padding = 'replicate'
save_imgs = True
data_mode_t = copy.deepcopy(data_mode)
if stage == 1 and data_mode_t != 'Vid4':
data_mode = 'REDS-EDVR_REDS_SR_L_flipx4'
save_folder = '../results/{}'.format(data_mode)
data_mode = copy.deepcopy(data_mode_t)
util.mkdirs(save_folder)
util.setup_logger('base',
save_folder,
'test',
level=logging.INFO,
screen=True,
tofile=True)
logger = logging.getLogger('base')
#### log info
logger.info('Data: {} - {}'.format(data_mode, test_dataset_folder))
logger.info('Padding mode: {}'.format(padding))
logger.info('Model path: {}'.format(model_path))
logger.info('Save images: {}'.format(save_imgs))
logger.info('Flip test: {}'.format(flip_test))
#### set up the models
print([a for a in dir(model)
if not callable(getattr(model, a))]) # not a.startswith('__') and
#model.load_state_dict(torch.load(model_path), strict=True)
raw_model.load_state_dict(torch.load(model_path), strict=True)
# model.load_state_dict(torch.load(model_path), strict=True)
#### change model so it can work with less input
model.nf = raw_model.nf
model.center = N_in // 2 # if center is None else center
model.is_predeblur = raw_model.is_predeblur
model.HR_in = raw_model.HR_in
model.w_TSA = raw_model.w_TSA
#ResidualBlock_noBN_f = functools.partial(arch_util.ResidualBlock_noBN, nf=nf)
#### extract features (for each frame)
if model.is_predeblur:
model.pre_deblur = raw_model.pre_deblur #Predeblur_ResNet_Pyramid(nf=nf, HR_in=self.HR_in)
model.conv_1x1 = raw_model.conv_1x1 #nn.Conv2d(nf, nf, 1, 1, bias=True)
else:
if model.HR_in:
model.conv_first_1 = raw_model.conv_first_1 #nn.Conv2d(3, nf, 3, 1, 1, bias=True)
model.conv_first_2 = raw_model.conv_first_2 #nn.Conv2d(nf, nf, 3, 2, 1, bias=True)
model.conv_first_3 = raw_model.conv_first_3 #nn.Conv2d(nf, nf, 3, 2, 1, bias=True)
else:
model.conv_first = raw_model.conv_first # nn.Conv2d(3, nf, 3, 1, 1, bias=True)
model.feature_extraction = raw_model.feature_extraction # arch_util.make_layer(ResidualBlock_noBN_f, front_RBs)
model.fea_L2_conv1 = raw_model.fea_L2_conv1 #nn.Conv2d(nf, nf, 3, 2, 1, bias=True)
model.fea_L2_conv2 = raw_model.fea_L2_conv2 #nn.Conv2d(nf, nf, 3, 1, 1, bias=True)
model.fea_L3_conv1 = raw_model.fea_L3_conv1 #nn.Conv2d(nf, nf, 3, 2, 1, bias=True)
model.fea_L3_conv2 = raw_model.fea_L3_conv2 #nn.Conv2d(nf, nf, 3, 1, 1, bias=True)
model.pcd_align = raw_model.pcd_align #PCD_Align(nf=nf, groups=groups)
######## Resize TSA
model.tsa_fusion.center = model.center
# temporal attention (before fusion conv)
model.tsa_fusion.tAtt_1 = raw_model.tsa_fusion.tAtt_1
model.tsa_fusion.tAtt_2 = raw_model.tsa_fusion.tAtt_2
# fusion conv: using 1x1 to save parameters and computation
#print(raw_model.tsa_fusion.fea_fusion.weight.shape)
#print(raw_model.tsa_fusion.fea_fusion.weight.shape)
#print(raw_model.tsa_fusion.fea_fusion.weight[127][639].shape)
#print("MAIN SHAPE(FEA): ", raw_model.tsa_fusion.fea_fusion.weight.shape)
model.tsa_fusion.fea_fusion = copy.deepcopy(
raw_model.tsa_fusion.fea_fusion)
model.tsa_fusion.fea_fusion.weight = copy.deepcopy(
torch.nn.Parameter(raw_model.tsa_fusion.fea_fusion.weight[:, 0:N_in *
128, :, :]))
#[:][] #nn.Conv2d(nframes * nf, nf, 1, 1, bias=True)
#model.tsa_fusion.fea_fusion.bias = raw_model.tsa_fusion.fea_fusion.bias
# spatial attention (after fusion conv)
model.tsa_fusion.sAtt_1 = copy.deepcopy(raw_model.tsa_fusion.sAtt_1)
model.tsa_fusion.sAtt_1.weight = copy.deepcopy(
torch.nn.Parameter(raw_model.tsa_fusion.sAtt_1.weight[:, 0:N_in *
128, :, :]))
#[:][] #nn.Conv2d(nframes * nf, nf, 1, 1, bias=True)
#model.tsa_fusion.sAtt_1.bias = raw_model.tsa_fusion.sAtt_1.bias
#print(N_in * 128)
#print(raw_model.tsa_fusion.fea_fusion.weight[:, 0:N_in * 128, :, :].shape)
print("MODEL TSA SHAPE: ", model.tsa_fusion.fea_fusion.weight.shape)
model.tsa_fusion.maxpool = raw_model.tsa_fusion.maxpool
model.tsa_fusion.avgpool = raw_model.tsa_fusion.avgpool
model.tsa_fusion.sAtt_2 = raw_model.tsa_fusion.sAtt_2
model.tsa_fusion.sAtt_3 = raw_model.tsa_fusion.sAtt_3
model.tsa_fusion.sAtt_4 = raw_model.tsa_fusion.sAtt_4
model.tsa_fusion.sAtt_5 = raw_model.tsa_fusion.sAtt_5
model.tsa_fusion.sAtt_L1 = raw_model.tsa_fusion.sAtt_L1
model.tsa_fusion.sAtt_L2 = raw_model.tsa_fusion.sAtt_L2
model.tsa_fusion.sAtt_L3 = raw_model.tsa_fusion.sAtt_L3
model.tsa_fusion.sAtt_add_1 = raw_model.tsa_fusion.sAtt_add_1
model.tsa_fusion.sAtt_add_2 = raw_model.tsa_fusion.sAtt_add_2
model.tsa_fusion.lrelu = raw_model.tsa_fusion.lrelu
#if model.w_TSA:
# model.tsa_fusion = raw_model.tsa_fusion[:][:128 * N_in][:][:] #TSA_Fusion(nf=nf, nframes=nframes, center=self.center)
#else:
# model.tsa_fusion = raw_model.tsa_fusion[:][:128 * N_in][:][:] #nn.Conv2d(nframes * nf, nf, 1, 1, bias=True)
# print(self.tsa_fusion)
#### reconstruction
model.recon_trunk = raw_model.recon_trunk # arch_util.make_layer(ResidualBlock_noBN_f, back_RBs)
#### upsampling
model.upconv1 = raw_model.upconv1 #nn.Conv2d(nf, nf * 4, 3, 1, 1, bias=True)
model.upconv2 = raw_model.upconv2 #nn.Conv2d(nf, 64 * 4, 3, 1, 1, bias=True)
model.pixel_shuffle = raw_model.pixel_shuffle # nn.PixelShuffle(2)
model.HRconv = raw_model.HRconv
model.conv_last = raw_model.conv_last
#### activation function
model.lrelu = raw_model.lrelu
#####################################################
model.eval()
model = model.to(device)
avg_ssim_l, avg_ssim_center_l, avg_ssim_border_l = [], [], []
subfolder_name_l = []
subfolder_l = sorted(glob.glob(osp.join(test_dataset_folder, '*')))
subfolder_GT_l = sorted(glob.glob(osp.join(GT_dataset_folder, '*')))
# for each subfolder
for subfolder, subfolder_GT in zip(subfolder_l, subfolder_GT_l):
subfolder_name = osp.basename(subfolder)
subfolder_name_l.append(subfolder_name)
save_subfolder = osp.join(save_folder, subfolder_name)
img_path_l = sorted(glob.glob(osp.join(subfolder, '*')))
max_idx = len(img_path_l)
print("MAX_IDX: ", max_idx)
if save_imgs:
util.mkdirs(save_subfolder)
#### read LQ and GT images
imgs_LQ = data_util.read_img_seq(subfolder)
img_GT_l = []
for img_GT_path in sorted(glob.glob(osp.join(subfolder_GT, '*'))):
img_GT_l.append(data_util.read_img(None, img_GT_path))
avg_ssim, avg_ssim_border, avg_ssim_center, N_border, N_center = 0, 0, 0, 0, 0
# process each image
for img_idx, img_path in enumerate(img_path_l):
img_name = osp.splitext(osp.basename(img_path))[0]
if data_mode == "blur":
select_idx = data_util.glarefree_index_generation(
img_idx, max_idx, N_in, padding=padding)
else:
select_idx = data_util.index_generation(
img_idx, max_idx, N_in, padding=padding) # HERE GOTCHA
print("SELECT IDX: ", select_idx)
imgs_in = imgs_LQ.index_select(
0, torch.LongTensor(select_idx)).unsqueeze(0).to(device)
if flip_test:
output = util.flipx4_forward(model, imgs_in)
else:
print("IMGS_IN SHAPE: ", imgs_in.shape) # check this
output = util.single_forward(model, imgs_in) # error here 1
output = util.tensor2img(output.squeeze(0))
if save_imgs:
cv2.imwrite(
osp.join(save_subfolder, '{}.png'.format(img_name)),
output)
# calculate SSIM
output = output / 255.
GT = np.copy(img_GT_l[img_idx])
# For REDS, evaluate on RGB channels; for Vid4, evaluate on the Y channel
if data_mode == 'Vid4': # bgr2y, [0, 1]
GT = data_util.bgr2ycbcr(GT, only_y=True)
output = data_util.bgr2ycbcr(output, only_y=True)
output, GT = util.crop_border([output, GT], crop_border)
crt_ssim = util.calculate_ssim(output * 255, GT * 255)
logger.info('{:3d} - {:25} \tSSIM: {:.6f} dB'.format(
img_idx + 1, img_name, crt_ssim))
if img_idx >= border_frame and img_idx < max_idx - border_frame: # center frames
avg_ssim_center += crt_ssim
N_center += 1
else: # border frames
avg_ssim_border += crt_ssim
N_border += 1
avg_ssim = (avg_ssim_center + avg_ssim_border) / (N_center + N_border)
avg_ssim_center = avg_ssim_center / N_center
avg_ssim_border = 0 if N_border == 0 else avg_ssim_border / N_border
avg_ssim_l.append(avg_ssim)
avg_ssim_center_l.append(avg_ssim_center)
avg_ssim_border_l.append(avg_ssim_border)
logger.info('Folder {} - Average SSIM: {:.6f} dB for {} frames; '
'Center SSIM: {:.6f} dB for {} frames; '
'Border SSIM: {:.6f} dB for {} frames.'.format(
subfolder_name, avg_ssim, (N_center + N_border),
avg_ssim_center, N_center, avg_ssim_border, N_border))
logger.info('################ Tidy Outputs ################')
for subfolder_name, ssim, ssim_center, ssim_border in zip(
subfolder_name_l, avg_ssim_l, avg_ssim_center_l,
avg_ssim_border_l):
logger.info('Folder {} - Average SSIM: {:.6f} dB. '
'Center SSIM: {:.6f} dB. '
'Border SSIM: {:.6f} dB.'.format(subfolder_name, ssim,
ssim_center, ssim_border))
logger.info('################ Final Results ################')
logger.info('Data: {} - {}'.format(data_mode, test_dataset_folder))
logger.info('Padding mode: {}'.format(padding))
logger.info('Model path: {}'.format(model_path))
logger.info('Save images: {}'.format(save_imgs))
logger.info('Flip test: {}'.format(flip_test))
logger.info('Total Average SSIM: {:.6f} dB for {} clips. '
'Center SSIM: {:.6f} dB. Border SSIM: {:.6f} dB.'.format(
sum(avg_ssim_l) / len(avg_ssim_l), len(subfolder_l),
sum(avg_ssim_center_l) / len(avg_ssim_center_l),
sum(avg_ssim_border_l) / len(avg_ssim_border_l)))
def create_test_png(model_path, device, gpu_id, opt, subfolder_l, save_folder, save_imgs,
frame_notation, N_in, PAD, flip_test, end, total_run_time, logger, padding):
model = EDVR_arch.EDVR(nf=opt['network_G']['nf'], nframes=opt['network_G']['nframes'],
groups=opt['network_G']['groups'], front_RBs=opt['network_G']['front_RBs'],
back_RBs=opt['network_G']['back_RBs'],
predeblur=opt['network_G']['predeblur'], HR_in=opt['network_G']['HR_in'],
w_TSA=opt['network_G']['w_TSA'])
model.load_state_dict(torch.load(model_path), strict=True)
model.eval()
model = model.to(device)
#if (torch.cuda.is_available()):
model = model.cuda(gpu_id)
for subfolder in subfolder_l:
input_subfolder = os.path.split(subfolder)[1]
# subfolder_GT = os.path.join(GT_dataset_folder,input_subfolder)
#if not os.path.exists(subfolder_GT):
# continue
print("Evaluate Folders: ", input_subfolder)
subfolder_name = osp.basename(subfolder)
#subfolder_name_l.append(subfolder_name)
save_subfolder = osp.join(save_folder, subfolder_name)
img_path_l = sorted(glob.glob(osp.join(subfolder, '*')))
max_idx = len(img_path_l)
if save_imgs:
util.mkdirs(save_subfolder)
#### read LQ and GT images
imgs_LQ = data_util.read_img_seq(subfolder) # Num x 3 x H x W
#img_GT_l = []
#for img_GT_path in sorted(glob.glob(osp.join(subfolder_GT, '*'))):
# img_GT_l.append(data_util.read_img(None, img_GT_path))
#avg_psnr, avg_psnr_border, avg_psnr_center, N_border, N_center = 0, 0, 0, 0, 0
# process each image
for img_idx, img_path in enumerate(img_path_l):
img_name = osp.splitext(osp.basename(img_path))[0]
# todo here handle screen change
select_idx, log1, log2, nota = data_util.index_generation_process_screen_change_withlog_fixbug(input_subfolder, frame_notation, img_idx, max_idx, N_in, padding=padding)
if not log1 == None:
logger.info('screen change')
logger.info(nota)
logger.info(log1)
logger.info(log2)
imgs_in = imgs_LQ.index_select(0, torch.LongTensor(select_idx)).unsqueeze(0).cuda(gpu_id) # 960 x 540
# here we split the input images 960x540 into 9 320x180 patch
gtWidth = 3840
gtHeight = 2160
intWidth_ori = imgs_in.shape[4] # 960
intHeight_ori = imgs_in.shape[3] # 540
split_lengthY = 180
split_lengthX = 320
scale = 4
intPaddingRight_ = int(float(intWidth_ori) / split_lengthX + 1) * split_lengthX - intWidth_ori
intPaddingBottom_ = int(float(intHeight_ori) / split_lengthY + 1) * split_lengthY - intHeight_ori
intPaddingRight_ = 0 if intPaddingRight_ == split_lengthX else intPaddingRight_
intPaddingBottom_ = 0 if intPaddingBottom_ == split_lengthY else intPaddingBottom_
pader0 = torch.nn.ReplicationPad2d([0, intPaddingRight_, 0, intPaddingBottom_])
print("Init pad right/bottom " + str(intPaddingRight_) + " / " + str(intPaddingBottom_))
intPaddingRight = PAD # 32# 64# 128# 256
intPaddingLeft = PAD # 32#64 #128# 256
intPaddingTop = PAD # 32#64 #128#256
intPaddingBottom = PAD # 32#64 # 128# 256
pader = torch.nn.ReplicationPad2d([intPaddingLeft, intPaddingRight, intPaddingTop, intPaddingBottom])
imgs_in = torch.squeeze(imgs_in, 0)# N C H W
imgs_in = pader0(imgs_in) # N C 540 960
imgs_in = pader(imgs_in) # N C 604 1024
assert (split_lengthY == int(split_lengthY) and split_lengthX == int(split_lengthX))
split_lengthY = int(split_lengthY)
split_lengthX = int(split_lengthX)
split_numY = int(float(intHeight_ori) / split_lengthY )
split_numX = int(float(intWidth_ori) / split_lengthX)
splitsY = range(0, split_numY)
splitsX = range(0, split_numX)
intWidth = split_lengthX
intWidth_pad = intWidth + intPaddingLeft + intPaddingRight
intHeight = split_lengthY
intHeight_pad = intHeight + intPaddingTop + intPaddingBottom
# print("split " + str(split_numY) + ' , ' + str(split_numX))
y_all = np.zeros((gtHeight, gtWidth, 3), dtype="float32") # HWC
for split_j, split_i in itertools.product(splitsY, splitsX):
# print(str(split_j) + ", \t " + str(split_i))
X0 = imgs_in[:, :,
split_j * split_lengthY:(split_j + 1) * split_lengthY + intPaddingBottom + intPaddingTop,
split_i * split_lengthX:(split_i + 1) * split_lengthX + intPaddingRight + intPaddingLeft]
# y_ = torch.FloatTensor()
X0 = torch.unsqueeze(X0, 0) # N C H W -> 1 N C H W
#X0 = X0.cuda(gpu_id)
if flip_test:
output = util.flipx4_forward(model, X0)
else:
output = util.single_forward(model, X0)
output_depadded = output[0, :, intPaddingTop * scale :(intPaddingTop+intHeight) * scale, intPaddingLeft * scale: (intPaddingLeft+intWidth)*scale]
output_depadded = output_depadded.squeeze(0)
output = util.tensor2img(output_depadded)
y_all[split_j * split_lengthY * scale :(split_j + 1) * split_lengthY * scale,
split_i * split_lengthX * scale :(split_i + 1) * split_lengthX * scale, :] = \
np.round(output).astype(np.uint8)
# plt.figure(0)
# plt.title("pic")
# plt.imshow(y_all)
if save_imgs:
cv2.imwrite(osp.join(save_subfolder, '{}.png'.format(img_name)), y_all)
print("*****************current image process time \t " + str(
time.time() - end) + "s ******************")
total_run_time.update(time.time() - end, 1)
# calculate PSNR
#y_all = y_all / 255.
#GT = np.copy(img_GT_l[img_idx])
# For REDS, evaluate on RGB channels; for Vid4, evaluate on the Y channel
#if data_mode == 'Vid4': # bgr2y, [0, 1]
# GT = data_util.bgr2ycbcr(GT, only_y=True)
# y_all = data_util.bgr2ycbcr(y_all, only_y=True)
#y_all, GT = util.crop_border([y_all, GT], crop_border)
#crt_psnr = util.calculate_psnr(y_all * 255, GT * 255)
#logger.info('{:3d} - {:25} \tPSNR: {:.6f} dB'.format(img_idx + 1, img_name, crt_psnr))
logger.info('{} : {:3d} - {:25} \t'.format(input_subfolder, img_idx + 1, img_name))
def main():
#################
# configurations
#################
parser = argparse.ArgumentParser()
parser.add_argument("--input_path", type=str, required=True)
# parser.add_argument("--gt_path", type=str, required=True)
parser.add_argument("--output_path", type=str, required=True)
parser.add_argument("--model_path", type=str, required=True)
parser.add_argument("--gpu_id", type=str, required=True)
parser.add_argument("--gpu_number", type=str, required=True)
parser.add_argument("--gpu_index", type=str, required=True)
parser.add_argument("--screen_notation", type=str, required=True)
parser.add_argument('--opt',
type=str,
required=True,
help='Path to option YAML file.')
args = parser.parse_args()
opt = option.parse(args.opt, is_train=False)
gpu_number = int(args.gpu_number)
gpu_index = int(args.gpu_index)
PAD = 32
total_run_time = AverageMeter()
# print("GPU ", torch.cuda.device_count())
device = torch.device('cuda')
# os.environ['CUDA_VISIBLE_DEVICES'] = '0'
os.environ['CUDA_VISIBLE_DEVICES'] = str(args.gpu_id)
print('export CUDA_VISIBLE_DEVICES=' + str(args.gpu_id))
data_mode = 'sharp_bicubic' # Vid4 | sharp_bicubic | blur_bicubic | blur | blur_comp
# Vid4: SR
# REDS4: sharp_bicubic (SR-clean), blur_bicubic (SR-blur);
# blur (deblur-clean), blur_comp (deblur-compression).
stage = 1 # 1 or 2, use two stage strategy for REDS dataset.
flip_test = False
# Input_folder = "/DATA7_DB7/data/4khdr/data/Dataset/train_sharp_bicubic"
# GT_folder = "/DATA7_DB7/data/4khdr/data/Dataset/train_4k"
# Result_folder = "/DATA7_DB7/data/4khdr/data/Results"
Input_folder = args.input_path
# GT_folder = args.gt_path
Result_folder = args.output_path
Model_path = args.model_path
# create results folder
if not os.path.exists(Result_folder):
os.makedirs(Result_folder, exist_ok=True)
############################################################################
#### model
# if data_mode == 'Vid4':
# if stage == 1:
# model_path = '../experiments/pretrained_models/EDVR_Vimeo90K_SR_L.pth'
# else:
# raise ValueError('Vid4 does not support stage 2.')
# elif data_mode == 'sharp_bicubic':
# if stage == 1:
# # model_path = '../experiments/pretrained_models/EDVR_REDS_SR_L.pth'
# else:
# model_path = '../experiments/pretrained_models/EDVR_REDS_SR_Stage2.pth'
# elif data_mode == 'blur_bicubic':
# if stage == 1:
# model_path = '../experiments/pretrained_models/EDVR_REDS_SRblur_L.pth'
# else:
# model_path = '../experiments/pretrained_models/EDVR_REDS_SRblur_Stage2.pth'
# elif data_mode == 'blur':
# if stage == 1:
# model_path = '../experiments/pretrained_models/EDVR_REDS_deblur_L.pth'
# else:
# model_path = '../experiments/pretrained_models/EDVR_REDS_deblur_Stage2.pth'
# elif data_mode == 'blur_comp':
# if stage == 1:
# model_path = '../experiments/pretrained_models/EDVR_REDS_deblurcomp_L.pth'
# else:
# model_path = '../experiments/pretrained_models/EDVR_REDS_deblurcomp_Stage2.pth'
# else:
# raise NotImplementedError
model_path = Model_path
if data_mode == 'Vid4':
N_in = 7 # use N_in images to restore one HR image
else:
N_in = 5
predeblur, HR_in = False, False
back_RBs = 40
if data_mode == 'blur_bicubic':
predeblur = True
if data_mode == 'blur' or data_mode == 'blur_comp':
predeblur, HR_in = True, True
if stage == 2:
HR_in = True
back_RBs = 20
model = EDVR_arch.EDVR(nf=opt['network_G']['nf'],
nframes=opt['network_G']['nframes'],
groups=opt['network_G']['groups'],
front_RBs=opt['network_G']['front_RBs'],
back_RBs=opt['network_G']['back_RBs'],
predeblur=opt['network_G']['predeblur'],
HR_in=opt['network_G']['HR_in'],
w_TSA=opt['network_G']['w_TSA'])
# model = EDVR_arch.EDVR(128, N_in, 8, 5, back_RBs, predeblur=predeblur, HR_in=HR_in)
#### dataset
if data_mode == 'Vid4':
test_dataset_folder = '../datasets/Vid4/BIx4'
GT_dataset_folder = '../datasets/Vid4/GT'
else:
if stage == 1:
# test_dataset_folder = '../datasets/REDS4/{}'.format(data_mode)
# test_dataset_folder = '/DATA/wangshen_data/REDS/val_sharp_bicubic/X4'
test_dataset_folder = Input_folder
else:
test_dataset_folder = '../results/REDS-EDVR_REDS_SR_L_flipx4'
print('You should modify the test_dataset_folder path for stage 2')
# GT_dataset_folder = '../datasets/REDS4/GT'
# GT_dataset_folder = '/DATA/wangshen_data/REDS/val_sharp'
# GT_dataset_folder = GT_folder
#### evaluation
crop_border = 0
border_frame = N_in // 2 # border frames when evaluate
# temporal padding mode
if data_mode == 'Vid4' or data_mode == 'sharp_bicubic':
padding = 'new_info'
else:
padding = 'replicate'
save_imgs = True
# save_folder = '../results/{}'.format(data_mode)
# save_folder = '/DATA/wangshen_data/REDS/results/{}'.format(data_mode)
save_folder = os.path.join(Result_folder, data_mode)
util.mkdirs(save_folder)
util.setup_logger('base',
save_folder,
'test',
level=logging.INFO,
screen=True,
tofile=True)
logger = logging.getLogger('base')
#### log info
logger.info('Data: {} - {}'.format(data_mode, test_dataset_folder))
logger.info('Padding mode: {}'.format(padding))
logger.info('Model path: {}'.format(model_path))
logger.info('Save images: {}'.format(save_imgs))
logger.info('Flip test: {}'.format(flip_test))
#### set up the models
model.load_state_dict(torch.load(model_path), strict=True)
model.eval()
model = model.to(device)
avg_psnr_l, avg_psnr_center_l, avg_psnr_border_l = [], [], []
subfolder_name_l = []
subfolder_l = sorted(glob.glob(osp.join(test_dataset_folder, '*')))
# subfolder_GT_l = sorted(glob.glob(osp.join(GT_dataset_folder, '*')))
# for each subfolder
# for subfolder, subfolder_GT in zip(subfolder_l, subfolder_GT_l):
end = time.time()
# load screen change notation
import json
with open(args.screen_notation) as f:
frame_notation = json.load(f)
subfolder_n = len(subfolder_l)
subfolder_l = subfolder_l[int(subfolder_n * gpu_index /
gpu_number):int(subfolder_n *
(gpu_index + 1) /
gpu_number)]
for subfolder in subfolder_l:
input_subfolder = os.path.split(subfolder)[1]
# subfolder_GT = os.path.join(GT_dataset_folder,input_subfolder)
#if not os.path.exists(subfolder_GT):
# continue
print("Evaluate Folders: ", input_subfolder)
subfolder_name = osp.basename(subfolder)
subfolder_name_l.append(subfolder_name)
save_subfolder = osp.join(save_folder, subfolder_name)
img_path_l = sorted(glob.glob(osp.join(subfolder, '*')))
max_idx = len(img_path_l)
if save_imgs:
util.mkdirs(save_subfolder)
#### read LQ and GT images
imgs_LQ = data_util.read_img_seq(subfolder) # Num x 3 x H x W
#img_GT_l = []
#for img_GT_path in sorted(glob.glob(osp.join(subfolder_GT, '*'))):
# img_GT_l.append(data_util.read_img(None, img_GT_path))
#avg_psnr, avg_psnr_border, avg_psnr_center, N_border, N_center = 0, 0, 0, 0, 0
# process each image
for img_idx, img_path in enumerate(img_path_l):
img_name = osp.splitext(osp.basename(img_path))[0]
# todo here handle screen change
select_idx, log1, log2, nota = data_util.index_generation_process_screen_change_withlog_fixbug(
input_subfolder,
frame_notation,
img_idx,
max_idx,
N_in,
padding=padding)
if not log1 == None:
logger.info('screen change')
logger.info(nota)
logger.info(log1)
logger.info(log2)
imgs_in = imgs_LQ.index_select(
0, torch.LongTensor(select_idx)).unsqueeze(0).to(
device) # 960 x 540
# here we split the input images 960x540 into 9 320x180 patch
gtWidth = 3840
gtHeight = 2160
intWidth_ori = imgs_in.shape[4] # 960
intHeight_ori = imgs_in.shape[3] # 540
split_lengthY = 180
split_lengthX = 320
scale = 4
intPaddingRight_ = int(float(intWidth_ori) / split_lengthX +
1) * split_lengthX - intWidth_ori
intPaddingBottom_ = int(float(intHeight_ori) / split_lengthY +
1) * split_lengthY - intHeight_ori
intPaddingRight_ = 0 if intPaddingRight_ == split_lengthX else intPaddingRight_
intPaddingBottom_ = 0 if intPaddingBottom_ == split_lengthY else intPaddingBottom_
pader0 = torch.nn.ReplicationPad2d(
[0, intPaddingRight_, 0, intPaddingBottom_])
print("Init pad right/bottom " + str(intPaddingRight_) + " / " +
str(intPaddingBottom_))
intPaddingRight = PAD # 32# 64# 128# 256
intPaddingLeft = PAD # 32#64 #128# 256
intPaddingTop = PAD # 32#64 #128#256
intPaddingBottom = PAD # 32#64 # 128# 256
pader = torch.nn.ReplicationPad2d([
intPaddingLeft, intPaddingRight, intPaddingTop,
intPaddingBottom
])
imgs_in = torch.squeeze(imgs_in, 0) # N C H W
imgs_in = pader0(imgs_in) # N C 540 960
imgs_in = pader(imgs_in) # N C 604 1024
assert (split_lengthY == int(split_lengthY)
and split_lengthX == int(split_lengthX))
split_lengthY = int(split_lengthY)
split_lengthX = int(split_lengthX)
split_numY = int(float(intHeight_ori) / split_lengthY)
split_numX = int(float(intWidth_ori) / split_lengthX)
splitsY = range(0, split_numY)
splitsX = range(0, split_numX)
intWidth = split_lengthX
intWidth_pad = intWidth + intPaddingLeft + intPaddingRight
intHeight = split_lengthY
intHeight_pad = intHeight + intPaddingTop + intPaddingBottom
# print("split " + str(split_numY) + ' , ' + str(split_numX))
y_all = np.zeros((gtHeight, gtWidth, 3), dtype="float32") # HWC
for split_j, split_i in itertools.product(splitsY, splitsX):
# print(str(split_j) + ", \t " + str(split_i))
X0 = imgs_in[:, :, split_j *
split_lengthY:(split_j + 1) * split_lengthY +
intPaddingBottom + intPaddingTop, split_i *
split_lengthX:(split_i + 1) * split_lengthX +
intPaddingRight + intPaddingLeft]
# y_ = torch.FloatTensor()
X0 = torch.unsqueeze(X0, 0) # N C H W -> 1 N C H W
if flip_test:
output = util.flipx4_forward(model, X0)
else:
output = util.single_forward(model, X0)
output_depadded = output[0, :, intPaddingTop *
scale:(intPaddingTop + intHeight) *
scale, intPaddingLeft *
scale:(intPaddingLeft + intWidth) *
scale]
output_depadded = output_depadded.squeeze(0)
output = util.tensor2img(output_depadded)
y_all[split_j * split_lengthY * scale :(split_j + 1) * split_lengthY * scale,
split_i * split_lengthX * scale :(split_i + 1) * split_lengthX * scale, :] = \
np.round(output).astype(np.uint8)
# plt.figure(0)
# plt.title("pic")
# plt.imshow(y_all)
if save_imgs:
cv2.imwrite(
osp.join(save_subfolder, '{}.png'.format(img_name)), y_all)
print("*****************current image process time \t " +
str(time.time() - end) + "s ******************")
total_run_time.update(time.time() - end, 1)
# calculate PSNR
#y_all = y_all / 255.
#GT = np.copy(img_GT_l[img_idx])
# For REDS, evaluate on RGB channels; for Vid4, evaluate on the Y channel
#if data_mode == 'Vid4': # bgr2y, [0, 1]
# GT = data_util.bgr2ycbcr(GT, only_y=True)
# y_all = data_util.bgr2ycbcr(y_all, only_y=True)
#y_all, GT = util.crop_border([y_all, GT], crop_border)
#crt_psnr = util.calculate_psnr(y_all * 255, GT * 255)
#logger.info('{:3d} - {:25} \tPSNR: {:.6f} dB'.format(img_idx + 1, img_name, crt_psnr))
logger.info('{} : {:3d} - {:25} \t'.format(input_subfolder,
img_idx + 1, img_name))
#if img_idx >= border_frame and img_idx < max_idx - border_frame: # center frames
# avg_psnr_center += crt_psnr
# N_center += 1
#else: # border frames
# avg_psnr_border += crt_psnr
# N_border += 1
#avg_psnr = (avg_psnr_center + avg_psnr_border) / (N_center + N_border)
#avg_psnr_center = avg_psnr_center / N_center
#avg_psnr_border = 0 if N_border == 0 else avg_psnr_border / N_border
#avg_psnr_l.append(avg_psnr)
#avg_psnr_center_l.append(avg_psnr_center)
#avg_psnr_border_l.append(avg_psnr_border)
#logger.info('Folder {} - Average PSNR: {:.6f} dB for {} frames; '
# 'Center PSNR: {:.6f} dB for {} frames; '
# 'Border PSNR: {:.6f} dB for {} frames.'.format(subfolder_name, avg_psnr,
# (N_center + N_border),
# avg_psnr_center, N_center,
# avg_psnr_border, N_border))
#logger.info('################ Tidy Outputs ################')
#for subfolder_name, psnr, psnr_center, psnr_border in zip(subfolder_name_l, avg_psnr_l,
# avg_psnr_center_l, avg_psnr_border_l):
# logger.info('Folder {} - Average PSNR: {:.6f} dB. '
# 'Center PSNR: {:.6f} dB. '
# 'Border PSNR: {:.6f} dB.'.format(subfolder_name, psnr, psnr_center,
# psnr_border))
#logger.info('################ Final Results ################')
logger.info('Data: {} - {}'.format(data_mode, test_dataset_folder))
logger.info('Padding mode: {}'.format(padding))
logger.info('Model path: {}'.format(model_path))
logger.info('Save images: {}'.format(save_imgs))
logger.info('Flip test: {}'.format(flip_test))
def __init__(self, opt, **kwargs):
super(VideoTestDataset, self).__init__()
self.scale = kwargs['scale']
idx = kwargs['idx'] if 'idx' in kwargs else None
self.opt = opt
self.cache_data = opt['cache_data']
self.half_N_frames = opt['N_frames'] // 2
if idx is None:
self.name = opt['name']
self.GT_root, self.LQ_root = opt['dataroot_GT'], opt['dataroot_LQ']
else:
self.name = opt['name'].split('+')[idx]
self.GT_root, self.LQ_root = opt['dataroot_GT'].split(
'+')[idx], opt['dataroot_LQ'].split('+')[idx]
self.data_type = self.opt['data_type']
self.data_info = {
'path_LQ': [],
'path_GT': [],
'folder': [],
'idx': [],
'border': []
}
if self.data_type == 'lmdb':
raise ValueError('No need to use LMDB during validation/test.')
#### Generate data info and cache data
self.imgs_LQ, self.imgs_GT = {}, {}
if self.name.lower() in ['vid4', 'reds', 'mm522']:
if self.name.lower() == 'vid4':
img_type = 'img'
subfolders_LQ = util.glob_file_list(
osp.join(self.LQ_root, 'X{}'.format(self.scale)))
subfolders_GT = util.glob_file_list(self.GT_root)
elif self.name.lower() == 'reds':
img_type = 'img'
list_hr_seq = util.glob_file_list(self.GT_root)
list_lr_seq = util.glob_file_list(
osp.join(self.LQ_root, 'X{}'.format(self.scale)))
subfolders_GT = [
k for k in list_hr_seq
if k.find('000') >= 0 or k.find('011') >= 0
or k.find('015') >= 0 or k.find('020') >= 0
]
subfolders_LQ = [
k for k in list_lr_seq
if k.find('000') >= 0 or k.find('011') >= 0
or k.find('015') >= 0 or k.find('020') >= 0
]
else:
img_type = 'img'
subfolders_LQ = util.glob_file_list(
osp.join(self.LQ_root, 'X{}'.format(self.scale)))
subfolders_GT = util.glob_file_list(self.GT_root)
for subfolder_LQ, subfolder_GT in zip(subfolders_LQ,
subfolders_GT):
subfolder_name = osp.basename(subfolder_GT)
img_paths_LQ = util.glob_file_list(subfolder_LQ)
img_paths_GT = util.glob_file_list(subfolder_GT)
max_idx = len(img_paths_LQ)
assert max_idx == len(
img_paths_GT
), 'Different number of images in LQ and GT folders'
self.data_info['path_LQ'].extend(img_paths_LQ)
self.data_info['path_GT'].extend(img_paths_GT)
self.data_info['folder'].extend([subfolder_name] * max_idx)
for i in range(max_idx):
self.data_info['idx'].append('{}/{}'.format(i, max_idx))
border_l = [0] * max_idx
for i in range(self.half_N_frames):
border_l[i] = 1
border_l[max_idx - i - 1] = 1
self.data_info['border'].extend(border_l)
if self.cache_data:
self.imgs_LQ[subfolder_name] = util.read_img_seq(
img_paths_LQ, img_type)
self.imgs_GT[subfolder_name] = util.read_img_seq(
img_paths_GT, img_type)
elif opt['name'].lower() in ['vimeo90k-test']:
pass # TODO
else:
raise ValueError(
'Not support video test dataset. Support Vid4, REDS4 and Vimeo90k-Test.'
)