cinema-spatial/离线渲染.py

#!/usr/bin/env python3
"""离线渲染 —— 把音频/视频处理成空间音频 FLAC。

链路与实时虚拟声卡一致:
    立体声 --pan上混--> 5.1 --sofalizer(H4 SOFA)--> 双耳立体声 --[混响卷积]--> 响度归一化 --> flac

输入:
    音频(mp3/flac/wav/m4a/ogg/opus/aac...) 直接处理;
    视频(mp4/mkv/mov/webm...) 只取音轨(ffmpeg 选流, 不落地中间文件);
    已经是 5.1 的片源不做上混, 直接进 HRTF。

★ 锁(老板要求):
    1) 渲染锁: 一轮任务处理期间**不接受新任务**(拿不到锁就退出, 不排队堆积);
    2) 文件就绪检查: 每个文件进队列前查"是否还在写入/下载"(两次 stat + ffprobe),
       没落盘完的直接跳过 —— 不会把半个文件扔进渲染。

用法:
    python3 离线渲染.py 片子.mp4 歌.mp3        # 渲染(串行, 实时进度)
    python3 离线渲染.py --状态                 # 看是否在跑 / 上次结果
    python3 离线渲染.py --混响 0.3 歌.mp3      # 指定混响湿量(0 = 不要混响)
    python3 离线渲染.py --目标 -16 歌.mp3      # 响度目标(dBFS)
    python3 离线渲染.py --无HRTF 歌.flac       # 只上混 5.1, 不做双耳空间化
    python3 离线渲染.py --输出目录 ~/音乐/空间 歌.mp3
"""
from __future__ import annotations

import argparse
import fcntl
import json
import os
import shutil
import subprocess
import sys
import time
from pathlib import Path
from typing import IO

ROOT = Path(__file__).resolve().parent


def _first_existing(cands: list[Path]) -> Path | None:
    """按顺序取第一个存在的文件 —— 开发目录与 deb 包里的素材布局不一样, 两种都要认。"""
    for c in cands:
        if c.is_file():
            return c
    return None


# SOFA: 开发目录在 sofa/H4_HRIR_SOFA/H4_HRIR_SOFA/, deb 包里直接躺在 sofa/ 下
SOFA_96K = _first_existing([
    ROOT / "sofa" / "H4_96K_24bit_512tap_FIR_SOFA.sofa",
    ROOT / "sofa" / "H4_HRIR_SOFA" / "H4_HRIR_SOFA" / "H4_96K_24bit_512tap_FIR_SOFA.sofa",
    Path("/usr/share/cinema-spatial/sofa/H4_96K_24bit_512tap_FIR_SOFA.sofa"),
])
# 混响 IR 同理(包里是 /usr/share/cinema-spatial/reverb/)
REVERB_IR = _first_existing([
    ROOT / "reverb" / "房间混响IR-96k.wav",
    Path("/usr/share/cinema-spatial/reverb/房间混响IR-96k.wav"),
])
LOCK_PATH = Path("/tmp/collaplex-render.lock")
STATE_PATH = Path(os.path.expanduser("~/.local/state/cinema-spatial/render.json"))

# ★ 进度文件(给面板轮询): 带 --进度json <路径> 时每 0.4 秒原子写一次, 面板读它画进度条。
PROGRESS_PATH: Path | None = None


def write_progress(data: dict) -> None:
    """原子写进度 JSON —— 面板在同时读, 不能让它读到半个文件。"""
    if PROGRESS_PATH is None:
        return
    try:
        tmp = PROGRESS_PATH.with_suffix(".tmp")
        with open(tmp, "w", encoding="utf-8") as fh:
            json.dump(data, fh, ensure_ascii=False)
        tmp.replace(PROGRESS_PATH)
    except OSError:
        pass

VIDEO_EXT = {".mp4", ".mkv", ".mov", ".webm", ".avi", ".flv", ".wmv", ".m4v", ".ts", ".mpg", ".mpeg", ".rmvb"}
AUDIO_EXT = {".mp3", ".flac", ".wav", ".m4a", ".aac", ".ogg", ".opus", ".wma", ".aiff", ".ape", ".dff", ".dsf"}
OUT_SR = 96000


class Opt:
    """一次渲染的选项。"""

    def __init__(self, wet: float, target: float, hrtf: bool, outdir: Path | None,
                 overwrite: bool, quiet: bool) -> None:
        self.wet = wet                  # 混响湿量(0 = 关)
        self.target = target            # 响度目标 dBFS
        self.hrtf = hrtf                # 是否做双耳 HRTF
        self.outdir = outdir            # 输出目录(None = 源文件同级)
        self.overwrite = overwrite      # 已存在是否覆盖
        self.quiet = quiet              # 安静模式(只报结果)


# ------------------------------------------------------------------ 锁

def acquire_lock() -> IO[str] | None:
    """拿渲染锁。拿不到 -> None(已有任务在跑, 调用方直接退出)。"""
    fd = os.open(str(LOCK_PATH), os.O_CREAT | os.O_RDWR, 0o644)
    try:
        fcntl.flock(fd, fcntl.LOCK_EX | fcntl.LOCK_NB)
    except OSError:
        os.close(fd)
        return None
    fh = os.fdopen(fd, "w+", encoding="utf-8")
    fh.truncate(0)
    fh.write("%d\n%s\n" % (os.getpid(), time.strftime("%F %T")))
    fh.flush()
    return fh


def runing_pid() -> str:
    """另一个进程是否正持有渲染锁; 返回它的 pid 描述(空 = 空闲)。"""
    if not LOCK_PATH.exists():
        return ""
    try:
        fd = os.open(str(LOCK_PATH), os.O_RDWR)
    except OSError:
        return ""
    try:
        fcntl.flock(fd, fcntl.LOCK_EX | fcntl.LOCK_NB)
        fcntl.flock(fd, fcntl.LOCK_UN)
        return ""
    except OSError:
        try:
            with open(LOCK_PATH, encoding="utf-8") as fh:
                first = fh.readline().strip()
            return first or "未知"
        except OSError:
            return "未知"
    finally:
        os.close(fd)


# ------------------------------------------------------------------ 探测

def ffprobe_json(path: Path) -> dict:
    """ffprobe -show_format -show_streams 的 JSON; 失败返回空 dict。"""
    try:
        r = subprocess.run(
            ["ffprobe", "-v", "error", "-print_format", "json",
             "-show_format", "-show_streams", str(path)],
            capture_output=True, text=True, timeout=60)
        return json.loads(r.stdout or "{}")
    except (subprocess.TimeoutExpired, json.JSONDecodeError, OSError):
        return {}


def audio_layout(path: Path) -> tuple[str, float]:
    """(音轨的声道布局, 时长秒); 读不到返回 ("", 0.0)。"""
    info = ffprobe_json(path)
    for st in info.get("streams") or []:
        if st.get("codec_type") == "audio":
            layout = str(st.get("channel_layout") or "")
            ch = int(st.get("channels") or 0)
            if not layout:
                layout = {1: "mono", 2: "stereo", 6: "5.1", 8: "7.1"}.get(ch, "stereo")
            dur = 0.0
            for key in ("duration",):
                try:
                    dur = float(st.get(key) or 0.0)
                except (TypeError, ValueError):
                    dur = 0.0
            if dur <= 0.0:
                try:
                    dur = float((info.get("format") or {}).get("duration") or 0.0)
                except (TypeError, ValueError):
                    dur = 0.0
            return (layout, dur)
    return ("", 0.0)


def file_ready(path: Path) -> tuple[bool, str]:
    """文件是否已经落盘完成。

    ★ 老板要求的第二道锁: 防"文件还在内存里没下载下来就被扔进渲染"。
      判据: ① 存在且是常规文件且非空
            ② 相隔 1.2 秒两次 stat, size 和 mtime 都不变(还在写会一直变)
            ③ ffprobe 能读出音轨(半个文件通常读不出)
    """
    if not path.exists():
        return (False, "文件不存在")
    if not path.is_file():
        return (False, "不是常规文件(目录/设备?)")
    try:
        s1 = path.stat()
    except OSError as e:
        return (False, "读不到文件属性: %s" % e)
    if s1.st_size == 0:
        return (False, "文件是空的(0 字节)")
    time.sleep(1.2)
    try:
        s2 = path.stat()
    except OSError as e:
        return (False, "文件在处理期间消失: %s" % e)
    if s2.st_size != s1.st_size or s2.st_mtime != s1.st_mtime:
        return (False, "文件还在写入/下载中(大小还在变), 等它传完再扔")
    layout, _dur = audio_layout(path)
    if not layout:
        return (False, "读不到音轨(可能没传完, 或不是音视频文件)")
    return (True, "")


# ------------------------------------------------------------------ 滤镜图

def build_filter(layout: str, src_idx: int, opt: Opt, with_reverb: bool) -> str:
    """拼 ffmpeg 滤镜图。

    与实时链对应关系:
      pan 5.1        <- 生成配置.py 里的上混矩阵(FL/FR 直通, FC=(FL+FR)*0.707,
                        LFE 同源衰减, BL/BR 反相)
      sofalizer      <- 12 个 HRTF 卷积 + mixL/mixR(6 路各 1/sqrt(6) 能量归一)
      afir(房间IR)   <- 混响湿路
      loudnorm       <- 响度归一化 DSP
    """
    src = "[%d:a]" % src_idx
    chain = "%saformat=sample_rates=%d" % (src, OUT_SR)

    if layout in ("5.1", "5.1(side)", "7.1", "6.1", "7.1(wide)"):
        chain += ",aformat=channel_layouts=5.1"          # 片源本就多声道, 不重上混
    elif layout == "mono":
        chain += ",aformat=channel_layouts=stereo,pan=5.1|FL=FL|FR=FR|FC=0.707*FL+0.707*FR|"
        chain += "LFE=0.5*FL+0.5*FR|BL=0.707*FL-0.707*FR|BR=0.707*FR-0.707*FL"
    else:
        chain += (",aformat=channel_layouts=stereo,pan=5.1|FL=FL|FR=FR|FC=0.707*FL+0.707*FR|"
                  "LFE=0.5*FL+0.5*FR|BL=0.707*FL-0.707*FR|BR=0.707*FR-0.707*FL")

    if opt.hrtf and SOFA_96K is not None:
        chain += ",sofalizer=sofa=%s:type=freq" % SOFA_96K
    else:
        chain += ",pan=stereo|FL=FL|FR=FR"               # 不做双耳: 5.1 下混成立体声

    graph = "%s[hr]" % chain
    if with_reverb:
        graph += (";[hr]asplit=2[rdry][rwet]"
                  ";[%d:a]aformat=sample_rates=%d:channel_layouts=mono[rir]"
                  ";[rwet][rir]afir=dry=0:wet=1:gtype=none:irnorm=1[rwet2]"
                  ";[rdry][rwet2]amix=inputs=2:weights=1 %.4f:normalize=0[hr2]"
                  % (src_idx + 1, OUT_SR, opt.wet))
        hr = "[hr2]"
    else:
        hr = "[hr]"
    graph += ";%sloudnorm=I=%.1f:TP=-1.5:LRA=11[out]" % (hr, opt.target)
    return graph


def render(src: Path, dst: Path, opt: Opt) -> bool:
    """渲染一个文件; 返回是否成功。实时打印进度。"""
    layout, dur = audio_layout(src)
    if not layout:
        print("  ✗ 读不到音轨, 跳过")
        return False
    with_reverb = opt.wet > 0.0 and REVERB_IR is not None
    graph = build_filter(layout, 0, opt, with_reverb)

    dst.parent.mkdir(parents=True, exist_ok=True)
    cmd = ["ffmpeg", "-hide_banner", "-v", "error", "-nostdin", "-y",
           "-i", str(src)]
    if with_reverb:
        cmd += ["-i", str(REVERB_IR)]
    cmd += ["-filter_complex", graph, "-map", "[out]",
            "-c:a", "flac", "-ar", str(OUT_SR), "-compression_level", "8",
            "-progress", "pipe:1", "-nostats", str(dst)]

    kind = "视频(取音轨)" if src.suffix.lower() in VIDEO_EXT else "音频"
    spk = " + 混响 %.2f" % opt.wet if with_reverb else ""
    print("  · %s  %s  %s  时长 %.1fs" % (kind, layout, spk, dur))
    t0 = time.time()
    proc = subprocess.Popen(cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True)
    last_draw = 0.0
    assert proc.stdout is not None
    for line in proc.stdout:
        # ★ ffmpeg 的 out_time_ms 实际单位是**微秒**(历史遗留命名), 用 out_time_us 才不歧义 ——
        #   按毫秒算会把 8 秒的素材显示成 4779 秒、进度瞬间冲到 100%(2026-09-13 实测)。
        if not line.startswith("out_time_us="):
            continue
        try:
            us = int(line.split("=", 1)[1].strip())
        except ValueError:
            continue
        if opt.quiet or dur <= 0.0:
            continue
        now = time.time()
        if now - last_draw < 0.4:
            continue
        last_draw = now
        sec = max(0.0, us / 1e6)
        pct = max(0.0, min(100.0, sec / dur * 100.0))
        speed = sec / max(0.001, now - t0)
        bar = "█" * int(pct / 4) + "·" * (25 - int(pct / 4))
        sys.stdout.write("\r    %s %5.1f%%  %.1fx  %4.0fs/%.0fs   "
                         % (bar, pct, speed, sec, dur))
        sys.stdout.flush()
        write_progress({"状态": "渲染中", "文件": src.name, "百分比": round(pct, 1),
                        "速度": round(speed, 2), "当前秒": round(sec, 1), "总秒": round(dur, 1)})
    proc.wait()
    err = (proc.stderr.read() if proc.stderr else "") or ""
    if not opt.quiet and dur > 0.0:
        sys.stdout.write("\r" + " " * 72 + "\r")
    if proc.returncode != 0 or not dst.exists():
        print("  ✗ 失败: %s" % (err.strip().split("\n")[-1] if err.strip() else "ffmpeg 退出码 %d" % proc.returncode))
        return False
    mb = dst.stat().st_size / 1048576.0
    print("  ✓ 完成 -> %s  (%.1f MB, 用时 %.1fs)" % (dst.name, mb, time.time() - t0))
    return True


# ------------------------------------------------------------------ 状态

def write_state(data: dict) -> None:
    try:
        STATE_PATH.parent.mkdir(parents=True, exist_ok=True)
        with open(STATE_PATH, "w", encoding="utf-8") as fh:
            json.dump(data, fh, ensure_ascii=False, indent=2)
    except OSError:
        pass


def show_state() -> None:
    pid = runing_pid()
    print("渲染锁: %s" % ("正被 PID %s 持有(有任务在跑)" % pid if pid else "空闲"))
    if STATE_PATH.exists():
        try:
            with open(STATE_PATH, encoding="utf-8") as fh:
                d = json.load(fh)
            print("上次: %s" % d.get("结束时间", "?"))
            print("  成功 %s / 失败 %s" % (d.get("成功", 0), d.get("失败", 0)))
            for f in d.get("文件", [])[:20]:
                print("   %s %s" % (f.get("状态", "?"), f.get("名称", "")))
        except (OSError, json.JSONDecodeError):
            print("上次: (状态文件读不出)")
    if not shutil.which("ffmpeg"):
        print("! 找不到 ffmpeg")


# ------------------------------------------------------------------ main

def main() -> int:
    ap = argparse.ArgumentParser(
        description="离线渲染: 音频/视频 -> 空间音频 FLAC",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="示例:\n  python3 离线渲染.py 片子.mp4 歌.mp3\n  python3 离线渲染.py --混响 0.3 歌.flac")
    ap.add_argument("files", nargs="*", help="要处理的音频/视频文件(可多个)")
    ap.add_argument("--状态", action="store_true", help="显示是否在跑 / 上次结果")
    ap.add_argument("--混响", type=float, default=0.3, help="混响湿量 0~1 (默认 0.3, 0 = 关)")
    ap.add_argument("--目标", type=float, default=-14.0, help="响度目标 dBFS (默认 -14)")
    ap.add_argument("--无HRTF", action="store_true", help="只上混 5.1, 不做双耳空间化")
    ap.add_argument("--输出目录", default="", help="输出目录(默认与源文件同级)")
    ap.add_argument("--覆盖", action="store_true", help="已存在同名输出时覆盖(默认跳过)")
    ap.add_argument("--安静", action="store_true", help="不打进度条, 只报结果")
    ap.add_argument("--进度json", default="", help="把进度写到该 JSON 文件(面板轮询用)")
    args = ap.parse_args()

    global PROGRESS_PATH
    if args.进度json:
        PROGRESS_PATH = Path(args.进度json).expanduser()

    if args.状态:
        show_state()
        return 0
    if not args.files:
        ap.print_help()
        return 1
    if not shutil.which("ffmpeg"):
        print("! 找不到 ffmpeg, 无法渲染")
        return 1
    if SOFA_96K is None and not args.无HRTF:
        print("! 找不到 SOFA 文件(开发目录 sofa/ 或 /usr/share/cinema-spatial/sofa/); "
              "可用 --无HRTF 跳过空间化")
        return 1

    # ★ 第一道锁: 渲染期间不接受新任务 —— 拿不到锁直接退出, 不排队堆积
    lock = acquire_lock()
    if lock is None:
        print("⛔ 已有一个渲染任务在跑(PID %s), 本次不接 —— 等它跑完再扔。" % runing_pid())
        return 2

    outdir = Path(args.输出目录).expanduser() if args.输出目录 else None
    opt = Opt(wet=max(0.0, min(1.0, args.混响)), target=args.目标, hrtf=not args.无HRTF,
              outdir=outdir, overwrite=args.覆盖, quiet=args.安静)

    print("离线渲染 —— 上混 5.1%s -> 响度目标 %.1f dBFS -> FLAC %d kHz"
          % (" -> SOFA 双耳 HRTF" if opt.hrtf else "", opt.target, OUT_SR // 1000))
    if opt.wet > 0:
        print("混响湿量 %.2f (%s)" % (opt.wet, REVERB_IR.name if REVERB_IR else "找不到混响 IR, 本次不叠加"))
    print("渲染锁已取得(PID %d); 期间不接受新任务\n" % os.getpid())

    ok: list[dict] = []
    skipped: list[dict] = []
    failed: list[dict] = []
    t_all = time.time()
    write_progress({"状态": "开始", "总数": len(args.files), "文件": "", "百分比": 0,
                    "已处理": 0, "成功": 0, "跳过": 0, "失败": 0})
    try:
        for i, raw in enumerate(args.files, start=1):
            src = Path(raw).expanduser()
            print("[%d/%d] %s" % (i, len(args.files), src.name))
            write_progress({"状态": "检查中", "文件": src.name, "百分比": 0, "序号": i,
                            "总数": len(args.files), "成功": len(ok), "跳过": len(skipped),
                            "失败": len(failed)})
            # ★ 第二道锁: 文件没落盘完(还在下载/写入)就不扔进渲染
            ready, why = file_ready(src)
            if not ready:
                print("  ⏭ 跳过: %s" % why)
                skipped.append({"名称": src.name, "状态": "跳过", "原因": why})
                continue
            if src.suffix.lower() not in AUDIO_EXT | VIDEO_EXT:
                print("  ⚠ 后缀不认识, 仍按音频尝试: %s" % src.suffix)
            dst = (opt.outdir or src.parent) / ("%s.空间音频.flac" % src.stem)
            if dst.exists() and not opt.overwrite:
                print("  ⏭ 跳过: 已存在 %s (要覆盖加 --覆盖)" % dst.name)
                skipped.append({"名称": src.name, "状态": "跳过", "原因": "输出已存在"})
                continue
            if render(src, dst, opt):
                ok.append({"名称": src.name, "状态": "成功", "输出": str(dst)})
            else:
                failed.append({"名称": src.name, "状态": "失败"})
    finally:
        lock.close()

    print("\n—— 完成: 成功 %d / 跳过 %d / 失败 %d, 总用时 %.1fs ——"
          % (len(ok), len(skipped), len(failed), time.time() - t_all))
    for f in failed:
        print("   失败: %s" % f["名称"])
    for f in skipped:
        print("   跳过: %s (%s)" % (f["名称"], f.get("原因", "")))
    write_progress({"状态": "完成", "成功": len(ok), "跳过": len(skipped),
                    "失败": len(failed), "百分比": 100})
    write_state({"结束时间": time.strftime("%F %T"), "成功": len(ok), "失败": len(failed),
                 "文件": ok + skipped + failed})
    return 0 if not failed else 3


if __name__ == "__main__":
    sys.exit(main())