提取视频训练集光流

我最近在处理视频数据集的光流图（UCF-101），因为之前TSN的获取方式过于复杂，我不得不探索新的获取方式，首先先给出tsn的获取光流图的方式：tsn，经过很长时间的调试，我都没解决所有的bug，我不得不考虑其他的方法获取光流图。
参考链接1：获取光流图
使用该链接的代码，提取出光流图,但是在opencv4.0+的版本上，需要使用cv2.optflow.DualTVL1OpticalFlow_create， 并且在命令行输入pip install opencv-contrib-python才能使用。

参考链接2：umat
考虑到在gpu上计算能加速，我使用cv2.UMat移动到gpu上

参考链接3：多线程
为了加快计算速度我尝试使用多线程

废话不多说，最后的整体代码如下：

import os
import numpy as np
import cv2
from glob import glob
from tqdm import tqdm
import threading
import multiprocessing
import math
_IMAGE_SIZE = 256

root_flow_path = '/data/zhengrui/dataset/ucf-flow'

def cal_for_frames(video_path):
    capture = cv2.VideoCapture(video_path)
    frame_count = int(capture.get(cv2.CAP_PROP_FRAME_COUNT))
    frame_width = int(capture.get(cv2.CAP_PROP_FRAME_WIDTH))
    frame_height = int(capture.get(cv2.CAP_PROP_FRAME_HEIGHT))
    _, prev = capture.read()
    prev = cv2.UMat(prev)
    prev = cv2.cvtColor(prev, cv2.COLOR_BGR2GRAY)
    flow = []
    for count in tqdm(range(1, frame_count)):
        retaining, frame = capture.read()
        frame = cv2.UMat(frame)
        frame = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
        curr = frame
        tmp_flow = compute_TVL1(prev, curr)
        flow.append(tmp_flow)
        prev = curr

    return flow


def compute_TVL1(prev, curr, bound=15):
    """Compute the TV-L1 optical flow."""
    TVL1 = cv2.optflow.DualTVL1OpticalFlow_create()
    flow = TVL1.calc(prev, curr, None)

    flow = cv2.UMat.get(flow)
    assert flow.dtype == np.float32

    flow = (flow + bound) * (255.0 / (2 * bound))
    flow = np.round(flow).astype(int)
    flow[flow >= 255] = 255
    flow[flow <= 0] = 0

    return flow


def save_flow(video_flows, flow_path):
    for i, flow in enumerate(video_flows):
        if not os.path.exists(flow_path+'_u'):
            print(flow_path+'_u')
        cv2.imwrite(os.path.join(flow_path+'_u', "{:06d}.jpg".format(i)),
                    flow[:, :, 0])
        cv2.imwrite(os.path.join(flow_path+'_v', "{:06d}.jpg".format(i)),
                    flow[:, :, 1])


def extract_flow(video_path, flow_path):
    flow = cal_for_frames(video_path)
    save_flow(flow, flow_path)
    return


def sort(video_paths):
    for x in tqdm(range(len(video_paths))):
        tar_path_u = os.path.join(root_flow_path, video_paths[x].split('/')[-3], video_paths[x].split('/')[-2],
                                  os.path.splitext(video_paths[x].split('/')[-1])[0] + '_u')
        tar_path_v = os.path.join(root_flow_path, video_paths[x].split('/')[-3], video_paths[x].split('/')[-2],
                                  os.path.splitext(video_paths[x].split('/')[-1])[0] + '_v')
        tar_path = os.path.join(root_flow_path, video_paths[x].split('/')[-3], video_paths[x].split('/')[-2],
                                os.path.splitext(video_paths[x].split('/')[-1])[0])
        if not os.path.exists(tar_path_u):
            os.makedirs(tar_path_u)
        if not os.path.exists(tar_path_v):
            os.makedirs(tar_path_v)
        flow_paths = tar_path
        extract_flow(video_paths[x], flow_paths)
    return

if __name__ == '__main__':
    thread_num = 16
    root_video_path = '/data/zhengrui/dataset/ucf101'
    video_paths = glob(root_video_path+'/*/*/*.avi')
    print(len(video_paths))
    n = int(math.ceil(len(video_paths) / float(thread_num)))
    pool = multiprocessing.Pool(processes=thread_num)
    result = []
    for i in tqdm(range(0, len(video_paths), n)):
        result.append(pool.apply_async(sort, (video_paths[i: i+n],)))
    pool.close()
    pool.join()
    print('finish!')

最后还有一部分的矩阵计算可以使用pycuda移到gpu上计算，我这没有使用可以考虑采用