我最近在处理视频数据集的光流图(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上计算,我这没有使用可以考虑采用