一、目的
兵马未动、粮草先行。
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)