视频抽帧工具：按需提取高质量数据集

引言

在视频数据处理中，抽帧是一项非常重要的任务。针对某些场景，我们需要在特定时间段提高抽帧频率（以捕获更多目标画面），而在其他时间段保持较低抽帧频率（以减少数据冗余并获取背景画面）。本文将分享一段代码，能够对指定视频按需抽帧，既保证了目标画面的多样性，又获取了一定量的背景数据。

功能介绍

这段代码主要实现了以下功能：

1. 支持自定义时间段的抽帧频率： 可以针对视频的不同时间段，设置不同的抽帧间隔。
2. 默认抽帧频率设置： 未指定时间段的帧，按照默认频率抽取。
3. 抽帧进度实时可见： 使用 tqdm 实现进度条显示，便于监控处理进度。
4. 抽取结果自动保存： 抽取的帧会以 .jpg 格式保存到指定文件夹中，文件名包含视频名称及帧编号。

使用场景

• 目标检测数据集制作： 在视频中捕获目标密集区更多帧，以提升目标检测模型的训练效果。
• 背景数据采集： 在目标稀疏或无目标的时段，按较低频率采集背景帧，减少数据冗余。
• 视频分析与处理： 提取关键帧用于后续视频内容分析。

代码实现

import os
import cv2
import math
from tqdm import tqdm


def extract_frames_from_single_video(video_path, output_folder, time_interval_frame_interval_dict, default_frame_interval):
    cap = cv2.VideoCapture(video_path)  # 打开视频文件
    if not cap.isOpened():  # 检查视频是否成功打开
        print(f"无法打开视频文件 {video_path}")
        return

    fps = math.ceil(cap.get(cv2.CAP_PROP_FPS))  # 获取并向上取整视频的fps

    # 设置计数器和帧数
    frame_count = 0
    frame_num = 0

    # 初始化当前抽帧间隔为默认间隔
    current_frame_interval = default_frame_interval
    total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))  # 计算视频总帧数
    video_name = os.path.splitext(os.path.basename(video_path))[0]

    with tqdm(total=total_frames, desc=f"Processing {video_name}") as pbar:
        while True:  # 对于每一帧
            ret, frame = cap.read()  # 读取视频帧
            if not ret:  # 判断是否到达视频末尾
                break

            current_time = frame_count // fps  # 计算当前帧在第几秒

            for interval, frame_interval in time_interval_frame_interval_dict.items():
                if interval[0] <= current_time < interval[1]:  # 如果当前时间在特定的范围内
                    current_frame_interval = frame_interval  # 更改抽帧间隔
                    break
            else:  # 如果没有在任何特定的范围内，使用默认的抽帧间隔
                current_frame_interval = default_frame_interval

            frame_count += 1
            # 根据抽帧间隔选择帧进行保存
            if frame_count % current_frame_interval == 0:  
                frame_num += 1
                frame_name = f"{video_name}_{frame_num}.jpg"
                frame_path = os.path.join(output_folder, frame_name)
                cv2.imwrite(frame_path, frame)
            
            pbar.update()  # 更新进度条

    cap.release()  # 释放视频对象


if __name__ == "__main__":
    # 视频文件夹的路径，输出文件夹的路径，视频文件名
    video_path = '/data/dataset/xxx.mp4'
    output_folder = '/data/dataset/img'

    # 特定时间段与抽帧频率的对应关系
    time_interval_frame_interval_dict = {(90, 115): 20,
                                        (116, 132): 10}  # 90-115秒，每20帧抽一张；116-132秒，每10帧抽一张

    # 默认的抽帧频率
    default_frame_interval = 30  # 其余时间每30帧抽一张
    extract_frames_from_single_video(video_path, output_folder, time_interval_frame_interval_dict, default_frame_interval)

参数说明

• video_path： 视频文件路径。
• output_folder： 抽帧结果保存路径。
• time_interval_frame_interval_dict： 特定时间段与抽帧频率的映射关系。例如，{(90, 115): 20} 表示 90 到 115 秒内，每 20 帧保存一张图片。
• default_frame_interval： 默认抽帧间隔，用于未指定时间段。

以上代码简单高效，适用于需要处理大量视频并生成数据集的场景。如果你觉得有帮助，欢迎点赞、评论或分享！