ffmpeg encoding & processing script through docker image

docker_ffmpeg_script.py

#!/usr/bin/env python3
# -*- coding: utf-8 -*-

"""
ffmpeg encoding & processing script through docker image: see https://github.com/jrottenberg/ffmpeg
"""
__author__ = 'Rodrigo Gonzalez del Cueto'
__copyright__ = 'Copyright 2020, ffmpeg processor'
__credits__ = ['Rodrigo Gonzalez del Cueto']
__license__ = '{GNU Lesser General Public License version 3}'
__version__ = '0.1.5'
__email__ = 'rdelcueto@gmail.com'

import os
import sys
import subprocess

import itertools
from collections import OrderedDict

import toml
import argparse
import glob

# Script debugging variables
debug = True
debugprint = print if debug else lambda *a, **k: None

# Script default configuration
default_config_toml = """
# Default configuration TOML file
title = "Docker ffmpeg automation script configuration"

[docker]
  uid = 1000
  gid = 65536
  device-bind = "--device /dev/dri:/dev/dri"
  cpu-config = "--cpuset-cpus=0,1"
  mem-config = "-m 1g"
  image-tag = "jrottenberg/ffmpeg:snapshot-vaapi"

[ffmpeg.hw-video]
  codec = "hevc_vaapi"
  cqp = 25
  #scale_w = 1280
  #scale_h = 720

[ffmpeg.video]
  codec = "libx265"
  tune-params = "fastdecode -preset fast"
  crf = 22
  #scale_w = 1280
  #scale_h = 720

[ffmpeg.stabilization.vidstabdetect]
  shakiness = 8
  accuracy = 10
  stepsize = 4

[ffmpeg.stabilization.vidstabtransform]
  smoothing = 6
  optalgo = "gauss"
  maxshift = -1
  maxangle = -1
  crop = "keep"
  optzoom = 2
  zoomspeed = 0.2
  interpol = "bicubic"

[ffmpeg.stabilization.unsharp]
  luma_msize_x = 3
  luma_msize_y = 3
  luma_amount = 0.8
  chroma_msize_x = 3
  chroma_msize_y = 3
  chroma_amount = 0.4

[ffmpeg.audio]
  #codec = "copy"
  #bitrate-param = ""
  codec = "libvorbis"
  bitrate-param = "-qscale:a 4"

[ffmpeg.output]
  container = "mp4"
"""

def parse_config (args):
    if args.config_file is None:
        parsed_toml = toml.loads(default_config_toml)
    else:
        if not os.path.isfile(args.config_file):
            debugprint ("Config file doesn't exist")
            return None
        else:
            debugprint ("Parsing TOML config file {}".format(args.config_file))
            parsed_toml = toml.load(args.config_file)

    return parsed_toml

def execute_docker_process (args, config, ffmpeg_cmd, cwd, environment):

    cmd = ("docker run --user {}:{}".format(config['docker']['uid'], config['docker']['gid']) +
           " {}".format(config['docker']['device-bind']) +
           " {}".format(config['docker']['cpu-config']) +
           " {}".format(config['docker']['mem-config']) +
           " -v {}:{} -w {}".format(cwd, cwd, cwd) +
           " {}".format(config['docker']['image-tag']) +
           " {}".format(ffmpeg_cmd))

    if args.verbose is True:
        print (cmd)

    if args.dry_run is False:
        cmd = cmd.split ()
        docker_ffmpeg_process = subprocess.Popen (cmd, env=environment, stdout=subprocess.PIPE)

        out, err = docker_ffmpeg_process.communicate ()
        print (out)

def process_input_files (args, config, input_files):
    # Set PATH
    environment = os.environ.copy ()
    environment["PATH"] = "/usr/sbin:/usr/bin:" + environment["PATH"]

    cwd = os.getcwd ()

    for file in input_files:
        print ("Processing {}...".format(file))

        if args.stabilize_video is False:
            if args.hwaccel == 'off':
                # Single Pass
                ffmpeg_cmd = get_single_pass_cmd (args, config, file, args.output_suffix)
                execute_docker_process (args, config, ffmpeg_cmd, cwd, environment)
            else:
                # Single Pass with Hardware acceleration decoding & encoding
                ffmpeg_cmd = get_vaapi_single_pass_cmd (args, config, file, args.output_suffix)
                execute_docker_process (args, config, ffmpeg_cmd, cwd, environment)
        else:
            if args.hwaccel == 'off':
                # First Pass
                ffmpeg_cmd = get_stabilized_1st_pass_cmd (args, config, file)
                execute_docker_process (args, config, ffmpeg_cmd, cwd, environment)

                # Second Pass
                ffmpeg_cmd = get_stabilized_2nd_pass_cmd (args, config, file, args.output_suffix + "_stabilized")
                execute_docker_process (args, config, ffmpeg_cmd, cwd, environment)
            else:
                # First Pass with Hardware acceleration decoding & encoding
                ffmpeg_cmd = get_vaapi_stabilized_1st_pass_cmd (args, config, file)
                execute_docker_process (args, config, ffmpeg_cmd, cwd, environment)

                # Second Pass with Hardware acceleration decoding & encoding
                ffmpeg_cmd = get_vaapi_stabilized_2nd_pass_cmd (args, config, file, args.output_suffix + "_stabilized")
                execute_docker_process (args, config, ffmpeg_cmd, cwd, environment)

def get_single_pass_cmd (args, config, input_file, output_suffix):
    input_filename_with_extension = os.path.basename(input_file)
    output_filename = os.path.splitext(input_filename_with_extension)[0]
    video_filters = []

    if config['ffmpeg']['video'].get('scale_w') and config['ffmpeg']['video'].get('scale_h'):
        scale = "scale=w={}:h={}".\
            format(config['ffmpeg']['video']['scale_w'], config['ffmpeg']['video']['scale_h'])
        video_filters.append (scale)

    video_filters = ','.join(video_filters) if video_filters else None

    cmd = (" -i {}".format(input_file) +
           " -c:v {}".format(config['ffmpeg']['video']['codec']) +
           ((" -vf " + video_filters) if video_filters else "") +
           " -tune {}".format(config['ffmpeg']['video']['tune-params']) +
           " -crf {}".format(config['ffmpeg']['video']['crf']) +
           " -c:a {}".format(config['ffmpeg']['audio']['codec']) +
           " {}".format(config['ffmpeg']['audio']['bitrate-param']) +
           (" -y " if args.force_overwrite else " -n ") +
           output_filename + output_suffix + '.' + config['ffmpeg']['output']['container'])

    return cmd

def get_stabilized_1st_pass_cmd (args, config, input_file):
    input_filename_with_extension = os.path.basename(input_file)
    output_filename = os.path.splitext(input_filename_with_extension)[0]
    video_filters = []

    if config['ffmpeg']['video'].get('scale_w') and config['ffmpeg']['video'].get('scale_h'):
        scale = "scale=w={}:h={}".\
            format(config['ffmpeg']['video']['scale_w'], config['ffmpeg']['video']['scale_h'])
        video_filters.append (scale)

    vidstabdetect = "vidstabdetect=shakiness={}:accuracy={}:stepsize={}:result={}".\
           format(config['ffmpeg']['stabilization']['vidstabdetect']['shakiness'],
                  config['ffmpeg']['stabilization']['vidstabdetect']['accuracy'],
                  config['ffmpeg']['stabilization']['vidstabdetect']['stepsize'],
                  (output_filename + '_vidstabdetect.trf'))

    video_filters.append (vidstabdetect)

    video_filters = ','.join(video_filters) if video_filters else None

    cmd = (" -i {}".format(input_file) +
           ((" -vf " + video_filters) if video_filters else "") +
           " -an -f null -")

    return cmd

def get_stabilized_2nd_pass_cmd (args, config, input_file, output_suffix):
    input_filename_with_extension = os.path.basename(input_file)
    output_filename = os.path.splitext(input_filename_with_extension)[0]
    video_filters = []

    if config['ffmpeg']['video'].get('scale_w') and config['ffmpeg']['video'].get('scale_h'):
        scale = "scale=w={}:h={}".\
            format(config['ffmpeg']['video']['scale_w'], config['ffmpeg']['video']['scale_h'])
        video_filters.append (scale)

    vidstabtransform = "vidstabtransform=smoothing={}:optalgo={}:maxshift={}:maxangle={}:crop={}:optzoom={}:zoomspeed={}:interpol={}:input={}".\
        format(config['ffmpeg']['stabilization']['vidstabtransform']['smoothing'],
               config['ffmpeg']['stabilization']['vidstabtransform']['optalgo'],
               config['ffmpeg']['stabilization']['vidstabtransform']['maxshift'],
               config['ffmpeg']['stabilization']['vidstabtransform']['maxangle'],
               config['ffmpeg']['stabilization']['vidstabtransform']['crop'],
               config['ffmpeg']['stabilization']['vidstabtransform']['optzoom'],
               config['ffmpeg']['stabilization']['vidstabtransform']['zoomspeed'],
               config['ffmpeg']['stabilization']['vidstabtransform']['interpol'],
               (output_filename + '_vidstabdetect.trf'))

    video_filters.append (vidstabtransform)

    unsharp = "unsharp={}:{}:{}:{}:{}:{}".\
        format(config['ffmpeg']['stabilization']['unsharp']['luma_msize_x'],
               config['ffmpeg']['stabilization']['unsharp']['luma_msize_y'],
               config['ffmpeg']['stabilization']['unsharp']['luma_amount'],
               config['ffmpeg']['stabilization']['unsharp']['chroma_msize_x'],
               config['ffmpeg']['stabilization']['unsharp']['chroma_msize_y'],
               config['ffmpeg']['stabilization']['unsharp']['chroma_amount'])

    video_filters.append (unsharp)

    video_filters = ','.join(video_filters) if video_filters else None

    cmd = (" -i {}".format(input_file) +
           ((" -vf " + video_filters) if video_filters else "") +
           " -c:v {}".format(config['ffmpeg']['video']['codec']) +
           " -tune {}".format(config['ffmpeg']['video']['tune-params']) +
           " -crf {}".format(config['ffmpeg']['video']['crf']) +
           " -c:a {}".format(config['ffmpeg']['audio']['codec']) +
           " {}".format(config['ffmpeg']['audio']['bitrate-param']) +
           (" -y " if args.force_overwrite else " -n ") +
           output_filename + output_suffix + '.' + config['ffmpeg']['output']['container'])

    return cmd

def get_vaapi_single_pass_cmd (args, config, input_file, output_suffix):
    input_filename_with_extension = os.path.basename(input_file)
    output_filename = os.path.splitext(input_filename_with_extension)[0]
    video_filters = []

    if config['ffmpeg']['hw-video'].get('scale_w') and config['ffmpeg']['hw-video'].get('scale_h'):
        scale = "scale_vaapi=w={}:h={}".\
            format(config['ffmpeg']['hw-video']['scale_w'], config['ffmpeg']['hw-video']['scale_h'])
        video_filters.append (scale)

    video_filters = ','.join(video_filters) if video_filters else None

    cmd = (" -hwaccel vaapi -hwaccel_output_format vaapi" +
           " -i {}".format(input_file) +
           ((" -vf " + video_filters) if video_filters else "") +
           " -c:v {}".format(config['ffmpeg']['hw-video']['codec']) +
           " -qp {}".format(config['ffmpeg']['hw-video']['cqp']) +
           " -c:a {}".format(config['ffmpeg']['audio']['codec']) +
           " {}".format(config['ffmpeg']['audio']['bitrate-param']) +
           (" -y " if args.force_overwrite else " -n ") +
           output_filename + output_suffix + '.' + config['ffmpeg']['output']['container'])

    return cmd

def get_vaapi_stabilized_1st_pass_cmd (args, config, input_file):
    input_filename_with_extension = os.path.basename(input_file)
    output_filename = os.path.splitext(input_filename_with_extension)[0]
    video_filters = []

    if config['ffmpeg']['hw-video'].get('scale_w') and config['ffmpeg']['hw-video'].get('scale_h'):
        scale = "scale_vaapi=w={}:h={}".\
            format(config['ffmpeg']['hw-video']['scale_w'], config['ffmpeg']['hw-video']['scale_h'])
        video_filters.append (scale)

    video_filters.append ("hwdownload")
    video_filters.append ("format=nv12")

    vidstabdetect = "vidstabdetect=shakiness={}:accuracy={}:stepsize={}:result={}".\
           format(config['ffmpeg']['stabilization']['vidstabdetect']['shakiness'],
                  config['ffmpeg']['stabilization']['vidstabdetect']['accuracy'],
                  config['ffmpeg']['stabilization']['vidstabdetect']['stepsize'],
                  (output_filename + '_vidstabdetect.trf'))

    video_filters.append (vidstabdetect)

    video_filters = ','.join(video_filters) if video_filters else None

    cmd = (" -hwaccel vaapi -hwaccel_output_format vaapi" +
           " -i {}".format(input_file) +
           ((" -vf " + video_filters) if video_filters else "") +
           " -an -f null -")

    return cmd

def get_vaapi_stabilized_2nd_pass_cmd (args, config, input_file, output_suffix):
    input_filename_with_extension = os.path.basename(input_file)
    output_filename = os.path.splitext(input_filename_with_extension)[0]
    video_filters = []

    if config['ffmpeg']['hw-video'].get('scale_w') and config['ffmpeg']['hw-video'].get('scale_h'):
        scale = "scale_vaapi=w={}:h={}".\
            format(config['ffmpeg']['hw-video']['scale_w'], config['ffmpeg']['hw-video']['scale_h'])
        video_filters.append (scale)

    video_filters.append ("hwdownload")
    video_filters.append ("format=nv12")

    vidstabtransform = "vidstabtransform=smoothing={}:optalgo={}:maxshift={}:maxangle={}:crop={}:optzoom={}:zoomspeed={}:interpol={}:input={}".\
        format(config['ffmpeg']['stabilization']['vidstabtransform']['smoothing'],
               config['ffmpeg']['stabilization']['vidstabtransform']['optalgo'],
               config['ffmpeg']['stabilization']['vidstabtransform']['maxshift'],
               config['ffmpeg']['stabilization']['vidstabtransform']['maxangle'],
               config['ffmpeg']['stabilization']['vidstabtransform']['crop'],
               config['ffmpeg']['stabilization']['vidstabtransform']['optzoom'],
               config['ffmpeg']['stabilization']['vidstabtransform']['zoomspeed'],
               config['ffmpeg']['stabilization']['vidstabtransform']['interpol'],
               (output_filename + '_vidstabdetect.trf'))

    video_filters.append (vidstabtransform)

    unsharp = "unsharp={}:{}:{}:{}:{}:{}".\
        format(config['ffmpeg']['stabilization']['unsharp']['luma_msize_x'],
               config['ffmpeg']['stabilization']['unsharp']['luma_msize_y'],
               config['ffmpeg']['stabilization']['unsharp']['luma_amount'],
               config['ffmpeg']['stabilization']['unsharp']['chroma_msize_x'],
               config['ffmpeg']['stabilization']['unsharp']['chroma_msize_y'],
               config['ffmpeg']['stabilization']['unsharp']['chroma_amount'])

    video_filters.append (unsharp)
    video_filters.append ("hwupload")

    video_filters = ','.join(video_filters) if video_filters else None

    cmd = (" -hwaccel vaapi -hwaccel_output_format vaapi" +
           " -i {}".format(input_file) +
           ((" -vf " + video_filters) if video_filters else "") +
           " -c:v {}".format(config['ffmpeg']['hw-video']['codec']) +
           " -qp {}".format(config['ffmpeg']['hw-video']['cqp']) +
           " -c:a {}".format(config['ffmpeg']['audio']['codec']) +
           " {}".format(config['ffmpeg']['audio']['bitrate-param']) +
           (" -y " if args.force_overwrite else " -n ") +
           output_filename + output_suffix + '.' + config['ffmpeg']['output']['container'])

    return cmd

def parse_input_files (args):
    if args.input_directory is None and args.input_files == [None]:
        print ("Error: No valid input defined!\n", file=sys.stderr)
        return None

    input_files_from_directories = []
    if args.input_directory is not None:
        glob_extensions = [(lambda x: "/*.{}".format(x))(file_type) for file_type in args.input_file_type]
        input_files_from_directories.extend(glob.glob(directory + ext) for ext in glob_extensions for directory in args.input_directory)
        input_files_from_directories = list (itertools.chain.from_iterable(input_files_from_directories))

    if args.input_files is not None:
        input_files = args.input_files
    else:
        input_files = []

    # Merge
    input_files = list (OrderedDict.fromkeys (input_files + input_files_from_directories))
    if args.verbose is not None:
        print (input_files)

    return input_files

def check_input_file (file):
    if not os.path.isfile(file):
        print("ERROR - Input file does not exist: {}".format(file), file=sys.stderr)
        return None
    return file

def check_input_directory (dir):
    if not os.path.isdir(dir):
        print("ERROR - Input directory does not exist: {}".format(dir), file=sys.stderr)
        return None

    return dir

def main():
    # Directory
    parser = argparse.ArgumentParser (description='Encode & process video files using ffmpeg from a docker container image')

    parser.add_argument ('-i', '--input-files',
                         nargs='?', action='append',
                         help="files to process", type=check_input_file)

    parser.add_argument ('-d', '--input-directory',
                         nargs='?', action='append',
                         help="directories to process", type=check_input_directory)

    parser.add_argument ('-t', '--input-file-type',
                         nargs='?', action='append',
                         default = ["mp4"],
                         help="file types to process. Default: mp4")

    parser.add_argument ('--output-suffix',
                         default="_converted",
                         help="output file suffix. Default = _converted")

    parser.add_argument ('-A', '--hwaccel',
                         choices=['off', 'vaapi'],
                         default='vaapi',
                         help="Use specified hardware acceleration for decoding and encoding. Default = vaapi")

    parser.add_argument ('-S', '--stabilize-video',
                         action='store_true',
                         default=False,
                         help="Analyze video stabilization/deshaking. Perform pass 1 with vidstabdetect filter, and vidstabtransform filter for pass 2. ")

    parser.add_argument ('-C', '--config-file',
                         default=None,
                         help="Config file to define ffmpeg parameters")

    parser.add_argument ('-n', '--dry-run',
                         action='store_true',
                         default=False,
                         help="perform a trial run with no processing of video files")

    parser.add_argument ('-v', '--verbose',
                         action='store_true',
                         default=False,
                         help="increase verbosity")

    parser.add_argument ('-f', '--force-overwrite',
                         action='store_true',
                         default=False,
                         help="overwrite output files")

    args = parser.parse_args ()
    debugprint (args)

    config = parse_config (args)
    if config is None:
        print ("Error reading config file")
        sys.exit (1)

    input_files = parse_input_files (args)

    if input_files is None:
        print ("Error: No input files to process!")
        sys.exit (1)

    process_input_files (args, config, input_files)
    sys.exit (0)

if __name__ == "__main__":
    main ()