FrameTour-RenderWorker/util/ffmpeg.py

import json
import logging
import os
import subprocess
from datetime import datetime
from typing import Optional, IO

from opentelemetry.trace import Status, StatusCode

from entity.ffmpeg import (
    FfmpegTask,
    ENCODER_ARGS,
    VIDEO_ARGS,
    AUDIO_ARGS,
    MUTE_AUDIO_INPUT,
    get_mp4toannexb_filter,
)
from telemetry import get_tracer

logger = logging.getLogger(__name__)


def re_encode_and_annexb(file):
    with get_tracer("ffmpeg").start_as_current_span("re_encode_and_annexb") as span:
        span.set_attribute("file.path", file)
        if not os.path.exists(file):
            span.set_status(Status(StatusCode.ERROR))
            return file
        logger.info("ReEncodeAndAnnexb: %s", file)
        has_audio = not not probe_video_audio(file)
        # 优先使用RE_ENCODE_VIDEO_ARGS环境变量，其次使用默认的VIDEO_ARGS
        if os.getenv("RE_ENCODE_VIDEO_ARGS", False):
            _video_args = tuple(os.getenv("RE_ENCODE_VIDEO_ARGS", "").split(" "))
        else:
            _video_args = VIDEO_ARGS
        # 优先使用RE_ENCODE_ENCODER_ARGS环境变量，其次使用默认的ENCODER_ARGS
        if os.getenv("RE_ENCODE_ENCODER_ARGS", False):
            _encoder_args = tuple(os.getenv("RE_ENCODE_ENCODER_ARGS", "").split(" "))
        else:
            _encoder_args = ENCODER_ARGS
        ffmpeg_process = subprocess.run(
            [
                "ffmpeg",
                "-y",
                "-hide_banner",
                "-i",
                file,
                *(set() if has_audio else MUTE_AUDIO_INPUT),
                "-fps_mode",
                "cfr",
                "-map",
                "0:v",
                "-map",
                "0:a" if has_audio else "1:a",
                *_video_args,
                "-bsf:v",
                get_mp4toannexb_filter(),
                *AUDIO_ARGS,
                "-bsf:a",
                "setts=pts=DTS",
                *_encoder_args,
                "-shortest",
                "-fflags",
                "+genpts",
                "-f",
                "mpegts",
                file + ".ts",
            ]
        )
        logger.info(" ".join(ffmpeg_process.args))
        span.set_attribute("ffmpeg.args", json.dumps(ffmpeg_process.args))
        logger.info(
            "ReEncodeAndAnnexb: %s, returned: %s", file, ffmpeg_process.returncode
        )
        span.set_attribute("ffmpeg.code", ffmpeg_process.returncode)
        if ffmpeg_process.returncode == 0:
            span.set_status(Status(StatusCode.OK))
            span.set_attribute("file.size", os.path.getsize(file + ".ts"))
            # os.remove(file)
            return file + ".ts"
        else:
            span.set_status(Status(StatusCode.ERROR))
            return file


# start_render函数已迁移到services/render_service.py中的DefaultRenderService
# 保留原有签名用于向后兼容，但建议使用新的服务架构


def start_render(ffmpeg_task):
    """
    已迁移到新架构，建议使用 DefaultRenderService.render()
    保留用于向后兼容
    """
    logger.warning(
        "start_render is deprecated, use DefaultRenderService.render() instead"
    )
    from services import DefaultRenderService

    render_service = DefaultRenderService()
    return render_service.render(ffmpeg_task)


def handle_ffmpeg_output(stdout: Optional[bytes]) -> str:
    out_time = "0:0:0.0"
    if stdout is None:
        print("[!]STDOUT is null")
        return out_time
    speed = "0"
    for line in stdout.split(b"\n"):
        if line == b"":
            break
        if line.strip() == b"progress=end":
            # 处理完毕
            break
        if line.startswith(b"out_time="):
            out_time = line.replace(b"out_time=", b"").decode().strip()
        if line.startswith(b"speed="):
            speed = line.replace(b"speed=", b"").decode().strip()
    print("[ ]Speed:", out_time, "@", speed)
    return out_time + "@" + speed


def duration_str_to_float(duration_str: str) -> float:
    _duration = datetime.strptime(duration_str, "%H:%M:%S.%f") - datetime(1900, 1, 1)
    return _duration.total_seconds()


def probe_video_info(video_file):
    tracer = get_tracer(__name__)
    with tracer.start_as_current_span("probe_video_info") as span:
        span.set_attribute("video.file", video_file)
        # 获取宽度和高度
        result = subprocess.run(
            [
                "ffprobe",
                "-v",
                "error",
                "-select_streams",
                "v:0",
                "-show_entries",
                "stream=width,height:format=duration",
                "-of",
                "csv=s=x:p=0",
                video_file,
            ],
            stderr=subprocess.STDOUT,
            **subprocess_args(True)
        )
        span.set_attribute("ffprobe.args", json.dumps(result.args))
        span.set_attribute("ffprobe.code", result.returncode)
        if result.returncode != 0:
            span.set_status(Status(StatusCode.ERROR))
            return 0, 0, 0
        all_result = result.stdout.decode("utf-8").strip()
        span.set_attribute("ffprobe.out", all_result)
        if all_result == "":
            span.set_status(Status(StatusCode.ERROR))
            return 0, 0, 0
        span.set_status(Status(StatusCode.OK))
        wh, duration = all_result.split("\n")
        width, height = wh.strip().split("x")
        return int(width), int(height), float(duration)


def probe_video_audio(video_file, type=None):
    tracer = get_tracer(__name__)
    with tracer.start_as_current_span("probe_video_audio") as span:
        span.set_attribute("video.file", video_file)
        args = [
            "ffprobe",
            "-hide_banner",
            "-v",
            "error",
            "-select_streams",
            "a",
            "-show_entries",
            "stream=index",
            "-of",
            "csv=p=0",
        ]
        if type == "concat":
            args.append("-safe")
            args.append("0")
            args.append("-f")
            args.append("concat")
        args.append(video_file)
        logger.info(" ".join(args))
        result = subprocess.run(args, stderr=subprocess.STDOUT, **subprocess_args(True))
        span.set_attribute("ffprobe.args", json.dumps(result.args))
        span.set_attribute("ffprobe.code", result.returncode)
        logger.info("probe_video_audio: %s", result.stdout.decode("utf-8").strip())
        if result.returncode != 0:
            return False
        if result.stdout.decode("utf-8").strip() == "":
            return False
        return True


# 音频淡出2秒
def fade_out_audio(file, duration, fade_out_sec=2):
    if type(duration) == str:
        try:
            duration = float(duration)
        except Exception as e:
            logger.error("duration is not float: %s", e)
            return file
    tracer = get_tracer(__name__)
    with tracer.start_as_current_span("fade_out_audio") as span:
        span.set_attribute("audio.file", file)
        if duration <= fade_out_sec:
            return file
        else:
            new_fn = file + "_.mp4"
            if os.path.exists(new_fn):
                os.remove(new_fn)
                logger.info("delete tmp file: " + new_fn)
            try:
                process = subprocess.run(
                    [
                        "ffmpeg",
                        "-i",
                        file,
                        "-c:v",
                        "copy",
                        "-c:a",
                        "aac",
                        "-af",
                        "afade=t=out:st="
                        + str(duration - fade_out_sec)
                        + ":d="
                        + str(fade_out_sec),
                        "-y",
                        new_fn,
                    ],
                    **subprocess_args(True)
                )
                span.set_attribute("ffmpeg.args", json.dumps(process.args))
                logger.info(" ".join(process.args))
                if process.returncode != 0:
                    span.set_status(Status(StatusCode.ERROR))
                    logger.error("FFMPEG ERROR: %s", process.stderr)
                    return file
                else:
                    span.set_status(Status(StatusCode.OK))
                    return new_fn
            except Exception as e:
                span.set_status(Status(StatusCode.ERROR))
                logger.error("FFMPEG ERROR: %s", e)
                return file


# Create a set of arguments which make a ``subprocess.Popen`` (and
# variants) call work with or without Pyinstaller, ``--noconsole`` or
# not, on Windows and Linux. Typical use::
#
#   subprocess.call(['program_to_run', 'arg_1'], **subprocess_args())
#
# When calling ``check_output``::
#
#   subprocess.check_output(['program_to_run', 'arg_1'],
#                           **subprocess_args(False))
def subprocess_args(include_stdout=True):
    # The following is true only on Windows.
    if hasattr(subprocess, "STARTUPINFO"):
        # On Windows, subprocess calls will pop up a command window by default
        # when run from Pyinstaller with the ``--noconsole`` option. Avoid this
        # distraction.
        si = subprocess.STARTUPINFO()
        si.dwFlags |= subprocess.STARTF_USESHOWWINDOW
        # Windows doesn't search the path by default. Pass it an environment so
        # it will.
        env = os.environ
    else:
        si = None
        env = None

    # ``subprocess.check_output`` doesn't allow specifying ``stdout``::
    #
    #   Traceback (most recent call last):
    #     File "test_subprocess.py", line 58, in <module>
    #       **subprocess_args(stdout=None))
    #     File "C:\Python27\lib\subprocess.py", line 567, in check_output
    #       raise ValueError('stdout argument not allowed, it will be overridden.')
    #   ValueError: stdout argument not allowed, it will be overridden.
    #
    # So, add it only if it's needed.
    if include_stdout:
        ret = {"stdout": subprocess.PIPE}
    else:
        ret = {}

    # On Windows, running this from the binary produced by Pyinstaller
    # with the ``--noconsole`` option requires redirecting everything
    # (stdin, stdout, stderr) to avoid an OSError exception
    # "[Error 6] the handle is invalid."
    ret.update({"stdin": subprocess.PIPE, "startupinfo": si, "env": env})
    return ret