一、目的

兵马未动、粮草先行。

cv领域，在设计、训练你的模型之前，最重要的事情莫过于处理数据了。

很多时候，我们获取到的原始数据并非别人整理好的图片，而是一段一段的视频；这个时候，就需要我们自己抽帧、过滤、然后标注，才能制作一批可用于训练的数据集。

其中，视频抽帧又是第一步。本文的目的，就是提供一个视频抽帧的脚本，可以同时处理多个视频，并可视化抽帧过程，同时还支持对每个视频设置是否旋转。

之所以加上旋转的功能，是由于有些时候我们拿到的视频画面不是正常的，而是顺时针、逆时针旋转了90°，或者180°，这对我们后续的数据处理带来了不便。

在此，提供我经常使用的视频抽帧方法，供大家参考。

二、主要特点及用法

主要特点：

  • 可批量处理多个视频文件夹中的视频；
  • 每个视频新建一个文件夹，保存该视频的帧；
  • 每个视频可以手动设置是否旋转以及旋转方式；
  • 可视化抽帧过程；

用法：

需要设置四个参数：if_rota 、paths、img_save_path、freq。

分别指定：是否手动设置旋转、视频文件夹路径（类型为list，每个元素为一个视频文件夹路径）、所有视频抽帧后的图片保存路径根目录、抽帧频率（帧/s）.

三、代码

# -*- coding: utf-8 -*-
"""

从文件夹读取视频列表，并逐个抽帧，保存到相应文件夹（每个视频一个文件夹，抽帧频率:1fps）

"""

import cv2
import os


# 顺时针旋转90度
def RotateClockWise90(img):
    trans_img = cv2.transpose(img)
    new_img = cv2.flip(trans_img, 1)
    return new_img


# 逆时针旋转90度
def RotateAntiClockWise90(img):
    trans_img = cv2.transpose(img)
    new_img = cv2.flip(trans_img, 0)
    return new_img


def list_dir(path, list_name):  # 传入存储的list
    for file in os.listdir(path):
        file_path = os.path.join(path, file)
        if os.path.isdir(file_path):
            list_dir(file_path, list_name)
        else:
            list_name.append(file_path)


def get_frames(video_path, save_path, frames_per_second, if_rota=False):
    cap = cv2.VideoCapture(video_path)
    cv2.namedWindow("Capture")

    fps = cap.get(5)  # 获取视频帧率
    freq = fps / frames_per_second
    total_frames = cap.get(7)
    digit_num = len(str(int(total_frames)))

    rotate_flag = 'n'
    while cap.isOpened():

        ok, frame = cap.read()
        if not ok:
            break

        frame_index = cap.get(1)  # 帧序
        h, w, _ = frame.shape

        cv2.imshow("Capture", cv2.resize(frame, (w // 2, h // 2)))

        if if_rota:
            if int(frame_index) == 1:
                if cv2.waitKey(0) & 0xFF == ord('n'):
                    print('不旋转')
                    pass
                elif cv2.waitKey(0) & 0xFF == ord('a'):
                    rotate_flag = 'a'
                    print('逆时针旋转90度')
                elif cv2.waitKey(0) & 0xFF == ord('c'):
                    rotate_flag = 'c'
                    print('顺时针旋转90度')
            if rotate_flag == 'n':
                pass
            elif rotate_flag == 'a':
                frame = RotateAntiClockWise90(frame)
            elif rotate_flag == 'c':
                frame = RotateClockWise90(frame)

        if frame_index % int(freq) != 0:
            continue
        else:
            # 保存当前帧
            img_name = video_path.split('/')[-1].split('.')[0] + '_frame_' + \
                       str(int(frame_index)).zfill(digit_num) + '.jpg'
            #            frame = cv2.resize(frame, (1080, 1920))
            print("saving frame:", img_name)
            cv2.imwrite(os.path.join(save_path, img_name), frame)

        if cv2.waitKey(1) & 0xFF == ord('q'):
            break
    cap.release()
    cv2.destroyAllWindows()


if __name__ == '__main__':

    if_rota = False
    paths = ["video_root_path1", "video_root_path2"]
    img_save_path = 'path_to_save_imgs'  # 本批次图片保存路径
    freq = 1  # 抽帧频率：freq张/s

    if not os.path.isdir(img_save_path):
        os.mkdir(img_save_path)

    list_videos = []
    # 获取路径中的视频列表
    for path in paths:
        path = path.replace('\\', '/')
        list_dir(path, list_videos)
    list_videos = [v for v in list_videos if v[-4:] in ['.mp4', '.avi']]

    #    # 针对每个视频，分别按照指定频率抽取帧，保存在文件夹中
    for video_path in list_videos:
        video_path = video_path.replace('\\', '/')
        print('正在处理：', video_path.split('\\')[-1].split('.')[0], '...')
        if if_rota:
            print('不旋转：单击“n”；逆时针旋转90度：双击“a”;顺时针旋转90度：三击“c”')
            save_path = os.path.join(img_save_path, video_path.split('\\')[-1].split('.')[0])

        if not os.path.isdir(save_path):
            os.mkdir(save_path)

        # 调用视频抽帧方法
        get_frames(video_path, save_path, frames_per_second=freq, if_rota=if_rota)