Python并行处理视频帧

参考链接：Speedy Computer Vision Pipelines using Parallelism

方案：

• 使用Python多进程编程，将视频分成多个小段，可按照CPU核数num_processes = multiprocessing.cpu_count()等分帧数。

• 通过cv2.CAP_PROP_FRAME_COUNT可以得到视频的总帧数(不准确),然后得到每一段视频的帧数大约为FRAMES_PER_SEG = cv2.CAP_PROP_FRAME_COUNT // multiprocessing.cpu_count()。

• 通过cv2.CAP_PROP_POS_FRAMES可以设置每一小段视频开始读取的起始帧位置：video.set(cv2.CAP_PROP_POS_FRAMES, seg_count * FRAMES_PER_SEG).

• 使用进程池multiprocessing.Pool(num_processes)

测试

测试视频：分辨率 800x450， 时长:00:01:56.40， FPS: 25, 总帧数：2908

• 常规操作：将视频存储为逐帧图片

import os
import sys
import time
import cv2
import argparse
 
 
def process_video(input, output):
    if not os.path.exists(output):
        os.makedirs(output)
 
    vid = cv2.VideoCapture(input)
    proc_frames = 0
    ret, frame = vid.read()
    while ret:
        out_file = output + '/{:0>6d}.png'.format(proc_frames)
        cv2.imwrite(out_file, frame)
        ret, frame = vid.read()
        proc_frames += 1
    vid.release()
 
 
if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("--input", default="heheda.mp4", type=str)
    parser.add_argument("--output", default="images", type=str)
    args = parser.parse_args()
 
    start_time = time.time()
    process_video(args.input, args.output)
 
    print(
        "Method {}: Input:{}, Time taken: {}".format(
            sys.argv[0], args.input, time.time() - start_time
        )
    )

测试耗时： 47.633s

• 多进程并行：将视频存储为逐帧图片

import os
import cv2
import sys
import time
import argparse
import multiprocessing as mp
 

def process_video(group_number):
    vid = cv2.VideoCapture(args.input)
    pos_frames = int(frame_jump_unit * group_number)
    vid.set(cv2.CAP_PROP_POS_FRAMES, pos_frames)
    proc_frames = 0
 
    if not os.path.exists(args.output):
        os.makedirs(args.output)
 
    if group_number != num_processes-1:
        not_last_seg = True
    else:
        not_last_seg = False
 
    ret, frame = vid.read()
    while ret:
        if not_last_seg and proc_frames == frame_jump_unit:
            break
        out_file = args.output + '/{:0>6d}.png'.format(proc_frames + pos_frames)
        cv2.imwrite(out_file, frame)
        ret, frame = vid.read()
        proc_frames += 1
    vid.release()
    return None
 
 
if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("--input", default="heheda.mp4", type=str)
    parser.add_argument("--output", default="images", type=str)
    args = parser.parse_args()
 
    start_time = time.time()
 
    num_processes = mp.cpu_count()
    vid = cv2.VideoCapture(args.input)
    frame_jump_unit = vid.get(cv2.CAP_PROP_FRAME_COUNT) // num_processes
    width = int(vid.get(cv2.CAP_PROP_FRAME_WIDTH))
    height = int(vid.get(cv2.CAP_PROP_FRAME_HEIGHT))
    fps = vid.get(cv2.CAP_PROP_FPS)
    vid.release()
 
    p = mp.Pool(num_processes)
    p.map(process_video, range(num_processes))
 
    print(
        "Method {}: Input:{}, Output:{}, Time taken: {}".format(
            sys.argv[0], args.input, args.output, time.time() - start_time
        )
    )

测试耗时： 20.775s

• 多进程并行：将视频处理后另存为视频
  代码来自github，将视频分段后并行处理各小段的视频帧，然后用ffmpeg进行视频拼接。

import cv2
import sys
import time
import argparse
import numpy as np
import subprocess as sp
import multiprocessing as mp


def process_video(group_number):
    cap = cv2.VideoCapture(args.input_file)
    cap.set(cv2.CAP_PROP_POS_FRAMES, frame_jump_unit * group_number)
    proc_frames = 0
    out = cv2.VideoWriter(
        "{}.{}".format(group_number, args.extension),
        cv2.VideoWriter_fourcc(*vid_fourcc[args.extension]),
        fps,
        (width, height),
    )

    while proc_frames < frame_jump_unit:
        ret, frame = cap.read()
        if ret == False:
            break
        out.write(cv2.filter2D(frame, -1, kernel))
        proc_frames += 1

    cap.release()
    out.release()
    return None


if __name__ == "__main__":
    kernel = np.ones((7,7),np.float32)/49
    parser = argparse.ArgumentParser()
    parser.add_argument("--input_file", default="Kiiara.mp4", type=str)
    parser.add_argument("--x264", default=False, type=bool)
    parser.add_argument("--extension", choices=["mp4", "avi", "mkv"], default="avi")
    parser.add_argument("--extra_flags", default="", type=str)
    args = parser.parse_args()
    start_time = time.time()
    
    num_processes = mp.cpu_count()
    vid_fourcc = {"avi": "MJPG", "mp4": "mp4v", "mkv": "mp4v"}
    cap = cv2.VideoCapture(args.input_file)
    frame_jump_unit = cap.get(cv2.CAP_PROP_FRAME_COUNT) // num_processes
    width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
    height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
    fps = cap.get(cv2.CAP_PROP_FPS)
    cap.release()

    p = mp.Pool(num_processes)
    p.map(process_video, range(num_processes))

    intermediate_files = ["{}.{}".format(i, args.extension) for i in range(num_processes)]
    with open("intermediate_files.txt", "w") as f:
        for t in intermediate_files:
            f.write("file {} \n".format(t))

    ffmpeg_command = "ffmpeg -y -loglevel error -f concat -safe 0 -i intermediate_files.txt {}".format(
        args.extra_flags
    ).strip()
    if args.x264 == False:
        ffmpeg_command += " -vcodec copy"
    else:
        ffmpeg_command += " -vcodec libx264"
    ffmpeg_command += " output_method2.{}".format(args.extension)
    t2 = time.time()
    sp.Popen(ffmpeg_command, shell=True).wait()
    t3 = time.time()

    from os import remove

    for f in intermediate_files:
        remove(f)
    remove("intermediate_files.txt")
    print(
        "Method {}: Input:{}, re-encode to x264:{}, output extension:{}, extra_flags:{}, Time taken: {}".format(
            sys.argv[0], args.input_file, args.x264, args.extension, args.extra_flags, t3 - start_time
        )
    )