MediaDash/transcode.py

import io
import json
import os
import shlex
import subprocess as SP
import sys
import time
import uuid

from tqdm import tqdm

from utils import handle_config

profiles = handle_config().get("transcode_profiles", {})

profiles[None] = {
    "command": "-vcodec copy -acodec copy -scodec copy -f null",
    "doc": "null output for counting frames",
}


def ffprobe(file):
    cmd = [
        "ffprobe",
        "-v",
        "error",
        "-print_format",
        "json",
        "-show_format",
        "-show_streams",
        file,
    ]
    try:
        out = SP.check_output(cmd)
    except KeyboardInterrupt:
        raise
    except BaseException:
        return file, None
    return file, json.loads(out)


def make_ffmpeg_command_line(infile, outfile, profile=None, **kwargs):
    default_opts = ["-v", "error", "-y", "-nostdin"]
    ffmpeg = (
        "C:\\Users\\Earthnuker\\scoop\\apps\\ffmpeg-nightly\\current\\bin\\ffmpeg.exe"
    )
    cmdline = profile["command"]
    opts = profile.get("defaults", {}).copy()
    opts.update(kwargs)

    if isinstance(cmdline, str):
        cmdline = shlex.split(cmdline)
    cmdline = list(cmdline or [])
    cmdline += ["-progress", "-", "-nostats"]
    ret = [ffmpeg, *default_opts, "-i", infile, *cmdline, outfile]
    ret = [v.format(**opts) for v in ret]
    return ret


def count_frames(file, **kwargs):
    total_frames = None
    for state in run_transcode(file, os.devnull, None):
        if state.get("progress") == "end":
            total_frames = int(state.get("frame", -1))
    if total_frames is None:
        return total_frames
    if total_frames <= 0:
        total_frames = None
    return total_frames


def run_transcode(file, outfile, profile, job_id=None, **kwargs):
    job_id = job_id or str(uuid.uuid4())
    stderr_fh = None
    if outfile != os.devnull:
        stderr_fh = open("{}.log".format(job_id), "w")
    proc = SP.Popen(
        make_ffmpeg_command_line(file, outfile, profiles[profile], **kwargs),
        stdout=SP.PIPE,
        stderr=stderr_fh,
        encoding="utf8",
    )
    state = {}
    poll = None
    while poll is None:
        poll = proc.poll()
        state["ret"] = poll
        if outfile != os.devnull:
            with open("{}.log".format(job_id), "r") as tl:
                state["stderr"] = tl.read()
        line = proc.stdout.readline().strip()
        if not line:
            continue
        try:
            key, val = line.split("=", 1)
        except ValueError:
            print(line)
            continue
        key = key.strip()
        val = val.strip()
        state[key] = val
        if key == "progress":
            yield state
    if stderr_fh:
        stderr_fh.close()
        os.unlink(stderr_fh.name)
    yield state


def transcode(file, outfile, profile, job_id=None, **kwargs):
    from pprint import pprint

    info = ffprobe(file)
    frames = count_frames(file)
    progbar = tqdm(
        desc="Processing {}".format(outfile),
        total=frames,
        unit=" frames",
        disable=False,
        leave=False,
    )
    for state in run_transcode(file, outfile, profile, job_id, **kwargs):
        if "frame" in state:
            progbar.n = int(state["frame"])
            progbar.update(0)
        state["total_frames"] = frames
        state["file"] = file
        state["outfile"] = outfile
        # progbar.write(state["stderr"])
        yield state
    progbar.close()


def preview_command(file, outfile, profile, **kwargs):
    return make_ffmpeg_command_line(file, outfile, profiles[profile], **kwargs)


if __name__ == "__main__":
    file = sys.argv[1]
    for profile in ["H.265 transcode", "H.264 transcode"]:
        for preset in ["ultrafast", "fast", "medium", "slow", "veryslow"]:
            for crf in list(range(10, 54, 4))[::-1]:
                outfile = os.path.join(
                    "E:\\",
                    "transcode",
                    profile,
                    "{}_{}.mkv".format(
                        crf,
                        preset))
                os.makedirs(os.path.dirname(outfile), exist_ok=True)
                if os.path.isfile(outfile):
                    print("Skipping", outfile)
                    continue
                for _ in transcode(
                    file, outfile, profile, "transcode", preset=preset, crf=crf
                ):
                    pass