Extract frame from a video using ffmpeg and txt file containing a list of timestamps expressed in seconds. One line is a timestamp.

frame_extractor.py

"""
Generates video screenshot from a text file of timestamps expressed in seconds.

Note: ffmpeg must be installed on your machine.

Tested on python3.7

Example usage:

```
$ python3 frame_extractor.py --video input.mp4 --output_dir frames/ --timestamps ts.txt
```

"""

import argparse
import subprocess
import pathlib
import logging
import sys

logger = logging.Logger(name='FrameExtractor', level='INFO')
sh = logging.StreamHandler(sys.stdout)
formatter = logging.Formatter('[%(asctime)s] %(levelname)s [%(filename)s.%(funcName)s:%(lineno)d] %(message)s', datefmt='%a, %d %b %Y %H:%M:%S')
sh.setFormatter(formatter)
logger.addHandler(sh)


def main(args) -> None:
    # 0. Sanity checks on file paths to read from / write to
    video_path = pathlib.Path(args.video)
    if not video_path.exists():
        print(f'Provide a valid video path, {video_path} is invalid')
        return
    timestamps_path = pathlib.Path(args.timestamps)
    if not timestamps_path.exists():
        print(f'Provide a valid timestamp file, {timestamps_path} is invalid')
        return
    output_dir = pathlib.Path(args.output_dir)
    if not output_dir.exists():
        print(f'Provide a valid output directory, {output_dir} is invalid')
        return

    # 1. read the timestamp file
    with timestamps_path.open(mode='r') as f:
        for line in f.readlines():
            ts: int = int(line.strip('\n'))
            img_file_name: str = f'{ts}.jpg'
            img_file_path: pathlib.Path = output_dir.joinpath(img_file_name)
            logger.info(f'Reading frame at {ts}s, writing to {img_file_path}')
            capture_frame(video_path, ts, img_file_path)
    return


def capture_frame(video_path: pathlib.Path, ts: int, img_path: pathlib.Path) -> None:
    """
    Capture the frame at second ts in the input video_path and write it to img_path.

    Args:
        video_path (pathlib.Path): path to the video to read from.
        ts (int): timestnap, in second of the frame to capture
        img_path (pathlib.Path): path to the image to write the frame to.

    Returns:
        None. This function side effect is to write an image to disk.
    """
    ffmpeg_cmd: str = f'ffmpeg -ss {ts} -i "{video_path.as_posix()}" -vframes 1 -q:v 2 "{img_path.as_posix()}" -hide_banner -loglevel error -y'
    _ = subprocess.check_output(
        ffmpeg_cmd,
        shell=True,
        stderr=subprocess.STDOUT
    )


if __name__ == '__main__':
    parser = argparse.ArgumentParser(description="Generate frames from a file of timestamps expressed in seconds")
    parser.add_argument('--timestamps', type=str, default="timestamps.txt",
                        help="Path to the input test file containing the timestamps, one timestamp per line, "
                             "each timestamp is an integer.")
    parser.add_argument('--output_dir', type=str, default='frames/',
                        help='Where to write the generated images')
    parser.add_argument('--video', type=str, default='input_video.mp4',
                        help="Video to capture the frames from.")
    args = parser.parse_args()
    main(args)