python保存指定范围帧内的图片

#! /usr/bin/env python
#-*- coding: utf-8 -*-


import argparse
import cv2
import os
import threading
import time
import multiprocessing
import sys
reload(sys)
sys.setdefaultencoding('utf-8')


def usage():
    parser = argparse.ArgumentParser()    
    parser.add_argument("-f", "--video_file",
            required=True,
            help="the video file, must give the path",
            dest="f")
    parser.add_argument("-r", "--frame_range",
            nargs='*',
            required=True,
            type=int,
            help="the range of the frame you want to get",
            dest="r")
    parser.add_argument("-m", "--mod",
            default=5,
            type=int,
            help="every mod frame get one",
            dest="m")
    parser.add_argument("-out", "--out_dir",
            default="./frames",
            help="the save directory of the frames",
            dest="out")
    parser.add_argument("-t", "--thread",
            default=2,
            type=int,
            help="the process num",
            dest="t")
    #import pdb; pdb.set_trace();    
    args = parser.parse_args()
    video_file = args.f
    frame_range = args.r
    mod = args.m
    out_dir = args.out
    thread = args.t

    return video_file, frame_range, mod, out_dir, thread

def get_frames(video_file, frame_range, mod):
    logging.info("in get_frames:")
    frame_list = []
    video_reader = cv2.VideoCapture(video_file)
    if not video_reader.isOpened():
        print "get %s frame failed"%video_file
        return
    i = 0
    frame_lower_bound = frame_range[0]
    if len(frame_range) == 1:
        frame_upper_bound = 99999999
    if len(frame_range) == 2:
        frame_upper_bound = frame_range[1]
    while True:
        i += 1
        print("  in while:%d"%(i))
        succ, frame = video_reader.read()
        if not succ:
            break
        if i < frame_lower_bound:
            continue
        if i > frame_upper_bound:
            break
        #每mod帧得到一帧图片
        if i % mod != 0:
            continue
        frame_list.append((frame, i))
    video_reader.release()
    return frame_list

def make_dst_file(dst_dir, file_name):
    base_name = os.path.basename(file_name)
    if not os.path.exists(dst_dir):
        os.makedirs(dst_dir)
    return os.path.join(dst_dir, base_name)

#保存图片
def save_frames(frame_list, out_dir, video_name):
    out_dir = "{}/{}".format(out_dir, video_name)
    for frame_pair in frame_list:
        frame_index = frame_pair[1]
        frame_index_str = (str(frame_index)).zfill(10)
        frame = frame_pair[0]
        h, w, _ = frame.shape
        #frame = cv2.resize(frame, (w/2, h/2))
        frame_name = "%s_%s.jpg"%(video_name, frame_index_str)
        frame_save_path = make_dst_file(out_dir, frame_name)
        print("save frame: %s"%(frame_save_path))
        cv2.imwrite(frame_save_path, frame)

def main():
    start_time = time.time()
    video_file, frame_range, mod, out_dir, thread_num = usage()
    #import pdb; pdb.set_trace();
    print("in main:video_file:%s,
            frame_range:%s, %s,mod:%d, out_dir:%s"
            %(video_file, str(frame_range[0]),
            str(frame_range[1]), mod, out_dir))
    #获取指定范围的帧，并存入列表
    frame_list = get_frames(video_file, frame_range, mod)
    video_name = video_file.split('/')[-1] 
    step = len(frame_list)/thread_num + 1
    proc = []
    for i in range(0, len(frame_list), step):    
        p = multiprocessing.Process(target=save_frames, args=(frame_list[i:i+step], out_dir, video_name,))
        proc.append(p)

    for p in proc:
        p.start()
    for p in proc:
        p.join()

    end_time = time.time()
    diff = end_time-start_time
    print "time cost:{}".format(diff)

if __name__ == "__main__":
    main()