ffmpeg_demo.py

# -*- coding: utf-8 -*-
'''Simple ffmpeg demo

Presented as a single file here, but realistically needs to be split
in to various modules as it best fits within the host project
'''

from __future__ import annotations

import subprocess
from dataclasses import dataclass
from enum import Enum, auto


FFMPEGEXE = '/Users/danielskovli/ffmpeg-4.4/ffmpeg'


# Models

class OutputFormat(Enum):
    mp4 = auto()
    mov = auto()

class ColorSpace:
    rec709 = 'bt709'
    gamma = 'gamma'
    gamma22 = 'gamma22'
    gamma28 = 'gamma28'
    linear = 'linear'
    sRgb = 'iec61966_2_1'

@dataclass()
class CompJob:
    inputFilePattern: str
    outputFilePath: str
    outputFormat: OutputFormat
    user: str
    show: str
    commitNotes: str
    framerate: float = 24
    startFrame: int = 0
    trc: str = None
    timecode: str = None

@dataclass()
class CompResult:
    success: bool = False
    output: str = None


# Factories

class FfmpegParamsFactory:
    '''Produce `ffmpeg` command params for the supplied `CompJob`'''

    # TODO: This belongs in a config class somewhere
    codecs = {
        OutputFormat.mp4: 'libx264 -pix_fmt yuv420p -preset slow -crf 20',
        OutputFormat.mov: 'prores_ks -profile:v 3 -vendor ap10 -c:a pcm_s16le -pix_fmt yuv422p10le'
    }
    forcedScales = {
        OutputFormat.mp4: 'scale=1920:1080:force_original_aspect_ratio=decrease,pad=1920:1080:(ow-iw)/2:(oh-ih)/2,setdar=dar=16/9'
    }


    def __init__(self, job: CompJob) -> None:
        self.job = job

    def getCodec(self) -> list[str]:
        codec = self.codecs[self.job.outputFormat] # Throws on error
        return ['-c:v'] + codec.split(' ')

    def getScale(self) -> str|tuple[str,str]:
        # FIXME: -vf filter is horribly outdated. Replace with -filter_complex or similar
        scale = self.forcedScales.get(self.job.outputFormat)
        if not scale:
            return ''

        return ('-vf', scale)

    def getInputFile(self) -> tuple[str,str]:
        return ('-i', self.job.inputFilePattern)

    def getFramerate(self) -> tuple[str,str]:
        return ('-framerate', f'{self.job.framerate}')

    def getStartFrame(self) -> tuple[str,str]:
        return ('-start_number', f'{job.startFrame}')

    def getTrc(self, includeFlag: bool=True) -> str|tuple[str,str]:
        if self.job.trc:
            return ('-apply_trc', self.job.trc)

        return ''

    def getTimecode(self, includeFlag: bool=True) -> tuple[str,str]:
        return ('-timecode', self.job.timecode or '00:00:00:00')

    def getMetadata(self) -> tuple[str,...]:
        # TODO: All of these injections need sanitizing
        metadata = (
            '-metadata', f'comment="{self.job.commitNotes}"',
            '-metadata', f'artist="{self.job.user}"',
            '-metadata', f'show="{self.job.show}"',
            '-metadata', 'publisher="YOUR COMPANY HERE"',
            '-metadata', 'copyright="YOUR COMPANY HERE"',
        )

        return metadata


# Methods

def compExrSequence(job: CompJob, overwrite: bool=False) -> CompResult:
    '''Set up and execute `ffmpeg` job as per incoming `CompJob`'''

    params = FfmpegParamsFactory(job)
    command = [
        '-y' if overwrite else '-n',
        *params.getTrc(),
        *params.getFramerate(),
        *params.getStartFrame(),
        *params.getInputFile(),
        *params.getCodec(),
        *params.getScale(),
        *params.getMetadata(),
        *params.getTimecode(),
        job.outputFilePath
    ]

    return _execute(command)


def _execute(args: list[str]):
    '''Execute `ffmpeg` system command with incoming args'''

    success: bool = False
    stdout: bytes = b''
    stderr: bytes = b''

    if not isinstance(args, (list, tuple)):
        raise TypeError(f'Invalid command list: {args}')

    try:
        print(f'Calling ffmpeg with args: {args}')
        args = [FFMPEGEXE] + [x for x in args if x] # Inject ffmpeg binaries and strip empty params
        proc = subprocess.Popen(args, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
        stdout, stderr = proc.communicate()
        success = proc.returncode == 0
    except Exception as e:
        print(f'Error encountered while processing ffmpeg command: {e}')

    # NOTE: ffmpeg communicates human-readable stuff on stderr instead of stdout, for whatever reason
    return CompResult(
        success=success,
        output=stderr.decode('utf-8')
    )



# Demo usage
if __name__ == '__main__':
    job = CompJob(
        inputFilePattern='/Users/danielskovli/Movies/imagesequence_test/frame_%04d.png',
        outputFilePath='/Users/danielskovli/Movies/frame_comp_test.mov',
        outputFormat=OutputFormat.mov,
        user='Daniel',
        show='Fancy Episodic Project',
        commitNotes='Cowbell increased by 25%',
        trc=ColorSpace.linear
    )

    result = compExrSequence(job, overwrite=True)
    if result.success:
        print('Sequence rendered successfully!')
    else:
        print(f'ffmpeg reports errors: \n{result.output}')