dnsmos.py

"""
dnsmos.py — Calidad percibida de la voz con DNSMOS P.835 (Microsoft).

Estima en una escala 1–5 lo que puntuarían oyentes humanos:
  SIG  = calidad de la voz (¿suena natural, sin distorsión?)
  BAK  = fondo (¿se oye ruido?)
  OVRL = valoración global
Trabaja a 16 kHz: evalúa la voz, no los agudos extremos.
Para que sea rápido en CPUs modestas se evalúan como mucho MAX_WINDOWS ventanas de 9 s
repartidas por el audio (la media se estabiliza con pocas).
Basado en DNSMOS/dnsmos_local.py del repositorio microsoft/DNS-Challenge.
"""

from __future__ import annotations

import os
import threading

import numpy as np

from models import dnsmos_paths

SR = 16000
WIN_SEC = 9.01
MAX_WINDOWS = 10

_lock = threading.Lock()
_sessions = None


def _load():
    global _sessions
    with _lock:
        if _sessions is None:
            import onnxruntime as ort

            paths = dnsmos_paths()
            opts = ort.SessionOptions()
            opts.log_severity_level = 3
            # En contenedores, onnxruntime puede lanzar más hilos que núcleos reales y ralentizarse.
            opts.intra_op_num_threads = max(1, min(4, os.cpu_count() or 1))
            opts.inter_op_num_threads = 1
            _sessions = ort.InferenceSession(str(paths["sig_bak_ovr.onnx"]), opts, providers=["CPUExecutionProvider"])
    return _sessions


def available() -> bool:
    try:
        _load()
        return True
    except Exception as e:
        print(f"[DNSMOS] No disponible: {e}")
        return False


def _polyfit(sig, bak, ovr):
    p_ovr = np.poly1d([-0.06766283, 1.11546468, 0.04602535])
    p_sig = np.poly1d([-0.08397278, 1.22083953, 0.0052439])
    p_bak = np.poly1d([-0.13166888, 1.60915514, -0.39604546])
    return p_sig(sig), p_bak(bak), p_ovr(ovr)


def score(mono: np.ndarray, sr: int) -> dict[str, float] | None:
    """Devuelve {'SIG','BAK','OVRL'} o None si el modelo no está disponible."""
    try:
        sess = _load()
    except Exception as e:
        print(f"[DNSMOS] No disponible: {e}")
        return None
    import soxr

    audio = soxr.resample(np.asarray(mono, dtype=np.float32), sr, SR, quality="HQ") if sr != SR else mono
    audio = np.asarray(audio, dtype=np.float32)
    if not np.any(audio):
        return None
    win = int(WIN_SEC * SR)
    while len(audio) < win:
        audio = np.concatenate([audio, audio])
    n_hops = int(np.floor(len(audio) / SR) - WIN_SEC) + 1
    starts = np.arange(n_hops) * SR
    if len(starts) > MAX_WINDOWS:
        starts = starts[np.linspace(0, len(starts) - 1, MAX_WINDOWS).round().astype(int)]
    rows = []
    for s in starts:
        seg = audio[int(s) : int(s) + win]
        if len(seg) < win:
            continue
        sig, bak, ovr = sess.run(None, {"input_1": seg[None, :]})[0][0]
        rows.append(_polyfit(sig, bak, ovr))
    if not rows:
        return None
    m = np.mean(rows, axis=0)
    return {"SIG": float(m[0]), "BAK": float(m[1]), "OVRL": float(m[2])}


def label(ovrl: float) -> str:
    if ovrl >= 3.6:
        return "excelente"
    if ovrl >= 3.2:
        return "buena"
    if ovrl >= 2.7:
        return "aceptable"
    return "mejorable"