modules/video_utils.py

import subprocess
import os
from typing import List, Optional

from modules.logger_util import get_logger
from modules.paths import TEMP_DIR

logger = get_logger()


def extract_frames(
    vid_input: str,
    output_temp_dir: str = TEMP_DIR,
    quality: int = 2,
    start_number: int = 0
):
    """
    Extract frames and save them into output_temp_dir. This needs FFmpeg installed.
    """
    os.makedirs(output_temp_dir, exist_ok=True)
    output_path = os.path.join(output_temp_dir, "%05d.jpg")

    command = [
        'ffmpeg',
        '-i', vid_input,
        '-q:v', str(quality),
        '-start_number', str(start_number),
        f'{output_path}'
    ]

    try:
        subprocess.run(command, check=True)
    except subprocess.CalledProcessError as e:
        logger.exception("Error occured while extracting frames from the video")
        raise f"An error occurred: {str(e)}"


def get_frames_from_dir(vid_dir: str,
                        available_extensions: Optional[List, str] = None) -> List:
    if available_extensions is None:
        available_extensions = [".jpg", ".jpeg", ".JPG", ".JPEG"]

    if isinstance(available_extensions, str):
        available_extensions = [available_extensions]

    frame_names = [
        p for p in os.listdir(vid_dir)
        if os.path.splitext(p)[-1] in available_extensions
    ]
    frame_names.sort(key=lambda x: int(os.path.splitext(x)[0]))
    return frame_names