visionary

#!/usr/bin/env bash

usage() {
cat - <<EOF
Usage: ${0##*/} -i <INPUT FILE>

-h|--help:          show this help message
-i|--input:         file name for input
-o|--output:        base file name for one million artifacts, do not include file extension
-k|--keep:          keep artifacts
-v|--verify:        verify stream with Dolby Vision toolkit, default disabled
-c|--crf:           set an alternate crf (defaults: x265: 16, svtav1: 22)
-p|--preset:        set an alternate encoder preset (defaults: x265: medium, svtav1: 7)
-t|--tune:          set an alternate encoder tune (defaults: x265: none, svtav1: 3)
-a|--audio:         transcode source audio to specified codec (opus, aac, flac)
--av1:              HIGHLY EXPERIMENTAL SUPPORT FOR AV1
--dual:             Run both an hevc AND an av1 encode in one pass (more efficient)
--debug:            print debugging information

EOF
}

HEVC=1

while [ $# -gt 0 ]; do
    OPT="${1}"
    case "${OPT}" in
    -i|--input)
        shift
        INPUT="${1}"
        shift;;
    -o|--output)
        shift
        OUTPUT="${1}"
        shift;;
    -k|--keep)
        KEEP=1
        shift;;
    -v|--verify)
        VERIFY=1
        shift;;
    -c|--crf)
        shift
        CRF="${1}"
        shift;;
    -p|--preset)
        shift
        PRESET="${1}"
        shift;;
    -t|--tune)
        shift
        TUNE="${1}"
        shift;;
    -a|--audio)
        shift
        AUDIO="${1}"
        shift;;
    --av1)
        AV1=1
        HEVC=0
        shift;;
    --dual)
        AV1=1
        shift;;
    --debug)
        DEBUG=1
        shift;;
    -h|--help)
        usage
        exit;;
    *) 
        echo "${0}: invalid option ${1}"
        usage
        exit 1;;
    esac
done

# Check for tool dependencies
SELF_FULLPATH=`realpath ${0}`
SELF_DIR=`dirname ${SELF_FULLPATH}`
# Override default tool list
TOOLS=(mkvmerge mkvextract mediainfo grep awk dovi_tool hdr10plus_tool jq mkfifo base64 ffprobe ffmpeg SvtAv1EncApp sed bc column)
[ -n "${VERIFY}" ] && TOOLS+=(dvesverifier)
# Depends on hevcmeup2 for transcoding
SCRIPTS=(hevcmeup2 striptags stripcovers addstats)
[ -n "${AUDIO}" ] && SCRIPTS+=(audiovise)

if [ ! -r ${SELF_DIR}/toolchecker ]; then
    echo "toolchecker not found."
    exit 1
else
    . ${SELF_DIR}/toolchecker
fi

## validate something was even provided for input
if [ -z "${INPUT}" ]; then
    echo "No input file provided."
    usage
    exit 1
fi

if [ -z "${OUTPUT}" ]; then
    echo "No output file base name provided."
    echo "Building output file base name from input file name..."
    OUTPUT=$(echo "${INPUT}" | sed -e 's/.\(mkv\|mp4\|ts\)/_visionary/g')
    echo "Output file name base: ${OUTPUT}"
fi

if [ -n "${AV1}" ] || [ -n "${DUAL}" ]; then
        CHECK_SVTAV1_VERSION=$(${SVTAV1ENCAPP} --version | grep PSY &>/dev/null && echo "PSY")
        if [ "${CHECK_SVTAV1_VERSION}" != "PSY" ]; then
            echo "This script requires the SVT-AV1-PSY fork for Dolby Vision AV1 encoding."
            echo "See https://github.com/gianni-rosato/svt-av1-psy for details."
            exit 1
        fi
fi

echo ">> Starting Dolby Vision/HDR10+ transcode"; echo

## this makes 'for crf in 16 18 20 22; do visionary -i A.Movie.mkv -o movie --crf $crf; done'
## work so much easier that I don't care if it makes sense to anyone else
if [ -n "${CRF}" ]; then
    OUTPUT="${OUTPUT}_${CRF}"
    echo ">> Appending crf value to output filenames: ${OUTPUT}"
fi

## build assorted artifact filenames and make sure they don't exist
## not all of these are used in all cases
## but I don't want to think about it very hard
FIFO="${OUTPUT}.fifo"
[ -p "${FIFO}" ] && fail "Named pipe already exists: ${FIFO}"
BL="${OUTPUT}.BL.hevc"
[ -f "${BL}" ] && fail "Base-layer file already exists: ${BL}"
BL10P="${OUTPUT}.BL10P.hevc"
[ -f "${BL10P}" ] && fail "HDR10+ base layer already exists: ${BL10P}"
EL="${OUTPUT}.EL.hevc"
[ -f "${EL}" ] && fail "Enhancement-layer file already exists: ${EL}"
RPU="${OUTPUT}.RPU.bin"
[ -f "${RPU}" ] && fail "RPU file already exists: ${RPU}"
RPU_8="${OUTPUT}.RPU_8.bin"
[ -f "${RPU_8}" ] && fail "Profile 8-compatible RPU file already exists: ${RPU_8}"
HDR10P="${OUTPUT}.hdr10p.json"
[ -f "${HDR10P}" ] && fail "HDR10+ JSON file already exists: ${HDR10P}"
BLX265="${OUTPUT}.BL.x265.hevc"
[ -f "${BLX265}" ] && fail "x265 transcoding destination file already exists: ${BLX265}"
WITH_RPU="${OUTPUT}.BL+RPU.x265.hevc"
[ -f "${WITH_RPU}" ] && fail "Transcoded file with RPU injection already exists: ${WITH_RPU}"
AV1_IVF="${OUTPUT}.av1.ivf"
[ -f "${AV1_IFV}" ] && fail "AV1 output file already exists: ${AV1_IVF}"
REMUXED="${OUTPUT}.mkv"
[ -f "${REMUXED}" ] && fail "Final remux target file already exists: ${REMUXED}"

#HDR=$($MEDIAINFO "${INPUT}" | $GREP ^HDR\ format | $AWK -F: '{print $2}')
FPS=$(${MEDIAINFO} --output=JSON "${INPUT}" | ${JQ} -r '.media.track[] | select(."@type" == "Video" ).FrameRate')
TRACK=$(${MEDIAINFO} --output=JSON "${INPUT}" | ${JQ} -r '.media.track[] | select (."@type"=="Video") | @base64')
HDRFORMAT=$(echo "${TRACK}" | ${BASE64} --decode | ${JQ} -r '.HDR_Format')
FORMATCOMPAT=$(echo "${TRACK}" | ${BASE64} --decode | ${JQ} -r '.HDR_Format_Compatibility')

ISDOVI=0
ISHDR10P=0

if [[ "${HDRFORMAT}" =~ "Dolby" ]]; then
    ISDOVI=1
    echo ">> Dolby Vision metadata present"
fi

if [[ "${FORMATCOMPAT}" =~ "HDR10+" ]]; then
    ISHDR10P=1
    echo ">> HDR10+ metadata present"
fi

let HDRSUM=${ISDOVI}+${ISHDR10P}

if [ $DEBUG ]; then
    echo " VARIABLE STATE:"
    echo "-----------------"
    for var in INPUT OUTPUT HDRFORMAT ISDOVI FORMATCOMPAT ISHDR10P HDRSUM AV1 DUAL KEEP VERIFY CRF PRESET TUNE AUDIO; do
        [ -n "${!var}" ] && echo "$var,${!var}" || echo "$var,<unset>"
    done | ${COLUMN} --separator , --table --table-name "VARIABLE STATE" --output-separator " | "
    echo
fi

if [ ${HDRSUM} -eq 0 ]; then
    fail ">> No dynamic metadata present. What are we even doing here, folks?"
fi

echo ">> Making named pipe"; echo
set -x
${MKFIFO} "${FIFO}"
{ set +x; } 2>/dev/null
echo

## Dolby Vision stream prep
if [ ${ISDOVI} -eq 1 ]; then
    DOVIPROFILE=$(echo "${TRACK}" | ${BASE64} --decode | ${JQ} -r '.HDR_Format_Profile' | ${AWK} '{print $1}')
    DOVILEVEL=$(echo "${TRACK}" | ${BASE64} --decode | ${JQ} -r '.HDR_Format_Level' | ${AWK} '{print $1}')

    echo ">> Dolby Vision profile: ${DOVIPROFILE}"
    echo ">> Dolby Vision level: ${DOVILEVEL}"
    echo ">> Base layer compatibility: ${FORMATCOMPAT}"; echo

    case "${DOVIPROFILE}" in
        "dvhe.05")
            [ -n "${DEBUG}" ] && echo "MEL-type profile 5"
            echo ">> Cleaning up named pipe"; echo
            set -x
            rm "${FIFO}"
            { set +x; } 2>/dev/null
            echo
            fail "<> I'm sorry Dave, I cannot currently encode dolby vision profile 5 content."
            #echo ">> NOTE: Profile 5 uses non-standard IPT-PQ-C2 colorspace in BL. Colors will not be properly rendered on playback devices that do not support Dolby Vision."

            ### create stripped base layer
            #echo ">> Creating stripped base layer"; echo
            #set -x
            #${MKVEXTRACT} "${INPUT}" tracks 0:"${FIFO}" &
            #${DOVI_TOOL} remove --output "${BL}" - <"${FIFO}" || fail "Base layer stripping failed."
            #{ set +x; } 2>/dev/null
            #echo

            ### extract the RPU
            #echo ">> Extracting RPU"; echo
            #set -x
            #${MKVEXTRACT} "${INPUT}" tracks 0:"${FIFO}" &
            #${DOVI_TOOL} --mode 3 extract-rpu --rpu-out "${RPU_8}" - <"${FIFO}" || fail "RPU extraction failed."
            #{ set +x; } 2>/dev/null
            #echo
            ;;
        "dvhe.07")
            ## since I can't figure out how to preserve a FEL, we're gonna demux the BL and EL and extract the RPU to 8-compatible format
            [ -n "${DEBUG}" ] && echo "FEL-type profile 7"
            [ -n "${DEBUG}" ] && echo "Cannot preserve EL through transcoding, will convert RPU to MEL-compatible profile 8 format"

            ## demux the layers
            echo ">> Demuxing layers"; echo
            set -x
            ${MKVEXTRACT} "${INPUT}" tracks 0:"${FIFO}" &
            ${DOVI_TOOL} demux --bl-out "${BL}" --el-out "${EL}" - <"${FIFO}"  || fail "Demuxing failed."
            { set +x; } 2>/dev/null
            echo

            ## extract the original RPU from the EL
            ## until I figure out how to preserve a FEL stack, no need for this step, but I'm keeping it in here
            #[ -n "${DEBUG}" ] && echo "Extracting RPU"
            #${DOVI_TOOL} extract-rpu --input ${EL} --rpu-out ${RPU}
            
            ## extract and convert the profile 7 RPU to a profile 8 RPU
            echo ">> Extracting RPU in profile 8 format"; echo
            set -x
            ${DOVI_TOOL} --mode 2 extract-rpu --input ${EL} --rpu-out ${RPU_8} || fail "RPU extraction failed."
            { set +x; } 2>/dev/null
            echo
            ;;
        "dvhe.08")
            ## great news, we're already in profile 8 territory
            [ -n "${DEBUG}" ] && echo "MEL-type profile 8"

            ## create stripped base layer
            echo ">> Removing Dolby Vision RPUs from base layer"; echo
            set -x
            ${MKVEXTRACT} "${INPUT}" tracks 0:"${FIFO}" &
            ${DOVI_TOOL} remove --output "${BL}" - <"${FIFO}" || fail "Base layer stripping failed."
            { set +x; } 2>/dev/null
            echo

            ## extract the RPU
            echo ">> Extracting RPU"; echo
            set -x
            ${MKVEXTRACT} "${INPUT}" tracks 0:"${FIFO}" &
            ${DOVI_TOOL} extract-rpu --rpu-out "${RPU_8}" - <"${FIFO}" || fail "RPU extraction failed."
            { set +x; } 2>/dev/null
            echo
            ;;
        *) 
            ## not gonna need that fifo after all I guess
            echo ">> Cleaning up named pipe"; echo
            set -x
            rm "${FIFO}"
            { set +x; } 2>/dev/null
            echo
            fail "Unrecognized Dolby Vision identifier: ${DOVIPROFILE}";;
    esac
fi

## HDR10+ stream prep
if [ ${ISHDR10P} -eq 1 ]; then
    if [ ! -f "${BL}" ]; then
        echo ">> Creating HDR10+-free base layer"; echo
        set -x
        ${MKVEXTRACT} "${INPUT}" tracks 0:"${FIFO}" &
        ${HDR10PLUS_TOOL} remove --output "${BL}" - <"${FIFO}" || fail "Base layer stripping failed."
        { set +x; } 2>/dev/null
        echo
        echo ">> Extracting HDR10+ metadata"; echo
        set -x
        ${MKVEXTRACT} "${INPUT}" tracks 0:"${FIFO}" &
        ${HDR10PLUS_TOOL} extract -o "${HDR10P}" - <"${FIFO}" || fail "HDR10+ extraction failed."
        { set +x; } 2>/dev/null
        echo
    else
        echo ">> Extracting HDR10+ metadata from existing base layer ${BL}"; echo
        set -x
        ${HDR10PLUS_TOOL} extract -o "${HDR10P}" -i "${BL}" || fail "HDR10+ extraction failed."
        { set +x; } 2>/dev/null
        echo
        echo ">> Removing HDR10+ metadata from existing base layer"; echo
        set -x
        ${HDR10PLUS_TOOL} remove -o "${BL%.hevc}.nohdr10p.hevc" -i "${BL}" || fail "HDR10+ extraction failed."
        rm "${BL}"
        mv "${BL%.hevc}.nohdr10p.hevc" "${BL}"
        { set +x; } 2>/dev/null
        echo
    fi
fi

## we're done with the fifo
echo ">> Cleaning up named pipe"; echo
set -x
rm "${FIFO}"
{ set +x; } 2>/dev/null
echo

## let's cook this sucker
echo ">> Starting encode..."; echo
if [ ${AV1} -eq 1 ]; then
    ## REFERENCES
    # https://codecalamity.com/encoding-uhd-4k-hdr10-videos-with-ffmpeg/
    ## THIS IS A DANK ASS TRICK THAT MAKES VARIABLES OUT OF ALL THE VIDEO METADATA
    eval "$(${FFPROBE} -analyzeduration 100 -probesize 100M -hide_banner -loglevel warning -select_streams v -print_format json -show_frames -read_intervals "%+#1" -show_entries "frame=hdr_format,color_range,color_space,color_primaries,color_transfer,side_data_list,pix_fmt" -i "${BL}" | ${GREP} -v -e frames -e '{' -e '\[' -e '}' -e '\]' -e 'side' | ${SED} -e 's/^\ *//g' -e 's/^"//g' -e 's/\":/:/g' -e 's/:\ /=/g' -e 's/,//g' -e 's/^/video_/g' -e 's/^.*=/\U&/g' -e 's/^\(.*\)\ \([A-Z].*\)=\(.*\)$/\1_\2=\3/g')"

    VIDEO_COLOR_SPACE_FIXED="${VIDEO_COLOR_SPACE/nc/-ncl}"

    ## Processing up the variables from ffprobe fractions to floats
    VIDEO_GREEN_X_BC=$(echo "scale=4;${VIDEO_GREEN_X}" | ${BC})
    VIDEO_GREEN_Y_BC=$(echo "scale=4;${VIDEO_GREEN_Y}" | ${BC})
    VIDEO_BLUE_X_BC=$(echo "scale=4;${VIDEO_BLUE_X}" | ${BC})
    VIDEO_BLUE_Y_BC=$(echo "scale=4;${VIDEO_BLUE_Y}" | ${BC})
    VIDEO_RED_X_BC=$(echo "scale=4;${VIDEO_RED_X}" | ${BC})
    VIDEO_RED_Y_BC=$(echo "scale=4;${VIDEO_RED_Y}" | ${BC})
    VIDEO_WHITE_POINT_X_BC=$(echo "scale=4;${VIDEO_WHITE_POINT_X}" | ${BC})
    VIDEO_WHITE_POINT_Y_BC=$(echo "scale=4;${VIDEO_WHITE_POINT_Y}" | ${BC})
    VIDEO_MAX_LUMINANCE_BC=$(echo "scale=4;${VIDEO_MAX_LUMINANCE}" | ${BC})
    VIDEO_MIN_LUMINANCE_BC=$(echo "scale=4;${VIDEO_MIN_LUMINANCE}" | ${BC})

    ## Common parameters
    SVTAV1_PARAMS="${SVTAV1_PARAMS} --enable-hdr 1"
    SVTAV1_PARAMS="${SVTAV1_PARAMS} --color-primaries ${VIDEO_COLOR_PRIMARIES}"
    SVTAV1_PARAMS="${SVTAV1_PARAMS} --transfer-characteristics ${VIDEO_COLOR_TRANSFER}"
    SVTAV1_PARAMS="${SVTAV1_PARAMS} --matrix-coefficients ${VIDEO_COLOR_SPACE_FIXED}"

    ## parameters specific to color transfers
    case "${VIDEO_COLOR_TRANSFER}" in
        smpte2084)
            echo ">> HDR10/SMPTE2084 color transfer"
            GREEN="${VIDEO_GREEN_X_BC},${VIDEO_GREEN_Y_BC}"
            BLUE="${VIDEO_BLUE_X_BC},${VIDEO_BLUE_Y_BC}"
            RED="${VIDEO_RED_X_BC},${VIDEO_RED_Y_BC}"
            WHITE_POINT="${VIDEO_WHITE_POINT_X_BC},${VIDEO_WHITE_POINT_Y_BC}"
            LUMINANCE="${VIDEO_MAX_LUMINANCE_BC},${VIDEO_MIN_LUMINANCE_BC}"
            CLL="${VIDEO_MAX_CONTENT},${VIDEO_MAX_AVERAGE}"
            SVTAV1_PARAMS="${SVTAV1_PARAMS} --mastering-display G(${GREEN})B(${BLUE})R(${RED})WP(${WHITE_POINT})L(${LUMINANCE}) --content-light ${CLL}";;
        arib-std-b67)
            echo ">> HLG/ARIB-STD-B67 color transfer";;
        *)
            echo "Unrecognized color transfer: ${VIDEO_COLOR_TRANSFER}"
            echo "Bailing out."
            exit 1;;
    esac

    SVTAV1_PARAMS="--enable-dlf 2 --enable-variance-boost 1 --enable-qm 1 --qm-min 0 --irefresh-type 2 --progress 3 -i stdin --preset ${PRESET:-7} --input-depth 10 --crf ${CRF:-22} --keyint 240 --tune ${TUNE:-3} ${SVTAV1_PARAMS}"

    [ ${ISDOVI} -eq 1 ] && SVTAV1_PARAMS="${SVTAV1_PARAMS} --dolby-vision-rpu ${RPU_8}"
    [ ${ISHDR10P} -eq 1 ] && SVTAV1_PARAMS="${SVTAV1_PARAMS} --hdr10plus-json ${HDR10P}"

    set -x
    nice -n 5 ${FFMPEG} -hide_banner -loglevel error -i "${BL}" -map 0:v:0 -pix_fmt yuv420p10le -f yuv4mpegpipe -strict -1 - | nice -n 5 ${SVTAV1ENCAPP} ${SVTAV1_PARAMS} --output ${AV1_IVF}
    { set +x; } 2>/dev/null
fi

## HEVC
if [ ${HEVC} -eq 1 ]; then
    OPTIONS="--input ${BL} --output ${BLX265} --fps ${FPS} ${CRF:+--crf $CRF} ${PRESET:+--preset $PRESET} ${TUNE:+--tune $TUNE}"
    [ ${ISDOVI} -eq 1 ] && OPTIONS="$OPTIONS --dovi-bl"
    set -x
    ${HEVCMEUP2} ${OPTIONS} || fail "Encode failed."
    { set +x; } 2>/dev/null
    COOKED="${BLX265}"
    ## time to put the sidecars back in the garage
    ## do HDR10+ injection first, if need be
    if [ ${ISHDR10P} -eq 1 ]; then
        echo ">> Injecting HDR10+ metadata into transcoded base layer..."; echo
        set -x
        ${HDR10PLUS_TOOL} inject --input ${COOKED} --json ${HDR10P} --output ${BL10P} || fail "HDR10+ injection failed."
        { set +x; } 2>/dev/null
        echo
        COOKED="${BL10P}"
        FINALVIDEO_HEVC="${BL10P}"
    fi
    ## dovi injection, if need be
    if [ ${ISDOVI} -eq 1 ]; then
        echo ">> Injecting RPU into transcoded BL..."; echo
        set -x
        ${DOVI_TOOL} inject-rpu --input ${COOKED} --rpu-in ${RPU_8} --output ${WITH_RPU} || fail "RPU injection failed."
        { set +x; } 2>/dev/null
        echo
        ## verify, if requested, you really shouldn't request it tho
        if [ -n "${VERIFY}" ]; then
            echo ">> Running Dolby Vision stream verifier"; echo
            set -x
            ${DVESVERIFIER} --input ${WITH_RPU} --rpu-compare ${RPU_8} --dv-profile ${DOVIPROFILE#dvhe.0}.${DOVILEVEL#0} --dv-level ${DOVILEVEL#0} --verbose 1 || fail "Dolby Vision verification failed."
            { set +x; } 2>/dev/null
            echo
        fi
        COOKED="${WITH_RPU}"
        FINALVIDEO_HEVC="${WITH_RPU}"
    fi
fi

if [[ "${AUDIO}" =~ (opus|aac|flac) ]]; then
    echo ">> Transcoding source audio to ${AUDIO}..."
    OTHERTRACKS="${OUTPUT}_${AUDIO}.mkv"
    set -x
    ${AUDIOVISE} -i ${INPUT} -o ${OTHERTRACKS} --${AUDIO} --novideo
    { set +x; } 2>/dev/null
    echo
elif [ -n "${AUDIO}" ] ; then
    echo ">> Audio codec \"${AUDIO}\" not recognized, ignoring"
    OTHERTRACKS="${INPUT}"
else
    OTHERTRACKS="${INPUT}"
fi

echo ">> Remuxing transcoded video with other source tracks..."; echo
set -x
[ $AV1 -eq 1 ] && ${MKVMERGE} --output "${REMUXED/mkv/av1.mkv}" "${AV1_IVF}" --no-video "${OTHERTRACKS}" || fail "AV1 merge failed."
[ $HEVC -eq 1 ] && ${MKVMERGE} --output "${REMUXED/mkv/hevc.mkv}" "${FINALVIDEO_HEVC}" --no-video "${OTHERTRACKS}" || fail "HEVC merge failed."
{ set +x; } 2>/dev/null

set -x
${STRIPTAGS} "${OUTPUT}.mkv"
${STRIPCOVERS} "${OUTPUT}.mkv"
${ADDSTATS} "${OUTPUT}.mkv"
{ set +x; } 2>/dev/null
echo

if [ -n "${KEEP}" ]; then
    echo ">> Preserving intermediate files"
else
    echo ">> Removing intermediate files..."; echo
    CLEANUP=(${BL} ${EL} ${RPU} ${RPU_8} ${HDR10P} ${BLX265} ${BL10P} ${WITH_RPU} ${FINALVIDEO_HEVC} ${AV1_IVF})
    [[ "${OTHERTRACKS}" != "${INPUT}" ]] && CLEANUP+=(${OTHERTRACKS})
    for FILE in ${CLEANUP[@]}; do
        if [ -f "${FILE}" ]; then
            set -x
            rm ${FILE}
            { set +x; } 2>/dev/null
        fi
    done
fi
echo

echo ">> Complete"