def main():
scale = 4
N_ot = args.N_out
N_in = 1 + N_ot // 2
os.environ['CUDA_VISIBLE_DEVICES'] = '0'
#### model
model_path = args.model
model = Sakuya_arch.LunaTokis(64, N_ot, 8, 5, 40)
#### extract the input video to temporary folder
save_folder = osp.join(osp.dirname(args.output), '.delme')
save_out_folder = osp.join(osp.dirname(args.output), '.hr_delme')
util.mkdirs(save_folder)
util.mkdirs(save_out_folder)
error = util.extract_frames(args.ffmpeg_dir, args.video, save_folder)
if error:
print(error)
exit(1)
# temporal padding mode
padding = 'replicate'
save_imgs = True
############################################################################
if torch.cuda.is_available():
device = torch.device('cuda')
else:
device = torch.device('cpu')
def single_forward(model, imgs_in):
with torch.no_grad():
# print(imgs_in.size()) # [1,5,3,270,480]
b, n, c, h, w = imgs_in.size()
h_n = int(4 * np.ceil(h / 4))
w_n = int(4 * np.ceil(w / 4))
imgs_temp = imgs_in.new_zeros(b, n, c, h_n, w_n)
imgs_temp[:, :, :, 0:h, 0:w] = imgs_in
model_output = model(imgs_temp)
model_output = model_output[:, :, :, 0:scale * h, 0:scale * w]
if isinstance(model_output, list) or isinstance(
model_output, tuple):
output = model_output[0]
else:
output = model_output
return output
model.load_state_dict(torch.load(model_path), strict=True)
model.eval()
model = model.to(device)
#### zsm images
imgs = util.read_seq_imgs(save_folder)
select_idx_list = util.test_index_generation(False, N_ot, len(imgs))
for select_idxs in select_idx_list:
# get input images
select_idx = select_idxs[0]
imgs_in = imgs.index_select(
0, torch.LongTensor(select_idx)).unsqueeze(0).to(device)
output = single_forward(model, imgs_in)
outputs = output.data.float().cpu().squeeze(0)
# save imgs
out_idx = select_idxs[1]
for idx, name_idx in enumerate(out_idx):
output_f = outputs[idx, ...].squeeze(0)
if save_imgs:
output = util.tensor2img(output_f)
cv2.imwrite(
osp.join(save_out_folder, '{:06d}.png'.format(name_idx)),
output)
# now turn output images to video
# generate mp4
util.create_video(args.ffmpeg_dir, save_out_folder, args.output, args.fps)
# remove tmp folder
rmtree(save_folder)
rmtree(save_out_folder)
exit(0)
def main():
scale = 4
N_ot = args.N_out
N_in = 1 + N_ot // 2
os.environ['CUDA_VISIBLE_DEVICES'] = '0'
# model
model_path = args.model
model = Sakuya_arch.VideoZoomModel(64, N_ot, 8, 5, 40)
# extract the input video to temporary folder
save_folder = 'dataset/input'
save_out_folder = 'dataset/output'
# util.mkdirs(save_folder)
# util.mkdirs(save_out_folder)
# error = util.extract_frames(args.ffmpeg_dir, args.video, save_folder)
# if error:
# print(error)
# exit(1)
# temporal padding mode
# padding = 'replicate'
save_imgs = True
############################################################################
if torch.cuda.is_available():
device = torch.device('cuda')
else:
device = torch.device('cpu')
# device = torch.device('cpu')
def single_forward(model, imgs_in):
with torch.no_grad():
# print(imgs_in.size()) # [1,5,3,270,480]
b, n, c, h, w = imgs_in.size()
h_n = int(4 * np.ceil(h / 4))
w_n = int(4 * np.ceil(w / 4))
imgs_temp = imgs_in.new_zeros(b, n, c, h_n, w_n)
imgs_temp[:, :, :, 0:h, 0:w] = imgs_in
del imgs_in
model_output = model(imgs_temp).cpu()
# pdb.set_trace() what's wrong ?
# torch.Size([1, 2, 3, 272, 480])
# (Pdb) model_output.size()
# torch.Size([1, 3, 3, 1088, 1920])
model_output = model_output[:, :, :, 0:scale * h, 0:scale * w]
# if isinstance(model_output, list) or isinstance(model_output, tuple):
# pdb.set_trace()
# output = model_output[0]
# else:
# pdb.set_trace()
# output = model_output
del imgs_temp
torch.cuda.empty_cache()
return model_output
model.load_state_dict(torch.load(model_path), strict=True)
model.eval()
model = model.to(device)
model = amp.initialize(model, opt_level='O1')
# zsm images
imgs = util.read_seq_imgs(save_folder)
select_idx_list = util.test_index_generation(False, N_ot, len(imgs))
# pdb.set_trace()
# (Pdb) type(select_idx_list),len(select_idx_list), select_idx_list[0], select_idx_list[1]
# (<class 'list'>, 99, [[0, 1], [0, 1, 2]], [[1, 2], [2, 3, 4]])
# (Pdb) imgs.size(), TxCxHxW
# torch.Size([100, 3, 270, 480])
for select_idxs in select_idx_list:
# get input images
print(select_idxs, " ...")
select_idx = select_idxs[0]
imgs_in = imgs.index_select(
0, torch.LongTensor(select_idx)).unsqueeze(0).to(device)
# (Pdb) imgs_in.size()
# torch.Size([1, 2, 3, 270, 480]), [B,T,C,H,W]
output = single_forward(model, imgs_in)
# (Pdb) output.size()
# torch.Size([1, 3, 3, 1080, 1920])
outputs = output.data.float().cpu().squeeze(0)
# save imgs
out_idx = select_idxs[1]
# pdb.set_trace()
for idx, name_idx in enumerate(out_idx):
output_f = outputs[idx, ...].squeeze(0)
if save_imgs:
output = util.tensor2img(output_f)
cv2.imwrite(
osp.join(save_out_folder, '{:06d}.png'.format(name_idx)),
output)
# now turn output images to video
# generate mp4
# util.create_video(args.ffmpeg_dir, save_out_folder, args.output, args.fps)
# remove tmp folder
# rmtree(save_folder)
# rmtree(save_out_folder)
exit(0)
def main():
scale = 4
N_ot = 5 # 3
N_in = 1 + N_ot // 2
os.environ['CUDA_VISIBLE_DEVICES'] = '0'
# model
# TODO: change your model path here
model_path = '../experiments/pretrained_models/xiang2020zooming.pth'
model = Sakuya_arch.LunaTokis(64, N_ot, 8, 5, 40)
# dataset
data_mode = 'Custom' # 'Vid4' #'SPMC'#'Middlebury'#
if data_mode == 'Vid4':
test_dataset_folder = '/data/xiang/SR/Vid4/LR/*'
if data_mode == 'SPMC':
test_dataset_folder = '/data/xiang/SR/spmc/*'
if data_mode == 'Custom':
test_dataset_folder = '../test_example/*' # TODO: put your own data path here
# evaluation
flip_test = False # True#
crop_border = 0
# temporal padding mode
padding = 'replicate'
save_imgs = False # True#
if 'Custom' in data_mode:
save_imgs = True
############################################################################
if torch.cuda.is_available():
device = torch.device('cuda')
else:
device = torch.device('cpu')
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')
model_params = util.get_model_total_params(model)
# 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('Model parameters: {} M'.format(model_params))
logger.info('Save images: {}'.format(save_imgs))
logger.info('Flip Test: {}'.format(flip_test))
def single_forward(model, imgs_in):
with torch.no_grad():
# imgs_in.size(): [1,n,3,h,w]
b, n, c, h, w = imgs_in.size()
h_n = int(4 * np.ceil(h / 4))
w_n = int(4 * np.ceil(w / 4))
imgs_temp = imgs_in.new_zeros(b, n, c, h_n, w_n)
imgs_temp[:, :, :, 0:h, 0:w] = imgs_in
t0 = time.time()
model_output = model(imgs_temp)
t1 = time.time()
logger.info('Single Test time: {}'.format(t1 - t0))
# model_output.size(): torch.Size([1, 3, 4h, 4w])
model_output = model_output[:, :, :, 0:scale * h, 0:scale * w]
if isinstance(model_output, list) or isinstance(
model_output, tuple):
output = model_output[0]
else:
output = model_output
return output
sub_folder_l = sorted(glob.glob(test_dataset_folder))
model.load_state_dict(torch.load(model_path), strict=True)
model.eval()
model = model.to(device)
avg_psnr_l = []
avg_psnr_y_l = []
sub_folder_name_l = []
# total_time = []
# for each sub-folder
for sub_folder in sub_folder_l:
gt_tested_list = []
sub_folder_name = sub_folder.split('/')[-1]
sub_folder_name_l.append(sub_folder_name)
save_sub_folder = osp.join(save_folder, sub_folder_name)
if data_mode == 'SPMC':
sub_folder = sub_folder + '/LR/'
img_LR_l = sorted(glob.glob(sub_folder + '/*'))
if save_imgs:
util.mkdirs(save_sub_folder)
# read LR images
imgs = util.read_seq_imgs(sub_folder)
# read GT images
img_GT_l = []
if data_mode == 'SPMC':
sub_folder_GT = osp.join(sub_folder.replace('/LR/', '/truth/'))
else:
sub_folder_GT = osp.join(sub_folder.replace('/LR/', '/HR/'))
if 'Custom' not in data_mode:
for img_GT_path in sorted(glob.glob(osp.join(sub_folder_GT, '*'))):
img_GT_l.append(util.read_image(img_GT_path))
avg_psnr, avg_psnr_sum, cal_n = 0, 0, 0
avg_psnr_y, avg_psnr_sum_y = 0, 0
if len(img_LR_l) == len(img_GT_l):
skip = True
else:
skip = False
if 'Custom' in data_mode:
select_idx_list = util.test_index_generation(
False, N_ot, len(img_LR_l))
else:
select_idx_list = util.test_index_generation(
skip, N_ot, len(img_LR_l))
# process each image
for select_idxs in select_idx_list:
# get input images
select_idx = select_idxs[0]
gt_idx = select_idxs[1]
imgs_in = imgs.index_select(
0, torch.LongTensor(select_idx)).unsqueeze(0).to(device)
output = single_forward(model, imgs_in)
print(select_idx, gt_idx, imgs_in.size(), output.size())
outputs = output.data.float().cpu().squeeze(0)
if flip_test:
# flip W
output = single_forward(model, torch.flip(imgs_in, (-1, )))
output = torch.flip(output, (-1, ))
output = output.data.float().cpu().squeeze(0)
outputs = outputs + output
# flip H
output = single_forward(model, torch.flip(imgs_in, (-2, )))
output = torch.flip(output, (-2, ))
output = output.data.float().cpu().squeeze(0)
outputs = outputs + output
# flip both H and W
output = single_forward(model, torch.flip(imgs_in, (-2, -1)))
output = torch.flip(output, (-2, -1))
output = output.data.float().cpu().squeeze(0)
outputs = outputs + output
outputs = outputs / 4
# save imgs
for idx, name_idx in enumerate(gt_idx):
if name_idx in gt_tested_list:
continue
gt_tested_list.append(name_idx)
output_f = outputs[idx, :, :, :].squeeze(0)
output = util.tensor2img(output_f)
if save_imgs:
cv2.imwrite(
osp.join(save_sub_folder,
'{:08d}.png'.format(name_idx + 1)), output)
if 'Custom' not in data_mode:
# calculate PSNR
output = output / 255.
GT = np.copy(img_GT_l[name_idx])
if crop_border == 0:
cropped_output = output
cropped_GT = GT
else:
cropped_output = output[crop_border:-crop_border,
crop_border:-crop_border, :]
cropped_GT = GT[crop_border:-crop_border,
crop_border:-crop_border, :]
crt_psnr = util.calculate_psnr(cropped_output * 255,
cropped_GT * 255)
cropped_GT_y = data_util.bgr2ycbcr(cropped_GT, only_y=True)
cropped_output_y = data_util.bgr2ycbcr(cropped_output,
only_y=True)
crt_psnr_y = util.calculate_psnr(cropped_output_y * 255,
cropped_GT_y * 255)
logger.info(
'{:3d} - {:25}.png \tPSNR: {:.6f} dB PSNR-Y: {:.6f} dB'
.format(name_idx + 1, name_idx + 1, crt_psnr,
crt_psnr_y))
avg_psnr_sum += crt_psnr
avg_psnr_sum_y += crt_psnr_y
cal_n += 1
if 'Custom' not in data_mode:
avg_psnr = avg_psnr_sum / cal_n
avg_psnr_y = avg_psnr_sum_y / cal_n
logger.info(
'Folder {} - Average PSNR: {:.6f} dB PSNR-Y: {:.6f} dB for {} frames; '
.format(sub_folder_name, avg_psnr, avg_psnr_y, cal_n))
avg_psnr_l.append(avg_psnr)
avg_psnr_y_l.append(avg_psnr_y)
if 'Custom' not in data_mode:
logger.info('################ Tidy Outputs ################')
for name, psnr, psnr_y in zip(sub_folder_name_l, avg_psnr_l,
avg_psnr_y_l):
logger.info(
'Folder {} - Average PSNR: {:.6f} dB PSNR-Y: {:.6f} dB. '.
format(name, psnr, psnr_y))
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 PSNR-Y: {:.6f} dB for {} clips. '.
format(
sum(avg_psnr_l) / len(avg_psnr_l),
sum(avg_psnr_y_l) / len(avg_psnr_y_l), len(sub_folder_l)))
def main():
counter = 0
time_values = []
temp_count = 0
# Deleting old files, if they exist
file_extention = os.path.splitext(args.video)[1]
if os.path.exists("/content/input{file_extention}"):
os.remove("/content/input{file_extention}")
if os.path.exists("/content/output-audio.aac"):
os.remove("/content/output-audio.aac")
if os.path.exists("/content/Colab-Zooming-Slow-Mo/extract"):
shutil.rmtree("/content/Colab-Zooming-Slow-Mo/extract")
if os.path.exists("/content/Colab-Zooming-Slow-Mo/tmp"):
shutil.rmtree("/content/Colab-Zooming-Slow-Mo/tmp")
scale = 4
N_ot = args.N_out
N_in = 1 + N_ot // 2
os.environ['CUDA_VISIBLE_DEVICES'] = '0'
#### model
model_path = args.model
model = Sakuya_arch.LunaTokis(64, N_ot, 8, 5, 40)
#### extract the input video to temporary folder
save_folder = "/content/Colab-Zooming-Slow-Mo/extract"
save_out_folder = "/content/Colab-Zooming-Slow-Mo/tmp"
util.mkdirs(save_folder)
util.mkdirs(save_out_folder)
error = util.extract_frames(args.ffmpeg_dir, args.video, save_folder)
if error:
print(error)
exit(1)
# temporal padding mode
padding = 'replicate'
save_imgs = True
############################################################################
if torch.cuda.is_available():
device = torch.device('cuda')
else:
device = torch.device('cpu')
def single_forward(model, imgs_in):
with torch.no_grad():
# print(imgs_in.size()) # [1,5,3,270,480]
b, n, c, h, w = imgs_in.size()
h_n = int(4 * np.ceil(h / 4))
w_n = int(4 * np.ceil(w / 4))
imgs_temp = imgs_in.new_zeros(b, n, c, h_n, w_n)
imgs_temp[:, :, :, 0:h, 0:w] = imgs_in
model_output = model(imgs_temp)
model_output = model_output[:, :, :, 0:scale * h, 0:scale * w]
if isinstance(model_output, list) or isinstance(
model_output, tuple):
output = model_output[0]
else:
output = model_output
return output
model.load_state_dict(torch.load(model_path), strict=True)
model.eval()
model = model.to(device)
#### zsm images
imgs = util.read_seq_imgs(save_folder)
select_idx_list = util.test_index_generation(False, N_ot, len(imgs))
rootdir = '/content/Colab-Zooming-Slow-Mo/extract'
files = glob.glob(rootdir + '/**/*.png', recursive=True)
final_amount_images = np.amax(np.amax(select_idx_list))
previous = 0
for select_idxs in select_idx_list:
start = time.time()
# get input images
select_idx = select_idxs[0]
imgs_in = imgs.index_select(
0, torch.LongTensor(select_idx)).unsqueeze(0).to(device)
output = single_forward(model, imgs_in)
outputs = output.data.float().cpu().squeeze(0)
# save imgs
out_idx = select_idxs[1]
for idx, name_idx in enumerate(out_idx):
output_f = outputs[idx, ...].squeeze(0)
if save_imgs:
output = util.tensor2img(output_f)
cv2.imwrite(
osp.join(save_out_folder, '{:06d}.png'.format(name_idx)),
output)
if (name_idx - previous) != 0:
previous = name_idx
end = time.time()
final_time = end - start
time_values.append(final_time)
counter += 1
print('Image %d out of %d' %
(counter, final_amount_images))
print(
f"********** Time per frame (avg): %f s Time left: %d s **********"
% ((sum(time_values) / len(time_values)),
(sum(time_values) / len(time_values) *
(final_amount_images - counter))))
exit(0)
