DVCheck/dv_check.py

import argparse
import json
import subprocess
import sys
from pathlib import Path
from pymediainfo import MediaInfo
from tkinter import Tk, messagebox
from typing import Optional, Tuple

__version__ = "0.1.0"


def exit_print(message: Optional[str], exit_code: Optional[int]) -> None:
    print(message if message else "")
    sys.exit(exit_code if exit_code else 0)


def cli() -> Tuple[Path, Path, Path, bool, bool]:
    parser = argparse.ArgumentParser(prog="DV Check")
    parser.add_argument("-f", "--ffmpeg", type=str, help="Path to FFMPEG executable.")
    parser.add_argument(
        "-d", "--dovi-tool", type=str, help="Path to dovi_tool executable."
    )
    parser.add_argument(
        "-p",
        "--prompt",
        action="store_true",
        help="Prompt user with a GUI prompt if FEL.",
    )
    parser.add_argument(
        "-c",
        "--create-fel-mkv",
        action="store_true",
        help="Create FEL output mkv if FEL layer exists (input name_el.ext).",
    )
    parser.add_argument("input", type=str, help="Path to input file.")

    args = parser.parse_args()

    if not args.input:
        exit_print("Program requires an input file", 1)

    if not args.ffmpeg:
        exit_print("Program requires FFEMPG", 1)

    if not args.dovi_tool:
        exit_print("Program requires dovi_tool", 1)

    return (
        Path(args.input),
        Path(args.ffmpeg),
        Path(args.dovi_tool),
        args.prompt,
        args.create_fel_mkv,
    )


def detect_el_layer(file_input: Path) -> bool:
    if not file_input.exists() or not file_input.is_file():
        exit_print("Input file not found", 1)

    if file_input.suffix not in (".mp4", ".mkv"):
        exit_print("Only MP4/MKV input's are supported", 1)

    media_info = MediaInfo.parse(file_input)
    try:
        video_track = media_info.video_tracks[0]

        hdr_format = video_track.hdr_format
        if not hdr_format:
            exit_print("No HDR detected", 0)

        if "Dolby Vision" in hdr_format:
            hdr_format_settings = video_track.hdr_format_settings
            if not hdr_format_settings:
                exit_print("No HDR detected", 0)

            if "EL" in hdr_format_settings:
                return True
            else:
                exit_print("EL layer not detected", 0)
        else:
            exit_print("No dolby vision metadata detected", 0)

    except IndexError:
        exit_print("Input has no video track", 1)

    return False


def read_rpu(
    file_input: Path, ffmpeg_path: Path, dovi_tool_path: Path
) -> Optional[str]:
    # output_el_hevc = file_input.with_suffix(".hevc").with_name(f"{file_input.stem}_el")
    output_bin = file_input.with_suffix(".bin")

    # Change the working directory to the directory containing the input file
    file_input_parent = file_input.parent

    # Define the ffmpeg command
    ffmpeg_command = [
        str(ffmpeg_path),
        "-ss",
        "00:00:10",
        "-i",
        str(file_input),
        "-t",
        "00:00:01",
        "-map",
        "0:v:0",
        "-c:v:0",
        "copy",
        "-bsf:v",
        "hevc_mp4toannexb",
        "-f",
        "hevc",
        "-",
        "-hide_banner",
        "-v",
        "panic",
    ]

    # Define the dovi_tool command
    dovi_tool_command = [dovi_tool_path, "extract-rpu", "-", "-o", str(output_bin)]

    # Run the ffmpeg command and pipe its output to the dovi_tool command
    with subprocess.Popen(
        ffmpeg_command, cwd=file_input_parent, stdout=subprocess.PIPE
    ) as ffmpeg_proc:
        with subprocess.Popen(
            dovi_tool_command,
            cwd=file_input_parent,
            stdin=ffmpeg_proc.stdout,
            stdout=subprocess.DEVNULL,
        ) as dovi_proc:
            ffmpeg_proc.stdout.close()  # Close the ffmpeg stdout to allow it to receive a SIGPIPE if dovi_proc exits
            dovi_proc.communicate()  # Wait for the dovi_tool process to complete

    if output_bin.exists():
        rpu_data_cmd = [dovi_tool_path, "info", "-f", "0", "-i", str(output_bin)]
        rpu_data_job = subprocess.run(rpu_data_cmd, check=True, capture_output=True)
        rpu_stdout_data = json.loads(
            rpu_data_job.stdout.decode().split("Parsing RPU file...")[1]
        )
        el_type = rpu_stdout_data.get("el_type")

        # clean up files
        output_bin.unlink()

        return el_type


def prompt_user() -> None:
    root = Tk()
    root.withdraw()
    root.attributes('-topmost', True)
    messagebox.showinfo(
        "Alert",
        "FEL source detected, please switch to a FEL enabled template to process this encode.",
        parent=root
    )
    root.destroy()
    root.mainloop()


def create_fel(file_input: Path, ffmpeg_path: Path, dovi_tool_path: Path) -> None:
    ffmpeg_command = [
        str(ffmpeg_path),
        "-i",
        str(file_input),
        "-map",
        "0:v:0",
        "-c:v:0",
        "copy",
        "-bsf:v",
        "hevc_mp4toannexb",
        "-f",
        "hevc",
        "-",
        "-hide_banner",
        "-v",
        "panic",
        "-stats",
    ]

    el_output_path = file_input.with_name(f"{file_input.stem}_el.hevc")

    dovi_tool_command = [
        dovi_tool_path,
        "demux",
        "--el-only",
        "-",
        "-e",
        str(el_output_path),
    ]

    with subprocess.Popen(ffmpeg_command, stdout=subprocess.PIPE) as ffmpeg_proc:
        with subprocess.Popen(
            dovi_tool_command,
            stdin=ffmpeg_proc.stdout,
        ) as dovi_proc:
            ffmpeg_proc.stdout.close()  # Close the ffmpeg stdout to allow it to receive a SIGPIPE if dovi_proc exits
            dovi_proc.communicate()  # Wait for the dovi_tool process to complete

    ffmpeg_mux_command = [
        str(ffmpeg_path),
        "-y",
        "-i",
        str(el_output_path),
        "-c",
        "copy",
        "-f",
        "hevc",
        str(el_output_path.with_suffix(".mkv")),
    ]
    subprocess.run(ffmpeg_mux_command)


if __name__ == "__main__":
    file_input, ffmpeg, dovi_tool, prompt, create_fel_mkv = cli()
    detect_el = detect_el_layer(file_input)
    if detect_el:
        rpu = read_rpu(file_input, ffmpeg, dovi_tool)
        if rpu:
            rpu = rpu.strip()
            print(rpu)
            if prompt and rpu == "FEL":
                prompt_user()
            if create_fel_mkv:
                create_fel(file_input, ffmpeg, dovi_tool)
    else:
        print("No detected EL layer")

    exit_print("", 0)