手把手写深度学习(0)：专栏文章导航
前言：当算法工程师拿到一个新的视频数据集的时候，需要首先查看一下这个数据集的基本特性，方便后续处理和模型训练。这篇博客提供自动化脚本，帮助统计视频数据集的基本特性（如帧率、帧数、秒数）等，然后根据自己的需求对视频数据集进行切分预处理。
目录
获取文件夹下的所有视频文件
使用decord读取视频文件
统计文件夹下视频的平均帧率和fps
完整统计视频特性代码
数据剪切
获取文件夹下的所有视频文件
def get_video_files(directory):
video_files = []
# 遍历指定目录
for root, dirs, files in os.walk(directory):
for file in files:
# 检查文件扩展名是否为视频格式
if file.lower().endswith(('.mp4', '.avi', '.mov', '.mkv', '.wmv', '.flv')):
video_files.append(os.path.join(root, file))
return video_files使用decord读取视频文件
def get_frame_count(video_path):
# 创建VideoReader对象
vr = VideoReader(video_path)
# 获取视频的帧数
frame_count = len(vr)
fps = vr.get_avg_fps()
return frame_count, fps统计文件夹下视频的平均帧率和fps
# 统计文件夹下视频的平均帧率和fps
def get_average_frame_fps(video_path):
video_files = get_video_files(video_paths)
print(f'total videos: {len(video_files)}')
frame_count_list = []
fps_list = []
second_list = []
for index, video_file in enumerate(video_files):
frame_count, fps = get_frame_count(video_file)
second = frame_count / fps
print(f'[{index}] The video has {frame_count} frames, fps is {fps}, seconds is {second}.')
frame_count_list.append(frame_count)
fps_list.append(fps)
second_list.append(second)
average_frame_count = sum(frame_count_list) / len(frame_count_list)
average_fps = sum(fps_list) / len(fps_list)
print(f"total videos numbers: {len(fps_list)}, average_frame_count is {average_frame_count}, average_fps is {average_fps}")
return average_frame_count, average_fps完整统计视频特性代码
import os
from decord import VideoReader
from decord import cpu, gpu
def get_video_files(directory):
video_files = []
# 遍历指定目录
for root, dirs, files in os.walk(directory):
for file in files:
# 检查文件扩展名是否为视频格式
if file.lower().endswith(('.mp4', '.avi', '.mov', '.mkv', '.wmv', '.flv')):
video_files.append(os.path.join(root, file))
return video_files
def get_frame_count(video_path):
# 创建VideoReader对象
vr = VideoReader(video_path)
# 获取视频的帧数
frame_count = len(vr)
fps = vr.get_avg_fps()
return frame_count, fps
# 统计文件夹下视频的平均帧率和fps
def get_average_frame_fps(video_path):
video_files = get_video_files(video_paths)
print(f'total videos: {len(video_files)}')
frame_count_list = []
fps_list = []
second_list = []
for index, video_file in enumerate(video_files):
frame_count, fps = get_frame_count(video_file)
second = frame_count / fps
print(f'[{index}] The video has {frame_count} frames, fps is {fps}, seconds is {second}.')
frame_count_list.append(frame_count)
fps_list.append(fps)
second_list.append(second)
average_frame_count = sum(frame_count_list) / len(frame_count_list)
average_fps = sum(fps_list) / len(fps_list)
print(f"total videos numbers: {len(fps_list)}, average_frame_count is {average_frame_count}, average_fps is {average_fps}")
return average_frame_count, average_fps
if __name__ == "__main__":
# video_paths = 'your_own_path'
get_average_frame_fps(video_paths)
至此，我们可以统计出视频数据集中的基本特性，下一步根据统计出来的特性剪切数据。
数据剪切
目前的视频生成模型大多数生成能力在30帧一下，所以训练的时候不能处理太长的视频；收集的视频一半的fps等于30或60，所以我们需要把收集到的视频进行剪切。
下面的函数实现了这一功能，frames_per_segment代表我们想要剪切视频的帧数量，fps是新保存视频的fps。
def clip_and_save_video(video_path, output_folder, frames_per_segment=60, fps=None):
video_name = video_path.split('/')[-1].split('.')[0]
os.makedirs(output_folder, exist_ok=True)
# 打开视频文件
cap = cv2.VideoCapture(video_path)
# 获取视频帧率
if fps is None:
fps = int(cap.get(cv2.CAP_PROP_FPS))
# 初始化帧计数器和片段索引
frame_count = 0
segment_index = 0
# 定义视频编码器和输出视频参数
fourcc = cv2.VideoWriter_fourcc(*'mp4v')
while cap.isOpened():
ret, frame = cap.read()
if not ret:
break
# 如果是片段的第一帧，创建新的视频写入对象
if frame_count % frames_per_segment == 0:
if frame_count > 0:
out.release()
segment_filename = os.path.join(output_folder, f'{video_name}_{segment_index:03d}.mp4')
out = cv2.VideoWriter(segment_filename, fourcc, fps, (int(cap.get(3)), int(cap.get(4))))
segment_index += 1
# 将当前帧写入视频
out.write(frame)
frame_count += 1
# 释放视频捕获和写入对象
cap.release()
if 'out' in locals():
out.release()
