Deadline Monitor job script to convert image sequences to Apple ProRes or H.264 formats

ConvertToProResH264.py

# FFmpeg: Convert to ProRes/H.264 in Deadline Monitor
#
# features:
# - quickly convert an image sequence to Apple ProRes (422 HQ, 4444 and 4444 XQ) or H.264 formats
# - right-click Job(s) > Scripts > Convert To ProRes/H.264
# - conversion job is dependent on source job, will be queued until all frames are rendered
# - option to burn-in the frame number at the bottom-center of each frame
#
# requirements: 
# - save this script to: [repo path]/custom/scripts/Jobs
# - FFmpeg must be installed and pathed in Deadline via Tools > Configure Plugins > FFmpeg
# - enable Multi-Select and customize script name/icon via Tools > Configure Script Menus > Job Scripts > Edit Selection
# - "ffmpeg" limit must be created, workers assigned

import os
import re
from DeadlineUI.Controls.Scripting.DeadlineScriptDialog import DeadlineScriptDialog
from Deadline.Scripting import MonitorUtils, RepositoryUtils, ClientUtils
from System.IO import Path, File
from collections import defaultdict

scriptDialog = None

class CodecSettings:
    def __init__(self, pixfmt, profileV, meta, vtag):
        self.pixfmt = pixfmt; self.profileV = profileV; self.meta = meta; self.vtag = vtag

# Dictionary of compression settings, indexed by the codec selection
codecSettingsDict = {
    "prores422": CodecSettings("yuv422p10le", 3, "encoder=Apple ProRes 422 HQ", "apch"),
    "prores4444": CodecSettings("yuva444p10le", 4, "encoder=Apple ProRes 4444", "ap4h"),
    "prores4444xq": CodecSettings("yuva444p10le", 5, "encoder=Apple ProRes 4444 XQ", "ap4x"),
    "h264": CodecSettings("yuv420p", "main", "encoder=avc1", "")
}

def __main__(*args):
    """ dialog for user input """

    global scriptDialog

    scriptDialog = DeadlineScriptDialog()
    scriptDialog.AllowResizingDialog(False)
    scriptDialog.SetTitle("Convert to ProRes/H.264")

    scriptDialog.AddGrid()
    codecLabel = scriptDialog.AddControlToGrid(             "codecLabel",       "LabelControl",     "Codec:",                           0, 0 )
    codecList = scriptDialog.AddComboControlToGrid(         "codecList",        "ComboControl",     "", list(codecSettingsDict.keys()), 0, 1 )
    frameRateLabel = scriptDialog.AddControlToGrid(         "frameRateLabel",   "LabelControl",     "Frame Rate:",                      1, 0 )
    frameRate = scriptDialog.AddRangeControlToGrid(         "frameRate",        "RangeControl",     24,    0, 9999, 0, 1,               1, 1, colSpan=2 )
    burnInCheck = scriptDialog.AddSelectionControlToGrid(   "burnInCheck",      "CheckBoxControl",  False, "Burn-in Frame Number",      2, 1, colSpan=3 )
    okButton = scriptDialog.AddControlToGrid(               "okButton",         "ButtonControl",    "OK",                               3, 2, expand=False )
    okButton.ValueModified.connect(convert_to_prores)
    cancelButton = scriptDialog.AddControlToGrid(           "cancelButton",     "ButtonControl",    "Cancel",                           3, 3, expand=False )
    cancelButton.ValueModified.connect(close_dialog)
    scriptDialog.EndGrid()

    scriptDialog.ShowDialog(True)

def close_dialog():
    """ generic close dialog """
    global scriptDialog
    scriptDialog.CloseDialog()

def get_first_image_of_each_sequence(file_path, file_prefix):
    """Gets the first image of each unique image sequence from the given folder, matching a given filename prefix.

    Returns:
        A dictionary mapping the name of each image sequence to the first image file in the sequence.
    """

    image_sequences = defaultdict(list)
    for filename in os.listdir(file_path):
        if filename.startswith(file_prefix) and (filename.endswith(".png") or filename.endswith(".exr")):
            seq_name = filename.split(".")[0][:-4].rstrip("_")
            image_sequences[seq_name].append(filename)

    if len(image_sequences) is 0:
            print("Warning: No matching image sequences found. Please check Output File Prefix in source job matches image filenames in Output File Path")
            print(f"Warning: Expected filename prefix from source job: {file_prefix}")
            print("Warning: First filename found in output filepath: " + os.listdir(file_path)[0])
            exit

    first_images = {}
    for seq_name, image_sequence in image_sequences.items():
        first_image_file = image_sequence[0]
        first_images[seq_name] = first_image_file

    return first_images

def convert_to_prores(*args):
    """ create our conversion job(s) """

    global scriptDialog
    frameRate = scriptDialog.GetValue("frameRate")
    selectedCodec = scriptDialog.GetValue("codecList")
    codecSettings = codecSettingsDict[selectedCodec]
    codecSettings.meta = f'"{codecSettings.meta}"'

    selectedJobs = MonitorUtils.GetSelectedJobs()

    # loop through selected jobs
    for job in selectedJobs:

        # get job properties
        jobName = str(job.Name)
        firstFrame = job.JobFramesList[0]
        
        # set batch name parameter on selected job if one doesn't exist
        if not job.JobBatchName:
            job.JobBatchName = jobName.split(".")[0]
            RepositoryUtils.SaveJob(job)

        # get output filepath and prefix from plugin info
        sourceFilepath = job.GetJobPluginInfoKeyValue("FilePath")
        sourcePrefix = job.GetJobPluginInfoKeyValue("FilePrefix")
        
        # get first image of each image sequence in the source job filepath
        firstImages = get_first_image_of_each_sequence(sourceFilepath, sourcePrefix)
        imageSequences = []
        for seqName, firstImageFile in firstImages.items():
            print(f"First image of sequence {seqName}: {firstImageFile}")
            pathtoFirstImage = os.path.join(sourceFilepath, firstImageFile)
            if pathtoFirstImage not in imageSequences:
                imageSequences.append(pathtoFirstImage)

        # loop through unique image sequences
        for i in range(0, len(imageSequences)):
            fileName = imageSequences[i]
            ffJobName = ("_").join([jobName, selectedCodec])

            # set ffmpeg job name based on image sequence name if multiple sequences
            seqName = os.path.basename(fileName).split(".")[0][:-4].rstrip("_")
            if len(imageSequences) > 1:
                ffJobName = ("_").join([seqName, selectedCodec])

            # check if valid output file type
            if os.path.splitext(fileName)[-1] in ["mov", "avi", "mp4"]:
                print("Warning: {} -- output filetype is not compatible!  Can only accept image sequences.".format(fileName))
                continue          

            # get filename, remove _####, add codec name, and append file extension
            outputFile = fileName.split(".")[0]
            regex = r"(_\d{4}|\d{4})$"
            outputFile = re.sub(regex, "", outputFile)
            outputFile = ("_").join([outputFile, selectedCodec])
            outputFile += ".mp4" if selectedCodec == "h264" else ".mov"

            # create argument for burn-in (if enabled)
            ARIAL_FONT_FILE = "C\\\:/Windows/fonts/arial.ttf"
            burnInStr = ""
            if scriptDialog.GetValue("burnInCheck") == True:
                frameNumStr = "'%{frame_num}'"
                burnInStr = f"-vf \"drawtext=fontfile={ARIAL_FONT_FILE}:text={frameNumStr}:start_number={firstFrame}:x=(w-tw)/2:y=h-(2*lh):fontcolor=black:fontsize=20:box=1:boxcolor=white:boxborderw=5\""

            # create job file for submission
            jobInfoFile = Path.Combine( ClientUtils.GetDeadlineTempPath(), "batch_job_info.job")
            writer = File.CreateText( jobInfoFile )

            writer.WriteLine("Frames=0")
            writer.WriteLine("Name={}".format(ffJobName))
            writer.WriteLine("UserName={}".format(job.JobUserName))
            writer.WriteLine("BatchName={}".format(job.JobBatchName))
            writer.WriteLine("Plugin=FFmpeg")
            writer.WriteLine("LimitGroups=%s" % "ffmpeg")
            writer.WriteLine("Pool=")
            writer.WriteLine("Group=")
            writer.WriteLine("Priority=50")
            writer.WriteLine("JobDependency0={}".format(job.JobId))
            writer.WriteLine(f"ExtraInfo0={job.JobFrames}") # optional: pass frame range to ExtraInfo
            writer.Close()

            # create plugin properties for submission
            pluginInfoFile = Path.Combine( ClientUtils.GetDeadlineTempPath(), "batch_plugin_info.job" )
            writer = File.CreateText( pluginInfoFile )

            # Common arguments
            common_args = f"-y -xerror -vendor apl0 -profile:v {codecSettings.profileV} -metadata:s {codecSettings.meta} -pix_fmt {codecSettings.pixfmt} -r {frameRate} {burnInStr}"

            # Codec-specific arguments (entries not listed use default)
            codec_specific_args = defaultdict(lambda: f"-vcodec prores_ks -bits_per_mb 8000 -vtag {codecSettings.vtag}", {
                "h264": f"-vcodec libx264 -g 1 -tune zerolatency -crf 9 -bf 0",
            })

            writer.WriteLine(f"OutputArgs={common_args} {codec_specific_args[selectedCodec]}")
            writer.WriteLine(f"InputArgs0=-f image2 -r {frameRate} -start_number {firstFrame} -thread_queue_size 4096")
            writer.WriteLine("InputFile0={}".format(os.path.join(sourceFilepath, fileName)))
            writer.WriteLine("OutputFile={}".format(outputFile))
            writer.WriteLine("UseSameInputArgs=True")
            writer.WriteLine("ReplacePadding0=True")
            writer.Close()

            cmdArgs = [jobInfoFile, pluginInfoFile]
            exitCode = ClientUtils.ExecuteCommand( cmdArgs )
            if(exitCode == 0):
                print("Successfully submitted ffmpeg job: {}".format(ffJobName))
            else:
                print("Failed to submit ffmpeg job: {}".format(ffJobName))

    scriptDialog.CloseDialog()