extract_keyframes.py

"""
This script extracts keyframes from videos or frames using a specified method.

Usage:
    python extract_keyframes.py input_dir output_dir [options]

Arguments:
    input_dir (str): Input directory containing videos or frames.
    output_dir (str): Output directory for saving keyframes.

Options:
    --frame_dataset: Flag to indicate if the input directory contains frame datasets.
    --fname_tmpl_input (str): Filename template for rawframe dataset. Default is "{:05}.jpg".
    --fname_tmpl_output (str): Filename template for extracted rawframes. Default is "{:05}.jpg".
    --method (str): Method for selecting keyframes. Choices are ["frame difference"]. Default is "frame difference".
    --num_workers (int): Number of worker threads for parallel processing. Default is 1.
"""

import argparse

import torch  # pylint: disable=unused-import

from src.keyframe_extractor.methods_factory import get_extractor
from src.keyframe_extractor.utils.file_handler import get_files


def parse_args():
    """
    Parse command line arguments.

    Returns:
        argparse.Namespace: Parsed command line arguments.
    """
    parser = argparse.ArgumentParser(
        description="Extract keyframes from videos"
    )

    # Positional arguments
    parser.add_argument(
        "input_dir",
        help="Input directory containing videos or frames",
    )
    parser.add_argument(
        "output_dir",
        help="Output directory for saving keyframes",
    )  ## TODO: # (append postfix "_keyframes" by default)

    # Optional arguments
    parser.add_argument(
        "--frame_dataset",
        action="store_true",
        help="Flag to indicate if the input directory contains frame datasets",
    )
    parser.add_argument(
        "--fname_tmpl_input",
        default="{:05}.jpg",
        type=str,
        help="filename_tmpl for rawframe dataset",
    )
    parser.add_argument(
        "--fname_tmpl_output",
        default="{:05}.jpg",
        type=str,
        help="filename_tmpl for extracted rawframes",
    )
    parser.add_argument(
        "--method",
        choices=["frame difference", "deep learning"],
        default="deep learning",
        help="Method for selecting keyframes",
    )
    parser.add_argument(
        "--num_workers",
        type=int,
        default=4,
        help="Number of worker threads for parallel processing",
    )
    parser.add_argument(
        "--output_type",
        choices=["frames", "video", "json"],
        default="csv",
        help="Output type. if 'json', save keyframe indices in a json file, otherwise save keyframes as images or video",
    )

    # Deep learning specifics
    parser.add_argument(
        "--num_keyframes",
        type=int,
        default=50,
        help="Number of keyframes to extract",
    )
    parser.add_argument(
        "--batch_size",
        type=int,
        default=16,
        help="Batch size for processing videos",
    )
    parser.add_argument(
        "--input_size",
        type=int,
        default=224,
        help="Input size (image) for deep learning method",
    )
    parser.add_argument(
        "--device",
        type=str,
        choices=["auto", "cpu", "gpu"],
        default="auto",
        help="Device for deep learning method. Default is 'auto' meaning GPU if avaiable.",
    )

    return parser.parse_args()


if __name__ == "__main__":
    args = parse_args()

    # Parse directory, creating a list of video files
    files = get_files(args.input_dir, args.frame_dataset)

    # Extract keyframes
    extractor = get_extractor(args)
    extractor.process(files)