airdef/tools/acoustics_check.py

#!/usr/bin/env python3
"""Проверка акустики без слуха (#96): энергия по полосам в дорожке с преградой и без неё.

    python3 tools/acoustics_check.py open.raw wall.raw            # s16le, 48000 Гц, стерео
    python3 tools/acoustics_check.py open.raw wall.raw --rate 48000 --channels 2

Дорожки берутся из `Audio::film_mix` (экспорт видео пишет `<имя>.audio.raw`) или из теста
`cargo test -p horizon-client --lib audio::tests::export_wall_tracks -- --ignored` (HORIZON_ACOUSTICS_OUT=папка).
Для каждой дорожки считается энергия левого канала в четырёх полосах (до 300 Гц, 300-1500, 1500-5000, выше
5000) и разница в дБ между дорожками. Преграда должна срезать верхи (выше 5 кГц не меньше чем на 15 дБ) и
оставить низы (до 300 Гц теряют не больше 8 дБ). Коды выхода: 0 — в порогах, 1 — нет, 2 — ошибка аргументов.
Режим `--echo` (стадия 2): в дорожке с одним щелчком ищутся пики огибающей после прямого звука, их задержки
(мс от прямого) сравниваются с ожидаемыми от геометрии:

    python3 tools/acoustics_check.py --echo gorge.raw --expect 187.3,254.1 [--tol 1.0]
    python3 tools/acoustics_check.py --echo field.raw --expect none
    python3 tools/acoustics_check.py --echo-dir папка        # gorge/field/room и expected.txt из export_echo_tracks

Дорожки пишет `audio::tests::export_echo_tracks` (HORIZON_ACOUSTICS_OUT=папка). Коды выхода те же.
Только стандартная библиотека.
"""
import argparse
import array
import cmath
import math
import sys

BANDS = [("<300", 0, 300), ("300-1500", 300, 1500), ("1500-5000", 1500, 5000), (">5000", 5000, 1e9)]


def fft(a):
    n = len(a)
    if n == 1:
        return a
    even, odd = fft(a[0::2]), fft(a[1::2])
    t = [cmath.exp(-2j * math.pi * k / n) * odd[k] for k in range(n // 2)]
    return [even[k] + t[k] for k in range(n // 2)] + [even[k] - t[k] for k in range(n // 2)]


def bands(path, rate, channels):
    raw = array.array("h")
    with open(path, "rb") as f:
        data = f.read()
    raw.frombytes(data[: len(data) // 2 * 2])
    left = [s / 32768.0 for s in raw[0::channels]]
    n = 1
    while n * 2 <= len(left):
        n *= 2
    n = min(n, 1 << 16)
    if n < 1024:
        sys.exit(f"{path}: too short ({len(left)} frames)")
    # A Hann window over the loudest n frames' start: the burst is at the start of the track.
    w = [left[i] * (0.5 - 0.5 * math.cos(2 * math.pi * i / (n - 1))) for i in range(n)]
    spec = fft([complex(v) for v in w])
    out = [0.0] * len(BANDS)
    for k in range(1, n // 2):
        hz = k * rate / n
        for i, (_, lo, hi) in enumerate(BANDS):
            if lo <= hz < hi:
                out[i] += abs(spec[k]) ** 2
    return out


def echo_peaks(path, rate, channels, min_gap_ms=10.0, rel=0.02):
    """Delays (ms from the first peak) of the peaks of the smoothed envelope of a track."""
    raw = array.array("h")
    with open(path, "rb") as f:
        data = f.read()
    raw.frombytes(data[: len(data) // 2 * 2])
    frames = [sum(abs(raw[i + c]) for c in range(channels)) / 32768.0 for i in range(0, len(raw) - channels + 1, channels)]
    win = max(1, rate // 1000)  # 1 ms
    env = []
    acc = 0.0
    for i, v in enumerate(frames):
        acc += v
        if i >= win:
            acc -= frames[i - win]
        env.append(acc / win)
    first = max(range(min(len(env), rate // 20)), key=env.__getitem__)  # the direct sound: loudest in the first 50 ms
    top = env[first]
    gap = int(min_gap_ms * rate / 1000)
    peaks, i = [], first + int(0.02 * rate)
    while i < len(env):
        if env[i] >= rel * top:
            j = max(range(i, min(len(env), i + gap)), key=env.__getitem__)
            peaks.append((j - first) * 1000.0 / rate)
            i = j + gap
        else:
            i += 1
    return peaks


def check_echo(path, expect, tol, rate, channels):
    peaks = echo_peaks(path, rate, channels)
    print(f"{path}: peaks at {', '.join(f'{p:.1f}' for p in peaks) or 'none'} ms; expected {', '.join(f'{e:.1f}' for e in expect) or 'none'} ms")
    ok = len(peaks) == len(expect) and all(abs(p - e) <= tol for p, e in zip(peaks, expect))
    print(f"  {'ok' if ok else 'FAIL'} (tolerance {tol} ms)")
    return ok


def db(a, b):
    return 10 * math.log10(max(a, 1e-20) / max(b, 1e-20))


def main():
    p = argparse.ArgumentParser(description=__doc__.split("\n")[0])
    p.add_argument("open", nargs="?")
    p.add_argument("wall", nargs="?")
    p.add_argument("--echo", help="track with one click: check the delays of its echoes")
    p.add_argument("--echo-dir", help="folder written by export_echo_tracks: check gorge, field and room")
    p.add_argument("--expect", default="none", help="expected delays in ms, comma separated, or none")
    p.add_argument("--tol", type=float, default=1.0, help="ms a peak may be off")
    p.add_argument("--rate", type=int, default=48000)
    p.add_argument("--channels", type=int, default=2)
    p.add_argument("--min-high-cut", type=float, default=15.0, help="dB the band above 5 kHz must drop")
    p.add_argument("--max-low-loss", type=float, default=8.0, help="dB the band below 300 Hz may drop")
    a = p.parse_args()
    if a.echo_dir:
        import os
        ok = True
        with open(os.path.join(a.echo_dir, "expected.txt")) as f:
            for line in f:
                name, _, vals = line.strip().partition(" ")
                want = [float(v) for v in vals.split(",") if v]
                # The ms in the room are for the walls the geometry says; only the first peak is compared there.
                ok &= check_echo(os.path.join(a.echo_dir, name + ".raw"), want, a.tol, a.rate, a.channels)
        return 0 if ok else 1
    if a.echo:
        want = [] if a.expect == "none" else [float(v) for v in a.expect.split(",")]
        return 0 if check_echo(a.echo, want, a.tol, a.rate, a.channels) else 1
    if not a.open or not a.wall:
        p.error("give open.raw and wall.raw, or --echo / --echo-dir")
    o, w = bands(a.open, a.rate, a.channels), bands(a.wall, a.rate, a.channels)
    print(f"{'band, Hz':>10} {'open dB':>9} {'wall dB':>9} {'change dB':>10}")
    ref = max(o)
    for (name, _, _), x, y in zip(BANDS, o, w):
        print(f"{name:>10} {db(x, ref):9.1f} {db(y, ref):9.1f} {db(y, x):10.1f}")
    low, high = db(w[0], o[0]), db(w[-1], o[-1])
    ok = high <= -a.min_high_cut and low >= -a.max_low_loss
    print(f"highs {high:.1f} dB (need <= -{a.min_high_cut}), lows {low:.1f} dB (need >= -{a.max_low_loss}): {'ok' if ok else 'FAIL'}")
    return 0 if ok else 1


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