def batch_shift_images(ori_root,
dest_root,
offset_x=0.,
offset_y=0.,
filename_template="{}.png"): # TODO
'''
function:
shifting images by (offset_x, offset_y) on (axis-x, axis-y) in batches
params:
ori_root: string, the dir of images that need to be processed
dest_root: string, the dir to save processed images
offset_x: float, offset pixels on axis-x
positive=left; negative=right
offset_y: float, offset pixels on axis-y
positive=up; negative=down
filename_template: string, the filename template for saving images
'''
handle_dir(dest_root)
offset_x, offset_y = float(offset_x), float(offset_y)
images_fname = sorted(listdir(ori_root))
for imf in images_fname:
img = cv2.imread(os.path.join(ori_root, imf)).astype('float32')
img = image_shift(img, offset_x=offset_x, offset_y=offset_y)
cv2.imwrite(
os.path.join(dest_root,
filename_template.format(get_fname_ext(imf)[0])), img)
print("Image", imf, "shift done !")
def batch_crop_img_with_padding(ori_root,
dest_root,
min_size=(100, 100),
padding=10):
'''
function:
cropping image to many patches with padding
it can be used for inferring large image
params:
ori_root: the dir of images that need to be processed
dest_root: the dir to save processed images
min_size: a tuple (h, w) the min size of crop, the border patch will be larger
padding: the padding size of each patch
notice:
filenames should not contain the character "-"
the crop flag "x-x-x-x" will be at the end of filename when cropping
'''
handle_dir(dest_root)
images_fname = sorted(listdir(ori_root))
for imf in images_fname:
img = cv2.imread(os.path.join(ori_root, imf))
img_cropped = crop_img_with_padding(img,
min_size=min_size,
padding=padding)
for k in img_cropped.keys():
cv2.imwrite(
os.path.join(
dest_root,
"{}_{}.png".format(os.path.basename(imf).split('.')[0],
k)), img_cropped[k])
print(imf, "crop done !")
def extra_frames_by_postfix(ori_root,
dest_root,
match_postfix='',
new_postfix=None,
match_ext='*'):
'''
extra frames from ori_root to dest_root by match_postfix and match_ext
params:
ori_root: the dir of videos that need to be processed
dest_root: the dir for saving matched files
match_postfix: the postfix to be matched
new_postfix: the postfix for matched files
default: None, that is keeping the ori postfix
match_ext: the ext to be matched
'''
if new_postfix is None:
new_postfix = match_postfix
handle_dir(dest_root)
video_list = listdir(ori_root)
for v in video_list:
file_regroup_utils.extra_files_by_postfix(
ori_root=os.path.join(ori_root, v),
dest_root=os.path.join(dest_root, v),
match_postfix=match_postfix,
new_postfix=new_postfix,
match_ext=match_ext)
def batch_traverse_crop_img(ori_root,
dest_root,
dsize=(100, 100),
interval=10):
'''
function:
traversing crop image to many patches with same interval
params:
ori_root: the dir of images that need to be processed
dest_root: the dir to save processed images
dsize: a tuple (h, w) the size of crop, the border patch will be overlapped for satisfing the dsize
interval: the interval when traversing
'''
handle_dir(dest_root)
images_fname = sorted(listdir(ori_root))
for imf in images_fname:
img = cv2.imread(os.path.join(ori_root, imf))
img_cropped = traverse_crop_img(img, dsize=dsize, interval=interval)
for i, cim in enumerate(img_cropped):
cv2.imwrite(
os.path.join(
dest_root,
"{}_{}.png".format(os.path.basename(imf).split('.')[0],
i)), cim)
print(imf, "crop done !")
def extra_frames_from_videos(ori_root, save_root, fname_template='%4d.png', start_end=None):
'''
function:
ext frames from videos
params:
ori_root: string, the dir of videos that need to be processed
dest_root: string, the dir to save processed videos
fname_template: the template for frames' filename
start_end: list, len=2, the start and end index for processed videos,
assert: len(start_end)=2
default: None, that is processing all videos
'''
handle_dir(save_root)
videos = sorted(listdir(ori_root))
if start_end is not None:
assert len(start_end) == 2, "only support len(start_end)=2"
videos = videos[start_end[0]:start_end[1]]
for v in videos:
vn = v[:-(len(v.split('.')[-1]) + 1)]
video_path = os.path.join(ori_root, v)
png_dir = os.path.join(save_root, vn)
png_path = os.path.join(png_dir, fname_template)
handle_dir(png_dir)
command = 'ffmpeg -i {} {}'.format(video_path, png_path)
os.system(command)
print("Extra frames from {}".format(video_path))
def remove_head_tail_frames(root, recycle_bin=None, num=0):
'''
remove num hean&tail frames from videos
params:
root: the dir of files that need to be processed
recycle_bin: the removed frames will be put here
defalut: None, that is putting the removed frames in root/_recycle_bin
num: the number of frames to be removed
'''
if recycle_bin is None:
recycle_bin = os.path.join(root, '_recycle_bin')
handle_dir(recycle_bin)
video_list = listdir(root)
for v in video_list:
img_list = sorted(glob_match(os.path.join(root, v, "*")))
handle_dir(os.path.join(recycle_bin, v))
for i in range(num):
src = img_list[i]
dest = os.path.join(recycle_bin, v, os.path.basename(src))
move_file(src, dest)
src = img_list[-(i + 1)]
dest = os.path.join(recycle_bin, v, os.path.basename(src))
move_file(src, dest)
def batch_matlab_resize_images(ori_root,
dest_root,
scale=1.0,
method='bicubic',
filename_template="{}.png"):
'''
function:
resizing images in batches, same as matlab2017 imresize
params:
ori_root: string, the dir of images that need to be processed
dest_root: string, the dir to save processed images
scale: float, the resize scale
method: string, the interpolation method,
optional: 'bilinear', 'bicubic'
default: 'bicubic'
filename_template: string, the filename template for saving images
'''
if method != 'bilinear' and method != 'bicubic':
raise Exception('Unknown method!')
handle_dir(dest_root)
scale = float(scale)
images_fname = sorted(listdir(ori_root))
for imf in images_fname:
img = cv2.imread(os.path.join(ori_root, imf)).astype('float32')
img = matlab_imresize(img, scalar_scale=scale, method=method)
cv2.imwrite(
os.path.join(dest_root,
filename_template.format(get_fname_ext(imf)[0])), img)
print("Image", imf, "resize done !")
def zip_frames_to_videos(ori_root, save_root, fname_template='%4d.png', video_ext='mp4', start_end=None):
'''
function:
zip frames to videos
params:
ori_root: string, the dir of videos that need to be processed
dest_root: string, the dir to save processed videos
fname_template: the template of frames' filename
start_end: list, len=2, the start and end index for processed videos,
assert: len(start_end)=2
default: None, that is processing all videos
'''
handle_dir(save_root)
videos_name = sorted(listdir(ori_root))
if start_end is not None:
assert len(start_end) == 2, "only support len(start_end)=2"
videos_name = videos_name[start_end[0]:start_end[1]]
for vn in videos_name:
imgs_path = os.path.join(ori_root, vn, fname_template)
video_path = os.path.join(save_root, '{}.{}'.format(vn, video_ext))
command = 'ffmpeg -i {} -vcodec libx264 -crf 16 -pix_fmt yuv420p {}'.format(imgs_path, video_path)
# command = 'ffmpeg -r 24000/1001 -i {} -vcodec libx265 -pix_fmt yuv422p -crf 10 {}'.format(
# imgs_path, video_path) # youku competition
os.system(command)
print("Zip frames to {}".format(video_path))
def batch_combine_img(ori_root, dest_root, padding=10):
'''
function:
combining many patches to image
it can be used to combine patches to image, when you finish inferring large image with cropped patches
params:
ori_root: the dir of images that need to be processed
dest_root: the dir to save processed images
padding: the padding size of each patch
notice:
filenames should not contain the character "-" except for the crop flag
the crop flag "x-x-x-x" should be at the end of filename when combining
'''
handle_dir(dest_root)
images_fname = [
fn[:-(len(fn.split('_')[-1]) + 1)] for fn in listdir(ori_root)
]
images_fname = list(set(images_fname))
for imf in images_fname:
croped_imgs_path = sorted(
glob_match(os.path.join(ori_root, "{}*".format(imf))))
croped_imgs = {}
for cip in croped_imgs_path:
img = cv2.imread(cip)
k = cip.split('.')[0].split('_')[-1]
croped_imgs[k] = img
img_combined = combine_img(croped_imgs, padding=padding)
cv2.imwrite(os.path.join(dest_root, "{}.png".format(imf)),
img_combined)
print("{}.png".format(imf), "combine done !")
def calc_video_PSNR_SSIM(output_root, gt_root, crop_border=4, shift_window_size=0, test_ycbcr=False, crop_GT=False):
'''
计算视频的 PSNR、SSIM，使用 EDVR 的计算方式
要求 output_root, gt_root 中的文件按顺序一一对应
'''
PSNR_sum = 0.
SSIM_sum = 0.
img_num = 0
video_PSNR = []
video_SSIM = []
video_list = sorted(listdir(output_root))
for v in video_list:
v_PSNR_list, v_SSIM_list, _ = calc_image_PSNR_SSIM(
output_root=os.path.join(output_root, v),
gt_root=os.path.join(gt_root, v),
crop_border=crop_border, shift_window_size=shift_window_size,
test_ycbcr=test_ycbcr, crop_GT=crop_GT
)
PSNR_sum += sum(v_PSNR_list)
SSIM_sum += sum(v_SSIM_list)
img_num += len(v_PSNR_list)
video_PSNR.append({
'video_name': v,
'psnr': v_PSNR_list
})
video_SSIM.append({
'video_name': v,
'ssim': v_SSIM_list
})
logs = []
PSNR_SSIM_csv_log = {
'col_names': [],
'row_names': [output_root],
'psnr_ssim': [[]]
}
for v_psnr, v_ssim in zip(video_PSNR, video_SSIM):
PSNR_SSIM_csv_log['col_names'].append('#{}'.format(v_psnr['video_name']))
PSNR_SSIM_csv_log['psnr_ssim'][0].append('{:.5}/{:.4}'.format(sum(v_psnr['psnr']) / len(v_psnr['psnr']),
sum(v_ssim['ssim']) / len(v_ssim['ssim'])))
log = 'Video: {} PSNR={:.5}, SSIM={:.4}'.format(v_psnr['video_name'],
sum(v_psnr['psnr']) / len(v_psnr['psnr']),
sum(v_ssim['ssim']) / len(v_ssim['ssim']))
print(log)
logs.append(log)
PSNR_SSIM_csv_log['col_names'].append('AVG')
PSNR_SSIM_csv_log['psnr_ssim'][0].append('{:.5}/{:.4}'.format(PSNR_sum / img_num, SSIM_sum / img_num))
log = 'Average PSNR={:.5}, SSIM={:.4}'.format(PSNR_sum / img_num, SSIM_sum / img_num)
print(log)
logs.append(log)
return PSNR_SSIM_csv_log, logs
def add_frames_postfix(root, postfix=''):
'''
add postfix to frames
params:
root: the dir of videos that need to be processed
postfix: the postfix to be added
'''
video_list = listdir(root)
for v in video_list:
file_regroup_utils.add_files_postfix(os.path.join(root, v),
postfix=postfix)
def remove_frames_postfix(root, postfix=''):
'''
remove postfix from frames
params:
root: the dir of videos that need to be processed
postfix: the postfix to be removed
'''
video_list = listdir(root)
for v in video_list:
file_regroup_utils.remove_files_postfix(os.path.join(root, v),
postfix=postfix)
def VideoFlag2FlagVideo(ori_root, dest_root, ori_flag, dest_flag=None):
'''
videos/type/frames --> type/videos/frames
params:
ori_root: the dir of files that need to be processed
dest_root: the dir for saving matched files
ori_flag: the ori video flag(e.g. blur)
dest_flag: the flag(e.g. blur) for saving videos
default: None, that is keeping the ori flag
'''
if dest_flag is None:
dest_flag = ori_flag
handle_dir(dest_root)
handle_dir(os.path.join(dest_root, dest_flag))
video_list = listdir(ori_root)
for v in video_list:
image_list = listdir(os.path.join(ori_root, v, ori_flag))
handle_dir(os.path.join(dest_root, dest_flag, v))
for im in image_list:
src = os.path.join(ori_root, v, ori_flag, im)
dst = os.path.join(dest_root, dest_flag, v, im)
copy_file(src, dst)
def resort_frames_index(root, template='{:0>4}', start_idx=0):
'''
resort frames' filename using template that index start from start_idx
params:
root: the dir of files that need to be processed
template: the template for processed filename
start_idx: the start index
'''
video_list = listdir(root)
for v in video_list:
file_regroup_utils.resort_files_index(os.path.join(root, v),
template=template,
start_idx=start_idx)
def copy_frames_for_fps(ori_root, save_root, mul=12, fname_template="{:0>4}", ext="png"):
'''
function:
copy frames for fps
params:
ori_root: string, the dir of videos that need to be processed
dest_root: string, the dir to save processed videos
mul: the multiple of copy
fname_template: the template of frames' filename
ext: the ext of frames' filename
'''
fname_template = fname_template + '.{}'
videos_name = sorted(listdir(ori_root))
handle_dir(save_root)
for vn in videos_name:
frmames = sorted(listdir(os.path.join(ori_root, vn)))
handle_dir(os.path.join(save_root, vn))
for i, f in enumerate(frmames):
for j in range(mul):
now_idx = i * mul + j
src = os.path.join(ori_root, vn, f)
dest = os.path.join(save_root, vn, fname_template.format(now_idx, ext))
copy_file(src, dest)
def calc_kernel_gradient_similarity(output_root, gt_root):
'''
计算 kernel 的 gradient similarity
要求 output_root, gt_root 中的文件按顺序一一对应
'''
GS_list = []
output_kernel_list = sorted(listdir(output_root))
gt_kernel_list = sorted(listdir(gt_root))
for o_k, g_k in zip(output_kernel_list, gt_kernel_list):
o_k_path = os.path.join(output_root, o_k)
g_k_path = os.path.join(gt_root, g_k)
kernel_GT = load_mat_kernel(g_k_path)
kernel_Gen = load_mat_kernel(o_k_path)
gs = kernel_base.Gradient_Similarity(kernel_Gen, kernel_GT)
GS_list.append(gs)
print("{} Gradient-Similarity={:.4}".format(o_k, gs))
log = 'Average Gradient-Similarity={:.4}'.format(sum(GS_list) / len(GS_list))
print(log)
return GS_list, log
def calc_video_LPIPS(output_root, gt_root, model=None, use_gpu=False, spatial=True):
'''
计算视频的 LPIPS
要求 output_root, gt_root 中的文件按顺序一一对应
'''
if model is None:
model = lpips_models.PerceptualLoss(model='net-lin', net='alex', use_gpu=use_gpu, spatial=spatial)
LPIPS_sum = 0.
img_num = 0
video_LPIPS = []
video_list = sorted(listdir(output_root))
for v in video_list:
v_LPIPS_list, _ = calc_image_LPIPS(
output_root=os.path.join(output_root, v),
gt_root=os.path.join(gt_root, v),
model=model, use_gpu=use_gpu, spatial=spatial
)
LPIPS_sum += sum(v_LPIPS_list)
img_num += len(v_LPIPS_list)
video_LPIPS.append({
'video_name': v,
'lpips': v_LPIPS_list
})
logs = []
LPIPS_csv_log = {
'col_names': [],
'row_names': [output_root],
'lpips': [[]]
}
for v_lpips in video_LPIPS:
LPIPS_csv_log['col_names'].append('#{}'.format(v_lpips['video_name']))
LPIPS_csv_log['lpips'][0].append('{:.4}'.format(sum(v_lpips['lpips']) / len(v_lpips['lpips'])))
log = 'Video: {} LPIPS={:.4}'.format(v_lpips['video_name'], sum(v_lpips['lpips']) / len(v_lpips['lpips']))
print(log)
logs.append(log)
LPIPS_csv_log['col_names'].append('AVG')
LPIPS_csv_log['lpips'][0].append('{:.4}'.format(LPIPS_sum / img_num))
log = 'Average LPIPS={:.4}'.format(LPIPS_sum / img_num)
print(log)
logs.append(log)
return LPIPS_csv_log, logs
def save_flow_pt2mat(ori_root, dest_root):
'''
把 ori_root 中视频的 flow 从 pt 文件转换为 mat 文件后保存到 dest_root
保存在 pt 文件中的 flow 的 shape=[1, 2, h, w]
'''
handle_dir(dest_root)
video_list = listdir(ori_root)
for v in video_list:
handle_dir(os.path.join(dest_root, v))
file_list = glob_match(os.path.join(ori_root, v, "*.pt"))
for pt_path in file_list:
fname, ext = get_fname_ext(pt_path)
mat_path = os.path.join(dest_root, v, "{}.mat".format(fname))
try:
flow = torch.load(pt_path)[0].permute(1, 2, 0).cpu().numpy()
flow_dict = {'flow': flow}
scio.savemat(mat_path, flow_dict)
print('save {} to {}'.format(pt_path, mat_path))
except:
print('skip file {}'.format(pt_path))
def batch_cv2_resize_images(ori_root,
dest_root,
scale=1.0,
method='bicubic',
filename_template="{}.png"):
'''
function:
resizing images in batches
params:
ori_root: string, the dir of images that need to be processed
dest_root: string, the dir to save processed images
scale: float, the resize scale
method: string, the interpolation method,
optional: 'nearest', 'bilinear', 'bicubic'
default: 'bicubic'
filename_template: string, the filename template for saving images
'''
if method == 'nearest':
interpolation = cv2.INTER_NEAREST
elif method == 'bilinear':
interpolation = cv2.INTER_LINEAR
elif method == 'bicubic':
interpolation = cv2.INTER_CUBIC
else:
raise Exception('Unknown method!')
handle_dir(dest_root)
scale = float(scale)
images_fname = sorted(listdir(ori_root))
for imf in images_fname:
img = cv2.imread(os.path.join(ori_root, imf)).astype('float32')
img = cv2.resize(img,
dsize=(0, 0),
fx=scale,
fy=scale,
interpolation=interpolation)
cv2.imwrite(
os.path.join(dest_root,
filename_template.format(get_fname_ext(imf)[0])), img)
print("Image", imf, "resize done !")
def batch_select_valid_patch(ori_root, dest_root, thre=7):
'''
function:
selecting valid patch that are not too smooth
params:
ori_root: the dir of patches that need to be selected
dest_root: the dir to save selected patch
thre: the threshold value of smooth
'''
handle_dir(dest_root)
images_fname = sorted(listdir(ori_root))
total_num = len(images_fname)
valid_num = 0
for imf in images_fname:
img = cv2.imread(os.path.join(ori_root, imf))
smooth = evaluate_smooth(img)
if smooth > thre:
cv2.imwrite(os.path.join(dest_root, imf), img)
valid_num += 1
else:
print(imf, "too smooth, smooth={}".format(smooth))
print("Total {} patches, valid {}, remove {}".format(
total_num, valid_num, total_num - valid_num))
def calc_kernel_gradient_similarity_video(output_root, gt_root):
'''
计算视频的 kernel 的 gradient similarity
要求 output_root, gt_root 中的文件按顺序一一对应
'''
GS_sum = 0.
kernel_num = 0
video_GS = []
video_list = sorted(listdir(output_root))
for v in video_list:
v_GS_list, _ = calc_kernel_gradient_similarity(
output_root=os.path.join(output_root, v),
gt_root=os.path.join(gt_root, v)
)
GS_sum += sum(v_GS_list)
kernel_num += len(v_GS_list)
video_GS.append({
'video_name': v,
'gradient_similarity': v_GS_list
})
logs = []
for v_lpips in video_GS:
log = 'Video: {} Gradient-Similarity={:.4}'.format(
v_lpips['video_name'],
sum(v_lpips['gradient_similarity']) / len(v_lpips['gradient_similarity']))
print(log)
logs.append(log)
log = 'Average Gradient-Similarity={:.4}'.format(GS_sum / kernel_num)
print(log)
logs.append(log)
return video_GS, logs
def batch_shift_videos(ori_root, dest_root, offset_x=0., offset_y=0., filename_template="{}.png"):
'''
function:
shifting videos by (offset_x, offset_y) on (axis-x, axis-y) in batches
params:
ori_root: string, the dir of videos that need to be processed
dest_root: string, the dir to save processed videos
offset_x: float, offset pixels on axis-x
positive=left; negative=right
offset_y: float, offset pixels on axis-y
positive=up; negative=down
filename_template: string, the filename template for saving images
'''
handle_dir(dest_root)
videos = listdir(ori_root)
for v in videos:
batch_shift_images(
ori_root=os.path.join(ori_root, v),
dest_root=os.path.join(dest_root, v),
offset_x=offset_x,
offset_y=offset_y,
filename_template=filename_template
)
print("Video", v, "shift done !")
def batch_matlab_resize_videos(ori_root, dest_root, scale=1.0, method='bicubic', filename_template="{}.png"):
'''
function:
resizing videos in batches, same as matlab2017 imresize
params:
ori_root: string, the dir of videos that need to be processed
dest_root: string, the dir to save processed videos
scale: float, the resize scale
method: string, the interpolation method,
optional: 'bilinear', 'bicubic'
default: 'bicubic'
filename_template: string, the filename template for saving images
'''
handle_dir(dest_root)
videos = listdir(ori_root)
for v in videos:
batch_matlab_resize_images(
ori_root=os.path.join(ori_root, v),
dest_root=os.path.join(dest_root, v),
scale=scale,
method=method,
filename_template=filename_template
)
print("Video", v, "resize done !")