_audit_noise.py

"""Score DNS5 noise files for music-likeness vs broadband-noise-likeness.

Music has structured pitch + rhythmic onset patterns that differentiate it
from HVAC / fan / wind / room-tone noise. We don't need a perfect classifier
— just a ranker that surfaces music-free candidates for a demo clip.

Features (computed on the first 10 s):
- spectral_flatness    : 1.0 = white noise; ~0.3 = pitched. Higher = noise-like.
- harmonicity          : ratio of largest non-trivial autocorrelation peak to
                          mean. High = pitched (music or speech). Lower = noise.
- chroma_entropy       : entropy of pitch-class distribution. High = energy
                          spread across all 12 classes (= noise). Low = a few
                          notes (= music).
- onset_periodicity    : autocorrelation peak of the onset-strength envelope at
                          beat-rate lags. High = music with steady beat.
- music_score          : weighted combo (low = clean noise, high = music).

Usage (inside Docker):
    LOCALVQE_CONTAINER=localvqe-audit ./train/scripts/docker-run.sh \\
        python /workspace/localvqe/space/_audit_noise.py
"""
from pathlib import Path

import numpy as np
import soundfile as sf

NOISE = Path("/workspace/localvqe/datasets_fullband/noise")
SR = 16000
CLIP = 10 * SR

# Top stationary-broadband candidates from the earlier ranking
# (high spectral flatness + low spectral flux). These are the files we'd
# pick from for the demo; just need to filter out the music-ones.
CANDIDATES = [
    "p1BmrTosB9w.flac", "ilVoxqe_z7E.flac", "iPGbDshbeuI.flac",
    "jyrMW_21Ocs.flac", "G3DdnBL0waA.flac", "G3fcLfe2HWs.flac",
    "EXM-wgH46OI.flac", "_dReWTQ4WqE.flac", "OuPtpakIs2k.flac",
    "BTSQlC4wVWg.flac", "r4XYK1WQKFo.flac", "_gjBxWy7v4Y.flac",
    "iqbiV-ttlA8.flac",
    # Throw in the bad pick for calibration
    "4Kud1014Cmg.flac",
]


def load_clip(path):
    wav, sr = sf.read(str(path), dtype="float32")
    if wav.ndim > 1:
        wav = wav.mean(axis=1)
    if sr != SR:
        # crude resample — fine for feature extraction
        from scipy.signal import resample_poly
        from math import gcd
        g = gcd(sr, SR)
        wav = resample_poly(wav, SR // g, sr // g).astype(np.float32)
    n = min(CLIP, len(wav))
    return wav[:n]


def spectral_flatness(wav, n_fft=1024, hop=512):
    n_frames = (len(wav) - n_fft) // hop + 1
    flat = []
    for i in range(n_frames):
        frame = wav[i * hop:i * hop + n_fft] * np.hanning(n_fft)
        spec = np.abs(np.fft.rfft(frame))
        spec = spec + 1e-12
        gm = np.exp(np.mean(np.log(spec)))
        am = np.mean(spec)
        flat.append(gm / am)
    return float(np.mean(flat))


def harmonicity(wav, n_fft=2048, hop=1024, min_lag=40, max_lag=400):
    """Ratio of strongest non-trivial autocorrelation peak to mean."""
    n_frames = (len(wav) - n_fft) // hop + 1
    ratios = []
    for i in range(n_frames):
        frame = wav[i * hop:i * hop + n_fft]
        if np.sqrt(np.mean(frame ** 2)) < 1e-4:
            continue
        # Auto-corr via FFT
        x = frame - frame.mean()
        ac = np.correlate(x, x, mode="full")[len(x) - 1:]
        ac = ac / (ac[0] + 1e-12)
        seg = ac[min_lag:max_lag]
        peak = float(seg.max()) if len(seg) > 0 else 0.0
        ratios.append(peak)
    return float(np.mean(ratios)) if ratios else 0.0


def chroma_entropy(wav, n_fft=2048, hop=1024):
    """Entropy of the chroma vector averaged over time. High = noise (flat
    pitch-class energy); low = music (energy in a few classes)."""
    n_frames = (len(wav) - n_fft) // hop + 1
    if n_frames <= 0:
        return float("nan")
    chroma_sum = np.zeros(12, dtype=np.float64)
    for i in range(n_frames):
        frame = wav[i * hop:i * hop + n_fft] * np.hanning(n_fft)
        spec = np.abs(np.fft.rfft(frame))
        # Map FFT bins to MIDI pitch -> chroma
        bins = np.arange(1, len(spec))
        freqs = bins * SR / n_fft
        midi = 69 + 12 * np.log2(freqs / 440.0 + 1e-12)
        valid = (midi >= 24) & (midi <= 108)
        midi = midi[valid].astype(int)
        spec = spec[1:][valid]
        for cls in range(12):
            chroma_sum[cls] += spec[midi % 12 == cls].sum()
    p = chroma_sum / (chroma_sum.sum() + 1e-12)
    p = p[p > 0]
    return float(-np.sum(p * np.log(p)) / np.log(12))  # normalised to [0, 1]


def onset_periodicity(wav, n_fft=1024, hop=256):
    """Autocorrelation peak of the onset-strength envelope at beat-rate lags
    (60-180 BPM band, lags 22-66 frames at this hop=256 / sr=16000)."""
    # Frame energy
    n_frames = (len(wav) - n_fft) // hop + 1
    E = np.zeros(n_frames, dtype=np.float32)
    for i in range(n_frames):
        E[i] = np.sqrt(np.mean(wav[i * hop:i * hop + n_fft] ** 2))
    # Onset envelope: rectified frame-to-frame energy difference
    onset = np.maximum(np.diff(E), 0)
    if onset.std() < 1e-6:
        return 0.0
    onset = (onset - onset.mean()) / (onset.std() + 1e-9)
    ac = np.correlate(onset, onset, mode="full")[len(onset) - 1:]
    ac = ac / (ac[0] + 1e-12)
    # Lags ~22-66 frames at hop=256, sr=16000 => 0.35-1.05s period => ~57-170 BPM
    band = ac[22:66]
    return float(band.max()) if len(band) > 0 else 0.0


def main():
    rows = []
    for name in CANDIDATES:
        p = NOISE / name
        if not p.exists():
            continue
        try:
            wav = load_clip(p)
            if len(wav) < SR:
                continue
            flat = spectral_flatness(wav)
            harm = harmonicity(wav)
            ce = chroma_entropy(wav)
            op = onset_periodicity(wav)
            # Music score: high harmonicity + low chroma entropy + high onset
            # periodicity all push up. Low spectral flatness (more pitched)
            # also pushes up. Hand-tuned weights (no fitted classifier; this
            # is a ranker, not a calibrated probability).
            music_score = (1.0 - flat) + harm + (1.0 - ce) + op
            rows.append((name, flat, harm, ce, op, music_score))
        except Exception as e:
            print(f"  {name}: error {e}")
    rows.sort(key=lambda r: r[5])  # ascending music_score = noise-first
    print(f"{'file':<24} {'flat':>5} {'harm':>5} {'ce':>5} {'beat':>5} {'music_score':>11}")
    for r in rows:
        print(f"{r[0]:<24} {r[1]:>5.2f} {r[2]:>5.2f} {r[3]:>5.2f} {r[4]:>5.2f} {r[5]:>11.2f}")


if __name__ == "__main__":
    main()