audio_analyzer.py

"""Deterministic audio feature extraction (CPU, no model).

The model never hears raw audio. It reasons over the precise numerical
description produced here. Everything in this file is reproducible AND robust:
no input (corrupt file, silence, NaN, single sample, clipping, stereo, wrong
sample rate, hours-long upload) can make it raise or emit a non-finite value.
"""
from __future__ import annotations

from dataclasses import dataclass, asdict

import numpy as np

SR = 22050
MAX_DURATION_S = 10.0
N_FFT = 2048  # librosa default frame; we pad shorter clips up to this


@dataclass
class AudioFeatures:
    duration_s: float
    rms_db: float                 # overall loudness
    rms_variance: float           # loudness variation (high = intermittent)
    zero_crossing_rate: float     # high = harsh/grinding, low = tonal
    spectral_centroid_hz: float   # high = bright/harsh, low = rumbling
    spectral_bandwidth_hz: float  # wide = complex/noisy
    spectral_rolloff_hz: float    # freq below which 85% of energy sits
    dominant_frequency_hz: float  # strongest fundamental
    harmonic_ratio: float         # 1.0 = pure tone, 0.0 = pure noise
    onset_rate_per_sec: float     # clicks/knocks per second
    has_regular_pattern: bool     # evenly spaced clicks (bearing signature)
    pattern_interval_ms: float    # interval between events if regular
    peak_db: float                # loudest instant (clipping risk)
    anomaly_score: float          # 0-1 heuristic "abnormality"
    signal_present: bool = True   # False = too quiet/short/empty to trust

    def to_dict(self) -> dict:
        return asdict(self)


# --- sanitization helpers ---------------------------------------------------
def _num(value, default: float, lo: float, hi: float) -> float:
    """Coerce to a finite float clamped to [lo, hi]; default if not finite."""
    try:
        v = float(value)
    except (TypeError, ValueError):
        return float(default)
    if not np.isfinite(v):
        return float(default)
    return float(min(max(v, lo), hi))


def _safe_db(x) -> float:
    try:
        return float(20 * np.log10(max(float(x), 0.0) + 1e-8))
    except Exception:
        return -120.0


NYQUIST = SR / 2.0

# Returned when audio is unusable (empty / silence / all-NaN / load failure).
_NEUTRAL = dict(
    duration_s=0.0, rms_db=-120.0, rms_variance=0.0, zero_crossing_rate=0.0,
    spectral_centroid_hz=0.0, spectral_bandwidth_hz=0.0, spectral_rolloff_hz=0.0,
    dominant_frequency_hz=0.0, harmonic_ratio=0.0, onset_rate_per_sec=0.0,
    has_regular_pattern=False, pattern_interval_ms=0.0, peak_db=-120.0,
    anomaly_score=0.0, signal_present=False,
)


def _neutral() -> AudioFeatures:
    return AudioFeatures(**_NEUTRAL)


def _load_audio(audio_path):
    """Load mono audio at SR, capped to MAX_DURATION_S. Returns y or None."""
    if not audio_path or not isinstance(audio_path, str):
        return None
    try:
        import librosa
        y, _ = librosa.load(audio_path, sr=SR, duration=MAX_DURATION_S, mono=True)
    except Exception:
        return None
    if y is None or len(y) == 0 or not np.any(np.isfinite(y)):
        return None
    y = np.nan_to_num(y, nan=0.0, posinf=0.0, neginf=0.0).astype(np.float32)
    # Reject effectively-silent input (peak below ~ -55 dBFS).
    if float(np.max(np.abs(y))) < 1.8e-3:
        return None
    return y


def extract_features(audio_path: str) -> AudioFeatures:
    """Extract ~14 deterministic features. Never raises; always finite."""
    import librosa

    y = _load_audio(audio_path)
    if y is None:
        return _neutral()

    duration_s = float(len(y) / SR)
    # Pad short clips so framed transforms have a full window to work with.
    y_proc = y if len(y) >= N_FFT else np.pad(y, (0, N_FFT - len(y)))

    def _agg(fn, default):
        try:
            return float(np.nanmean(fn()))
        except Exception:
            return float(default)

    rms = None
    try:
        rms = librosa.feature.rms(y=y_proc)[0]
        rms_db = _safe_db(np.nanmean(rms))
        rms_var = _num(np.nanvar(rms), 0.0, 0.0, 1e6)
    except Exception:
        rms_db, rms_var = -120.0, 0.0

    zcr = _agg(lambda: librosa.feature.zero_crossing_rate(y=y_proc), 0.0)
    centroid = _agg(lambda: librosa.feature.spectral_centroid(y=y_proc, sr=SR), 0.0)
    bandwidth = _agg(lambda: librosa.feature.spectral_bandwidth(y=y_proc, sr=SR), 0.0)
    rolloff = _agg(lambda: librosa.feature.spectral_rolloff(y=y_proc, sr=SR), 0.0)

    # Dominant fundamental via pitch tracking (pyin can be all-NaN or raise).
    dominant_f0 = 0.0
    try:
        f0, _, _ = librosa.pyin(y_proc, fmin=30, fmax=4000, sr=SR)
        if f0 is not None and np.any(np.isfinite(f0)):
            dominant_f0 = float(np.nanmean(f0))
    except Exception:
        dominant_f0 = 0.0

    # Harmonic vs percussive energy split.
    try:
        y_harm, _ = librosa.effects.hpss(y_proc)
        harm_ratio = float(np.mean(np.abs(y_harm)) / (np.mean(np.abs(y_proc)) + 1e-8))
    except Exception:
        harm_ratio = 0.0

    # Onset detection (clicks, knocks). delta>0 suppresses noise-floor peaks.
    try:
        onsets = librosa.onset.onset_detect(y=y_proc, sr=SR, units="time", delta=0.3)
    except Exception:
        onsets = np.array([])
    onset_rate = float(len(onsets) / duration_s) if duration_s > 0 else 0.0

    # Regular spacing of onsets = mechanical periodicity (bearing fault).
    has_pattern = False
    pattern_interval = 0.0
    if len(onsets) > 2:
        intervals = np.diff(onsets)
        mean_iv = float(np.mean(intervals))
        if np.isfinite(mean_iv) and mean_iv > 0 and float(np.std(intervals)) < 0.05 * mean_iv:
            has_pattern = True
            pattern_interval = mean_iv * 1000.0

    peak_db = _safe_db(np.max(np.abs(y)))

    # Heuristic anomaly score (transparent, not a model output).
    anomaly = rms_var * 10 + abs(centroid - 2000) / 5000 + onset_rate / 20

    return AudioFeatures(
        duration_s=_num(duration_s, 0.0, 0.0, MAX_DURATION_S),
        rms_db=_num(rms_db, -120.0, -120.0, 20.0),
        rms_variance=_num(rms_var, 0.0, 0.0, 1e6),
        zero_crossing_rate=_num(zcr, 0.0, 0.0, 1.0),
        spectral_centroid_hz=_num(centroid, 0.0, 0.0, NYQUIST),
        spectral_bandwidth_hz=_num(bandwidth, 0.0, 0.0, SR),
        spectral_rolloff_hz=_num(rolloff, 0.0, 0.0, NYQUIST),
        dominant_frequency_hz=_num(dominant_f0, 0.0, 0.0, NYQUIST),
        harmonic_ratio=_num(harm_ratio, 0.0, 0.0, 1.0),
        onset_rate_per_sec=_num(onset_rate, 0.0, 0.0, 1000.0),
        has_regular_pattern=bool(has_pattern),
        pattern_interval_ms=_num(pattern_interval, 0.0, 0.0, 60000.0),
        peak_db=_num(peak_db, -120.0, -120.0, 6.0),
        anomaly_score=_num(anomaly, 0.0, 0.0, 1.0),
        signal_present=True,
    )