Av1an/av1an/concat.py

import os
import platform
import shlex
import subprocess
import sys
from pathlib import Path
from subprocess import PIPE, STDOUT

from av1an.logger import log

if platform.system() == "Linux":
    import resource


def vvc_concat(temp: Path, output: Path):
    """
    Concatenates vvc files

    :param temp: the temp directory
    :param output: the output video
    :return: None
    """
    encode_files = sorted((temp / "encode").iterdir())
    bitstreams = [x.as_posix() for x in encode_files]
    bitstreams = " ".join(bitstreams)
    cmd = f"vvc_concat  {bitstreams} {output.as_posix()}"

    subprocess.run(cmd, shell=True, check=True)


def concatenate_ffmpeg(temp: Path, output: Path, encoder: str):
    """
    Uses ffmpeg to concatenate encoded segments into the final file

    :param temp: the temp directory
    :param output: the final output file
    :param encoder: the encoder
    :return: None
    """

    log("Concatenating")

    with open(temp / "concat", "w") as f:

        encode_files = sorted((temp / "encode").iterdir())
        f.writelines(
            f'file {shlex.quote("file:"+str(file.absolute()))}\n'
            for file in encode_files
        )

    # Add the audio/subtitles/else file if one was extracted from the input
    audio_file = temp / "audio.mkv"
    if audio_file.exists() and audio_file.stat().st_size > 1024:
        audio = ("-i", audio_file.as_posix(), "-c", "copy", "-map", "1")
    else:
        audio = ()

    if encoder == "x265":

        cmd = [
            "ffmpeg",
            "-y",
            "-fflags",
            "+genpts",
            "-hide_banner",
            "-loglevel",
            "error",
            "-f",
            "concat",
            "-safe",
            "0",
            "-i",
            (temp / "concat").as_posix(),
            *audio,
            "-c",
            "copy",
            "-movflags",
            "frag_keyframe+empty_moov",
            "-map",
            "0",
            "-f",
            "mp4",
            output.as_posix(),
        ]
        concat = subprocess.run(cmd, stdout=PIPE, stderr=STDOUT, check=True).stdout

    else:
        cmd = [
            "ffmpeg",
            "-y",
            "-hide_banner",
            "-loglevel",
            "error",
            "-f",
            "concat",
            "-safe",
            "0",
            "-i",
            (temp / "concat").as_posix(),
            *audio,
            "-c",
            "copy",
            "-map",
            "0",
            output.as_posix(),
        ]

        concat = subprocess.run(cmd, stdout=PIPE, stderr=STDOUT, check=True).stdout

    if len(concat) > 0:
        log(concat.decode())
        print(concat.decode())
        tb = sys.exc_info()[2]
        raise RuntimeError.with_traceback(tb)


def concatenate_mkvmerge(temp: Path, output):
    """
    Uses mkvmerge to concatenate encoded segments into the final file

    :param temp: the temp directory
    :param output: the final output file
    :return: None
    """

    log("Concatenating")

    output = shlex.quote(output.as_posix())

    encode_files = sorted(
        (temp / "encode").iterdir(),
        key=lambda x: int(x.stem) if x.stem.isdigit() else x.stem,
    )
    encode_files = [shlex.quote(f.as_posix()) for f in encode_files]

    if platform.system() == "Linux":
        file_limit, _ = resource.getrlimit(resource.RLIMIT_NOFILE)
        cmd_limit = os.sysconf(os.sysconf_names["SC_ARG_MAX"])
    else:
        file_limit = -1
        cmd_limit = 32767

    audio_file = temp / "audio.mkv"
    audio = audio_file.as_posix() if audio_file.exists() else ""

    if len(encode_files) > 1:
        encode_files = [
            _concatenate_mkvmerge(encode_files, output, file_limit, cmd_limit)
        ]

    cmd = ["mkvmerge", "-o", output, encode_files[0]]

    if audio:
        cmd.append(audio)

    concat = subprocess.Popen(cmd, stdout=PIPE, universal_newlines=True)
    message, _ = concat.communicate()
    concat.wait()

    if concat.returncode != 0:
        log(message)
        print(message)
        tb = sys.exc_info()[2]
        raise RuntimeError.with_traceback(tb)

    # remove temporary files used by recursive concat
    if os.path.exists("{}.tmp0.mkv".format(output)):
        os.remove("{}.tmp0.mkv".format(output))

    if os.path.exists("{}.tmp1.mkv".format(output)):
        os.remove("{}.tmp1.mkv".format(output))


def _concatenate_mkvmerge(files, output, file_limit, cmd_limit, flip=False):
    tmp_out = "{}.tmp{}.mkv".format(output, int(flip))
    cmd = ["mkvmerge", "-o", tmp_out, files[0]]

    remaining = []
    for i, file in enumerate(files[1:]):
        new_cmd = cmd + ["+{}".format(file)]
        if sum(len(s) for s in new_cmd) < cmd_limit and (
            file_limit == -1 or i < max(1, file_limit - 10)
        ):
            cmd = new_cmd
        else:
            remaining = files[i + 1 :]
            break

    concat = subprocess.Popen(cmd, stdout=PIPE, universal_newlines=True)
    message, _ = concat.communicate()
    concat.wait()

    if concat.returncode != 0:
        log(message)
        print(message)
        tb = sys.exc_info()[2]
        raise RuntimeError.with_traceback(tb)

    if len(remaining) > 0:
        return _concatenate_mkvmerge(
            [tmp_out] + remaining, output, file_limit, cmd_limit, not flip
        )
    return tmp_out