Python使用OpenCV按自定义帧率提取视频帧并保存

在做室外语义分割、视觉导航与定位的时候，通常会用对一个连续的视频帧进行测试，除去常用数据集外，也经常会自己制作一些数据集，这个工具类可以按需求对视频进行分帧提取，封装好了直接可以使用，根据需求更改视频路径，保存路径和帧率即可。 

import os
import cv2
import shutil

VIDEO_PATH = 'test.mp4'  # 视频地址
EXTRACT_FOLDER = 'video'  # 存放帧图片的位置
EXTRACT_FREQUENCY = 1  # 帧提取频率

def extract_frames(video_path, dst_folder, index):
    # 主操作
    video = cv2.VideoCapture()
    if not video.open(video_path):
        print("can not open the video")
        exit(1)
    count = 1
    while True:
        _, frame = video.read()
        if frame is None:
            break
        if count % EXTRACT_FREQUENCY == 0:
            save_path = "{}/{:>03d}.jpg".format(dst_folder, index)
            cv2.imwrite(save_path, frame)
            index += 1
        count += 1
    video.release()
    # 打印出所提取帧的总数
    print("Totally save {:d} pics".format(index - 1))
def main():
    # 递归删除之前存放帧图片的文件夹，并新建一个
    try:
        shutil.rmtree(EXTRACT_FOLDER)
    except OSError:
        pass
    if not os.path.exists(EXTRACT_FOLDER):
        os.mkdir(EXTRACT_FOLDER)
    # 抽取帧图片，并保存到指定路径
    extract_frames(VIDEO_PATH, EXTRACT_FOLDER, 1)
if __name__ == '__main__':
    main()

另外再po一个工具类，用来获取指定路径下的文件夹或者文件名，因为我在使用网络进行验证的时候，要对多个数据集进行验证，一个一个图片去预测验证的时候太慢了，所以就获取整个文件夹的名称进行预测。

def getFiles(rootDir):
    files = []
    for lists in os.listdir(rootDir):
        path = os.path.join(rootDir, lists)
        files.append(path)
        # print(path)
        if os.path.isdir(path):
            getFiles(path)
    return files