evOo-local/infra/media-worker/app/ops/audio.py

"""audioFeatures without ML: ffmpeg ebur128 + silencedetect, librosa beat tracking.

Audio is decoded once to mono 22.05 kHz float32 and processed in five-minute
chunks, so memory stays bounded and the request can stop between chunks.
"""

from __future__ import annotations

import re
import subprocess
import threading
from typing import Any

import numpy as np

from ..errors import Cancelled, MediaError
from ..media_io import FFMPEG, parse_ebur128, run_cancellable
from ..schemas import PathRequest
from .common import Ctx, open_input

SAMPLE_RATE = 22050
CHUNK_SECONDS = 300
ENERGY_HOP_MS = 250
SILENCE_DB = -40
SILENCE_MIN_S = 0.5


def loudness_and_silence(
    path: str, cancel: threading.Event
) -> tuple[float | None, float | None, list[dict[str, int]]]:
    result = run_cancellable(
        [
            FFMPEG,
            "-hide_banner",
            "-nostats",
            "-nostdin",
            "-i",
            path,
            "-map",
            "0:a:0",
            "-vn",
            "-af",
            f"ebur128=peak=true,silencedetect=noise={SILENCE_DB}dB:d={SILENCE_MIN_S}",
            "-f",
            "null",
            "-",
        ],
        cancel,
    )
    stderr = result.stderr.decode("utf-8", "replace")
    if result.returncode != 0:
        raise MediaError("internal", "loudness measurement failed")
    lufs, peak = parse_ebur128(stderr)
    silences: list[dict[str, int]] = []
    start: float | None = None
    for line in stderr.splitlines():
        if (m := re.search(r"silence_start:\s*(-?[\d.]+)", line)) is not None:
            start = max(float(m.group(1)), 0.0)
        elif (m := re.search(r"silence_end:\s*([\d.]+)", line)) is not None and start is not None:
            silences.append(
                {"startMs": int(round(start * 1000)), "endMs": int(round(float(m.group(1)) * 1000))}
            )
            start = None
    if start is not None:
        silences.append({"startMs": int(round(start * 1000)), "endMs": -1})
    return lufs, peak, silences


def decode_chunks(path: str, cancel: threading.Event) -> Any:
    argv = [
        FFMPEG,
        "-v",
        "error",
        "-nostdin",
        "-i",
        path,
        "-map",
        "0:a:0",
        "-vn",
        "-ac",
        "1",
        "-ar",
        str(SAMPLE_RATE),
        "-f",
        "f32le",
        "pipe:1",
    ]
    proc = subprocess.Popen(argv, stdout=subprocess.PIPE, stderr=subprocess.DEVNULL, shell=False)
    chunk_bytes = SAMPLE_RATE * CHUNK_SECONDS * 4
    assert proc.stdout is not None
    try:
        while True:
            if cancel.is_set():
                raise Cancelled()
            buf = proc.stdout.read(chunk_bytes)
            if not buf:
                break
            usable = len(buf) - len(buf) % 4
            yield np.frombuffer(buf[:usable], dtype="<f4").astype(np.float32)
        if proc.wait() != 0:
            raise MediaError("internal", "audio decode failed")
    finally:
        if proc.poll() is None:
            proc.kill()
            proc.wait()


def run(ctx: Ctx, body: PathRequest, cancel: threading.Event) -> dict[str, Any]:
    import librosa  # deferred: heavy import (numba), only needed here

    real, info = open_input(ctx, body.path, cancel, need_video=False)
    ctx.registry.runtime("audioFeatures")
    if not info.has_audio:
        return {
            "integratedLufs": None,
            "truePeakDb": None,
            "tempoBpm": None,
            "beatsMs": [],
            "energy": [],
            "silences": [],
        }
    lufs, peak, silences = loudness_and_silence(real, cancel)

    hop = int(SAMPLE_RATE * ENERGY_HOP_MS / 1000)
    beats: list[int] = []
    tempos: list[tuple[float, float]] = []
    rms_values: list[float] = []
    offset_s = 0.0
    for samples in decode_chunks(real, cancel):
        length_s = samples.size / SAMPLE_RATE
        if samples.size >= SAMPLE_RATE * 2 and float(np.abs(samples).max()) > 1e-4:
            tempo, frames = librosa.beat.beat_track(y=samples, sr=SAMPLE_RATE, units="frames")
            times = librosa.frames_to_time(frames, sr=SAMPLE_RATE)
            beats.extend(int(round((offset_s + t) * 1000)) for t in times)
            value = float(np.atleast_1d(tempo)[0])
            if value > 0:
                tempos.append((value, length_s))
        for i in range(0, samples.size, hop):
            window = samples[i : i + hop]
            rms_values.append(float(np.sqrt(np.mean(window.astype(np.float64) ** 2))) if window.size else 0.0)
        offset_s += length_s

    duration_ms = int(round(offset_s * 1000))
    for silence in silences:
        if silence["endMs"] < 0:
            silence["endMs"] = duration_ms
    peak_rms = max(rms_values) if rms_values else 0.0
    energy = [
        {"tMs": i * ENERGY_HOP_MS, "value": (v / peak_rms) if peak_rms > 0 else 0.0}
        for i, v in enumerate(rms_values)
    ]
    tempo_bpm = None
    if tempos:
        weights = np.array([w for _, w in tempos])
        tempo_bpm = float(np.average([t for t, _ in tempos], weights=weights))
    return {
        "integratedLufs": lufs,
        "truePeakDb": peak,
        "tempoBpm": tempo_bpm,
        "beatsMs": sorted(set(beats)),
        "energy": energy,
        "silences": silences,
    }