【OpenCV】基于opencv的视频间隔抽帧脚本

1. 通过指定frame_interval来确定帧间隔；

2. 使用了joblib来加快处理速度。

代码如下所示：

import os
import cv2
import torchvision.transforms as T
from PIL import Image
from tqdm import tqdm
from torchvision.utils import save_image
from joblib import Parallel, delayed

preprocess = T.Compose([T.Resize(256), T.CenterCrop(256), T.ToTensor()])

def process_one_video(video_file, input_video_folder, target_fold, frame_interval=5):
    input_video_path = os.path.join(input_video_folder, video_file)
    video_id         = os.path.splitext(video_file)[0]

    vidcap           = cv2.VideoCapture(input_video_path)
    fps              = vidcap.get(cv2.CAP_PROP_FPS)
    total_frame_num  = int(vidcap.get(cv2.CAP_PROP_FRAME_COUNT))

    images_folder    = os.path.join(target_fold, video_id)
    os.makedirs(images_folder, exist_ok=True)

    for i in range(0, total_frame_num, frame_interval):
        vidcap.set(cv2.CAP_PROP_POS_FRAMES, i)
        _, frame     = vidcap.read()
        cv2_im       = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
        pil_im       = Image.fromarray(cv2_im)
        image_tensor = preprocess(pil_im)

        save_image(image_tensor, os.path.join(images_folder, '{:06d}.jpg'.format(i)))

    # Handle the last frames if they are less than frame_interval
    last_frame_index      = total_frame_num - 1
    
    if last_frame_index % frame_interval != 0:
        vidcap.set(cv2.CAP_PROP_POS_FRAMES, last_frame_index)
        _, last_frame     = vidcap.read()
        cv2_im_last       = cv2.cvtColor(last_frame, cv2.COLOR_BGR2RGB)
        pil_im_last       = Image.fromarray(cv2_im_last)
        image_tensor_last = preprocess(pil_im_last)
        save_image(image_tensor_last, os.path.join(images_folder, '{:06d}.jpg'.format(last_frame_index)))

    print("Video {} has been processed. Total frames: {}".format(video_id, total_frame_num))

#-----------------------------------#
# 对视频进行并行处理
#-----------------------------------#
def extract_frames_parallel(input_video_folder, target_fold, frame_interval=5):
    video_files = [f for f in os.listdir(input_video_folder) if f.endswith('.mp4')]

    Parallel(n_jobs=256)(delayed(process_one_video)(video_file, input_video_folder, target_fold, frame_interval) for video_file in tqdm(video_files))

if __name__ == "__main__":
    from argparse import ArgumentParser
    parser = ArgumentParser(description="Python script to extract frames from videos in a folder and save as jpgs.")
    parser.add_argument("-input_video_folder", type=str, default='/you/input/', help="The folder containing input videos.")
    parser.add_argument("-target_fold", type=str, default='/you/output/', help="The place to store the video frames.")
    parser.add_argument("-frame_interval", type=int, default=5, help="Interval between frames to extract.")
    args   = parser.parse_args()

    extract_frames_parallel(args.input_video_folder, args.target_fold, args.frame_interval)