def __getitem__(self, 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]
lq_path = self.data_info['lq_path'][index]
select_idx = generate_frame_indices(idx,
max_idx,
self.opt['num_frame'],
padding=self.opt['padding'])
if self.cache_data:
imgs_lq = self.imgs_lq[folder].index_select(
0, torch.LongTensor(select_idx))
img_gt = self.imgs_gt[folder][idx]
else:
img_paths_lq = [self.imgs_lq[folder][i] for i in select_idx]
imgs_lq = read_img_seq(img_paths_lq)
img_gt = read_img_seq([self.imgs_gt[folder][idx]])
img_gt.squeeze_(0)
return {
'lq': imgs_lq, # (t, c, h, w)
'gt': img_gt, # (c, h, w)
'folder': folder, # folder name
'idx': self.data_info['idx'][index], # e.g., 0/99
'border': border, # 1 for border, 0 for non-border
'lq_path': lq_path # center frame
}
def __getitem__(self, 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]
lq_path = self.data_info['lq_path'][index]
select_idx = generate_frame_indices(idx,
max_idx,
self.opt['num_frame'],
padding=self.opt['padding'])
if self.cache_data:
if self.opt['use_duf_downsampling']:
# read imgs_gt to generate low-resolution frames
imgs_lq = self.imgs_gt[folder].index_select(
0, torch.LongTensor(select_idx))
imgs_lq = duf_downsample(imgs_lq,
kernel_size=13,
scale=self.opt['scale'])
else:
imgs_lq = self.imgs_lq[folder].index_select(
0, torch.LongTensor(select_idx))
img_gt = self.imgs_gt[folder][idx]
else:
if self.opt['use_duf_downsampling']:
img_paths_lq = [self.imgs_gt[folder][i] for i in select_idx]
# read imgs_gt to generate low-resolution frames
imgs_lq = read_img_seq(img_paths_lq,
require_mod_crop=True,
scale=self.opt['scale'])
imgs_lq = duf_downsample(imgs_lq,
kernel_size=13,
scale=self.opt['scale'])
else:
img_paths_lq = [self.imgs_lq[folder][i] for i in select_idx]
imgs_lq = read_img_seq(img_paths_lq)
img_gt = read_img_seq([self.imgs_gt[folder][idx]],
require_mod_crop=True,
scale=self.opt['scale'])
img_gt.squeeze_(0)
return {
'lq': imgs_lq, # (t, c, h, w)
'gt': img_gt, # (c, h, w)
'folder': folder, # folder name
'idx': self.data_info['idx'][index], # e.g., 0/99
'border': border, # 1 for border, 0 for non-border
'lq_path': lq_path # center frame
}
def __getitem__(self, index):
lq_path = self.data_info['lq_path'][index]
gt_path = self.data_info['gt_path'][index]
imgs_lq = read_img_seq(lq_path)
img_gt = read_img_seq([gt_path])
img_gt.squeeze_(0)
return {
'lq': imgs_lq, # (t, c, h, w)
'gt': img_gt, # (c, h, w)
'folder': self.data_info['folder'][index], # folder name
'idx': self.data_info['idx'][index], # e.g., 0/843
'border': self.data_info['border'][index], # 0 for non-border
'lq_path': lq_path[self.opt['num_frame'] // 2] # center frame
}
def __getitem__(self, index):
clip_name = self.data_info['clip_name'][index]
max_idx = self.data_info['max_idx'][index]
select_idx = range(int(max_idx))
lq_frames_path = [self.lq_frames[clip_name][i] for i in select_idx]
gt_frames_path = [self.gt_frames[clip_name][i] for i in select_idx]
frame_list = list(range(len(lq_frames_path)))
lq_frames = read_img_seq(lq_frames_path)
gt_frames = read_img_seq(gt_frames_path)
return {
'lq': lq_frames, # (t, c, h, w)
'gt': gt_frames, # (t, c, h, w)
'clip_name': clip_name,
'frame_list': frame_list
}
def __init__(self, opt):
super(VideoTestDataset, 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 = {
'lq_path': [],
'gt_path': [],
'folder': [],
'idx': [],
'border': []
}
# file client (io backend)
self.file_client = None
self.io_backend_opt = opt['io_backend']
assert self.io_backend_opt[
'type'] != 'lmdb', 'No need to use lmdb during validation/test.'
logger = get_root_logger()
logger.info(f'Generate data info for VideoTestDataset - {opt["name"]}')
self.imgs_lq, self.imgs_gt = {}, {}
if 'meta_info_file' in opt:
with open(opt['meta_info_file'], 'r') as fin:
subfolders = [line.split(' ')[0] for line in fin]
subfolders_lq = [
osp.join(self.lq_root, key) for key in subfolders
]
subfolders_gt = [
osp.join(self.gt_root, key) for key in subfolders
]
else:
subfolders_lq = sorted(glob.glob(osp.join(self.lq_root, '*')))
subfolders_gt = sorted(glob.glob(osp.join(self.gt_root, '*')))
if opt['name'].lower() in ['vid4', 'reds4', 'redsofficial']:
for subfolder_lq, subfolder_gt in zip(subfolders_lq,
subfolders_gt):
# get frame list for lq and gt
subfolder_name = osp.basename(subfolder_lq)
img_paths_lq = sorted(
list(scandir(subfolder_lq, full_path=True)))
img_paths_gt = sorted(
list(scandir(subfolder_gt, full_path=True)))
max_idx = len(img_paths_lq)
assert max_idx == len(img_paths_gt), (
f'Different number of images in lq ({max_idx})'
f' and gt folders ({len(img_paths_gt)})')
self.data_info['lq_path'].extend(img_paths_lq)
self.data_info['gt_path'].extend(img_paths_gt)
self.data_info['folder'].extend([subfolder_name] * max_idx)
for i in range(max_idx):
self.data_info['idx'].append(f'{i}/{max_idx}')
border_l = [0] * max_idx
for i in range(self.opt['num_frame'] // 2):
border_l[i] = 1
border_l[max_idx - i - 1] = 1
self.data_info['border'].extend(border_l)
# cache data or save the frame list
if self.cache_data:
logger.info(
f'Cache {subfolder_name} for VideoTestDataset...')
self.imgs_lq[subfolder_name] = read_img_seq(img_paths_lq)
self.imgs_gt[subfolder_name] = read_img_seq(img_paths_gt)
else:
self.imgs_lq[subfolder_name] = img_paths_lq
self.imgs_gt[subfolder_name] = img_paths_gt
else:
raise ValueError(
f'Non-supported video test dataset: {type(opt["name"])}')
def main():
parser = argparse.ArgumentParser()
parser.add_argument(
'--model_path',
type=str,
default='experiments/pretrained_models/BasicVSRPP_REDS4.pth')
parser.add_argument('--input_path',
type=str,
default='datasets/REDS4/sharp_bicubic/000',
help='input test image folder')
parser.add_argument('--save_path',
type=str,
default='results/BasicVSRPP/000',
help='save image path')
parser.add_argument('--interval',
type=int,
default=100,
help='interval size')
args = parser.parse_args()
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
# set up model
model = BasicVSRPlusPlus(mid_channels=64, num_blocks=7)
model.load_state_dict(torch.load(args.model_path)['params'], strict=True)
model.eval()
model = model.to(device)
os.makedirs(args.save_path, exist_ok=True)
# extract images from video format files
input_path = args.input_path
use_ffmpeg = False
if not os.path.isdir(input_path):
use_ffmpeg = True
video_name = os.path.splitext(os.path.split(args.input_path)[-1])[0]
input_path = os.path.join('./BasicVSRPP_tmp', video_name)
os.makedirs(os.path.join('./BasicVSRPP_tmp', video_name),
exist_ok=True)
os.system(
f'ffmpeg -i {args.input_path} -qscale:v 1 -qmin 1 -qmax 1 -vsync 0 {input_path} /frame%08d.png'
)
# load data and inference
imgs_list = sorted(glob.glob(os.path.join(input_path, '*')))
num_imgs = len(imgs_list)
if len(imgs_list
) <= args.interval: # too many images may cause CUDA out of memory
imgs, imgnames = read_img_seq(imgs_list, return_imgname=True)
imgs = imgs.unsqueeze(0).to(device)
inference(imgs, imgnames, model, args.save_path)
else:
for idx in range(0, num_imgs, args.interval):
interval = min(args.interval, num_imgs - idx)
imgs, imgnames = read_img_seq(imgs_list[idx:idx + interval],
return_imgname=True)
imgs = imgs.unsqueeze(0).to(device)
inference(imgs, imgnames, model, args.save_path)
# delete ffmpeg output images
if use_ffmpeg:
shutil.rmtree(input_path)