UnheardCity/ackit/audio.py

"""Portable audio helpers (no external programs needed).

Audio is handled as float32 numpy arrays shaped (frames, channels).
Decoding uses libsndfile (WAV, FLAC, OGG/Vorbis/Opus, MP3, AIFF, ...); other formats (M4A, WMA, ...)
fall back to an ffmpeg found on the system, if any. Resampling uses soxr.
"""
import io
import os
import shutil
import subprocess

import numpy as np
import soundfile as sf


class AudioError(Exception):
    pass


def _as2d(x):
    x = np.asarray(x, dtype=np.float32)
    return x[:, None] if x.ndim == 1 else x


def load(path):
    """Read an audio file -> (float32 array (n, ch), rate)."""
    try:
        x, sr = sf.read(path, dtype="float32", always_2d=True)
        return x, sr
    except (sf.LibsndfileError, RuntimeError, TypeError) as e:
        ff = shutil.which("ffmpeg")
        if not ff:
            raise AudioError("Cannot read '%s' (%s). Supported: WAV, FLAC, OGG, OPUS, MP3, AIFF. "
                             "Convert the file to one of these, or install ffmpeg for other formats."
                             % (os.path.basename(path), e))
        raw = subprocess.run([ff, "-v", "error", "-i", path, "-f", "wav", "-"], capture_output=True)
        if raw.returncode:
            raise AudioError("ffmpeg could not read '%s': %s" % (path, raw.stderr.decode(errors="replace")[-300:]))
        x, sr = sf.read(io.BytesIO(raw.stdout), dtype="float32", always_2d=True)
        return x, sr


def save_wav(path, x, rate):
    sf.write(path, _as2d(x), rate, subtype="PCM_16")


def resample(x, rate, new_rate):
    x = _as2d(x)
    if rate == new_rate or len(x) == 0:
        return x
    import soxr
    return soxr.resample(x, rate, new_rate, quality="HQ").astype(np.float32)


def to_channels(x, channels):
    """Up/down-mix. Up: mono is copied to every channel, other sources keep their channels and the new ones get
    the average. Down: nothing is dropped - every channel goes into the mono sum, or the channels beyond the
    kept ones (e.g. the centre and surrounds of 5.1) are added to every kept one at -3 dB (the standard
    fold-down); a result louder than the source's peak is scaled back to it, so it does not clip."""
    x = _as2d(x)
    cur = x.shape[1]
    if cur == channels:
        return x
    if cur == 1:
        return np.repeat(x, channels, axis=1)
    if channels > cur:
        out = np.empty((len(x), channels), dtype=np.float32)
        out[:, :cur] = x
        out[:, cur:] = x.mean(axis=1, keepdims=True)
        return out
    if channels == 1:
        out = x.sum(axis=1, keepdims=True)
    else:
        out = x[:, :channels] + 0.7071 * x[:, channels:].sum(axis=1, keepdims=True)
    peak_in, peak_out = float(np.max(np.abs(x), initial=0.0)), float(np.max(np.abs(out), initial=0.0))
    if peak_out > peak_in > 0:
        out = out * (peak_in / peak_out)
    return out.astype(np.float32)


def trim(x, rate, start=None, end=None):
    x = _as2d(x)
    a = int(round((start or 0) * rate))
    b = len(x) if end is None else int(round(end * rate))
    return x[max(0, a):max(a, min(len(x), b))]


def gain(x, db):
    return (_as2d(x) * (10 ** (db / 20.0))).astype(np.float32)


def peak_db(x):
    p = float(np.max(np.abs(x))) if len(x) else 0.0
    return -120.0 if p <= 0 else 20 * np.log10(p)


def rms_db(x):
    r = float(np.sqrt(np.mean(np.square(x)))) if len(x) else 0.0
    return -120.0 if r <= 0 else 20 * np.log10(r)


def normalize(x, peak=-1.0):
    """Scale so the highest peak is `peak` dBFS."""
    p = peak_db(x)
    return _as2d(x) if p <= -119 else gain(x, peak - p)


def match_loudness(x, reference):
    """Scale x to the RMS loudness of `reference` (limited so it does not clip)."""
    g = rms_db(reference) - rms_db(x)
    g = min(g, -0.3 - peak_db(x))
    return gain(x, g)


def trim_silence(x, threshold_db=-50.0, pad=0.02, rate=48000):
    x = _as2d(x)
    lvl = np.max(np.abs(x), axis=1)
    idx = np.nonzero(lvl > 10 ** (threshold_db / 20.0))[0]
    if len(idx) == 0:
        return x
    p = int(pad * rate)
    return x[max(0, idx[0] - p):min(len(x), idx[-1] + p)]


def conform(x, rate, new_rate, channels):
    """Resample + channel conversion + clip guard, ready for encoding."""
    y = to_channels(resample(x, rate, new_rate), channels)
    return np.clip(y, -1.0, 1.0)


def decode_wem(data: bytes):
    """Decode a game WEM -> (float32 array (n, ch), rate), using the game's codebooks."""
    from .convert import wem_to_ogg
    from .wem import Wem
    ogg, _ = wem_to_ogg(Wem.parse(data))
    x, sr = sf.read(io.BytesIO(ogg), dtype="float32", always_2d=True)
    return x, sr


def duration(x, rate):
    return len(x) / float(rate) if rate else 0.0