itproject-code/pipeline/video_producers.py

"""Classes for producing videos"""

from abc import ABC
import logging
import subprocess
import tempfile

class VideoProducer(ABC):
    """Video producer interface."""
    def __init__(self, features):
        pass
    def produce(self):
        pass


class FfmpegVideoProducer(VideoProducer):
    """Produce videos using ffmpeg"""
    # TODO: consider output filename options

    def __init__(self, features):
        if not features:
            raise ValueError("No features provided")
        # TODO: consider if we want to permit empty features (producing no video)
        self.features = features

    def _ffmpeg_feature_to_clip(self, feature=None, output_filepath=None):
        """use ffmpeg to produve a video clip from a feature"""
        OVERWRITE = True # TODO: consider making this a config option
        if not feature or not feature.interval:
            raise ValueError("No feature provided")

        if not output_filepath:
            raise ValueError("No output filepath provided")

        ffmpeg_prefix = ["ffmpeg", "-y"] if OVERWRITE else ["ffmpeg"]
        ffmpeg_suffix = ["-r", "60", "-c:v", "libx264", "-crf", "26", "-c:a", "aac", "-preset", "ultrafast"]
        # TODO: match framerate of input video
        # TODO: adjustable encoding options
        seek = ["-ss", str(feature.interval.start)]
        duration = ["-t", str(feature.interval.duration)]
        ffmpeg_args = ffmpeg_prefix + seek + ["-i"] + [feature.path] + duration + ffmpeg_suffix + [output_filepath]
        logging.info(f"ffmpeg_args: {ffmpeg_args}")
        subprocess.run(ffmpeg_args, stdout=None, stderr=None)

    def _ffmpeg_concat_clips(self, clips=None, output_filepath=None):
        """use ffmpeg to concatenate clips into a single video"""
        OVERWRITE = True
        ffmpeg_prefix = ["ffmpeg"]
        ffmpeg_prefix += ["-y"] if OVERWRITE else []
        ffmpeg_prefix += ["-f", "concat", "-safe", "0", "-i"]

        # there is a method to do this via process substitution, but it's not portable
        # so we'll use the input file list method

        if not clips:
            raise ValueError("No clips provided")

        if not output_filepath:
            raise ValueError("No output filepath provided")

        # generate a temporary file with the list of clips
        join_file = tempfile.NamedTemporaryFile(mode="w")
        for clip in clips:
            join_file.write(f"file '{clip}'\n")
        join_file.flush()

        ffmpeg_args = ffmpeg_prefix + [join_file.name] + ["-c", "copy", output_filepath]
        logging.info(f"ffmpeg_args: {ffmpeg_args}")
        subprocess.run(ffmpeg_args, stdout=None, stderr=None)
        join_file.close()

    def produce(self):
        OUTPUT_DIR = "/tmp/" # TODO: make this a config option

        clips = []
        for num, feature in enumerate(self.features):
            output_filepath = f"{OUTPUT_DIR}/highlight_{num}.mp4"
            self._ffmpeg_feature_to_clip(feature, output_filepath)
            clips.append(output_filepath)

        # concatenate the clips
        output_filepath = f"{OUTPUT_DIR}/highlights.mp4"
        self._ffmpeg_concat_clips(clips, output_filepath)
        logging.info(f"Produced video: {output_filepath}")