Extract a sampled range of frames from a video using OpenCV

extract_frames_from_video.py

import os
from typing import Optional

import cv2
from tqdm import tqdm


def extract_frames_from_video(
        target_dir: str,
        video_path: str,
        start_frame: Optional[int]=0, # Let us skipping initial frames
        end_frame: Optional[int]=None,
        save_freq: Optional[int]=1,   # Let us sample the video in a lower frame-rate
        image_format: Optional[str]="jpg"
) -> None:
    """Extracts frames from a video using OpenCV

    Args:
        target_dir (str): Path to the target directory for extracted frames.
        video_path (str): Path to the input video file.
        start_frame (int, optional): Frame number to start extraction from
            (default: 0).
        end_frame (int, optional): Frame number to start extraction from
            (default: end).
        save_freq (int, optional): How many frames to skip between saving
            (default: 1).
        image_format (str, optional): Format for saving frames (e.g., "jpg",
            "png", default: "jpg").
    
    Raises:
        FileNotFoundError: If the input video file is not found.
        ValueError:
            - If the target directory already exists and contains files.
            - If the save_freq value is not a positive integer.
            - If the start_frame value is greater than the total frame count of the video.
        RuntimeError: If the input video cannot be opened.
    """

    # Check if input video file exists
    if not os.path.isfile(video_path):
        raise FileNotFoundError(f"Input video file not found: {video_path}")
    
    # Create output directory if it doesn't exist
    if not os.path.exists(target_dir):
        os.makedirs(target_dir)
    else:
        files = os.listdir(target_dir)
        if files:
            raise ValueError(f"Output directory {target_dir} already exists "
                             "and contains files. Choose a different directory.")

    # Validate save_freq
    if not isinstance(save_freq, int) or save_freq <= 0:
        raise ValueError("save_freq must be a positive integer.")
    
    # Open the video using OpenCV
    cap = cv2.VideoCapture(video_path)

    # Verify valid video capture
    if not cap.isOpened():
        raise RuntimeError(f"Could not open video file: {video_path}")

    # Get video properties
    fps = cap.get(cv2.CAP_PROP_FPS)
    total_frame_count = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))

    # Set the end_frame to total_frame_count if it is not specified
    if end_frame is None:
        end_frame = total_frame_count
   
    processing_frame_count = end_frame - start_frame

    # Validate start_frame and end_frame
    if processing_frame_count < 0:
        raise ValueError(f"Invalid frame range: start_frame={start_frame}, end_frame={end_frame}."
                         "The end_frame must be greater than or equal to the start_frame.")

    print(f"Video properties for {video_path}:")
    print(f"\t- FPS: {fps}")
    print(f"\t- Total frame count: {total_frame_count}")
    print(f"\t- Processing frame count: {processing_frame_count}")

    # Set starting frame
    cap.set(cv2.CAP_PROP_POS_FRAMES, start_frame)

    # Initialize progress bar
    pbar = tqdm(total = processing_frame_count, desc="Extracting frames")

    # Capture and extract frames
    frame_idx = start_frame
    while cap.isOpened() and frame_idx < end_frame:
        ret, frame = cap.read()

        if not ret:
            # frame is not returned due to an error or the end of the video
            break

        if (frame_idx - start_frame) % save_freq == 0:
            # save the frame as an image file
            file_name = f"frame_{frame_idx}.{image_format}"
            try:
                cv2.imwrite(os.path.join(target_dir, file_name), frame)
            except Exception as e:
                print(f"Frame {frame_idx} is not saved:\n{e}")
            
        pbar.update()
        frame_idx += 1

    # Release resources
    cap.release()
    cv2.destroyAllWindows() # OpenCV might create invisible internal windows
    pbar.close()