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 !")