Upload app.py

app.py

@@ -0,0 +1,525 @@
+import io
+import math
+import tempfile
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Dict, Optional, Tuple
+
+import gradio as gr
+import librosa
+import matplotlib.pyplot as plt
+import numpy as np
+import onnxruntime as ort
+import soundfile as sf
+from PIL import Image
+
+# -----------------------------
+# Configuration
+# -----------------------------
+MAX_SECONDS = 10.0
+ONNX_DIR = Path("./onnx")
+
+
+@dataclass(frozen=True)
+class ModelSpec:
+    name: str
+    sr: int
+    onnx_path: str
+
+
+# -----------------------------
+# Model discovery and metadata
+# -----------------------------
+def _infer_model_meta(model_name: str) -> int:
+    normalized = model_name.lower().replace("-", "_")
+
+    if "48khz" in normalized or "48k" in normalized or "48hr" in normalized:
+        return 48000
+
+    # Fallback for unknown 16 kHz DPDFNet variants
+    return 16000
+
+
+def _display_label(spec: ModelSpec) -> str:
+    khz = int(spec.sr // 1000)
+    return f"{spec.name} ({khz} kHz)"
+
+
+def discover_model_presets() -> Dict[str, ModelSpec]:
+    ordered_names = [
+        "baseline",
+        "dpdfnet2",
+        "dpdfnet4",
+        "dpdfnet8",
+        "dpdfnet2_48khz_hr",
+        "dpdfnet8_48khz_hr",
+    ]
+
+    found_paths = {p.stem: p for p in ONNX_DIR.glob("*.onnx") if p.is_file()}
+    presets: Dict[str, ModelSpec] = {}
+
+    for name in ordered_names:
+        p = found_paths.get(name)
+        if p is None:
+            continue
+        sr = _infer_model_meta(name)
+        spec = ModelSpec(
+            name=name,
+            sr=sr,
+            onnx_path=str(p),
+        )
+        presets[_display_label(spec)] = spec
+
+    # Include any additional ONNX files not in the canonical order list.
+    for name, p in sorted(found_paths.items()):
+        if name in ordered_names:
+            continue
+        sr = _infer_model_meta(name)
+        spec = ModelSpec(
+            name=name,
+            sr=sr,
+            onnx_path=str(p),
+        )
+        presets[_display_label(spec)] = spec
+
+    return presets
+
+
+MODEL_PRESETS = discover_model_presets()
+DEFAULT_MODEL_KEY = next(iter(MODEL_PRESETS), None)
+
+
+# -----------------------------
+# ONNX Runtime + frontend cache
+# -----------------------------
+_SESSIONS: Dict[str, ort.InferenceSession] = {}
+_INIT_STATES: Dict[str, np.ndarray] = {}
+
+
+def resolve_model_path(local_path: str) -> str:
+    p = Path(local_path)
+    if p.exists():
+        return str(p)
+    raise gr.Error(
+        f"ONNX model not found at: {local_path}. "
+        "Expected local models under ./onnx/."
+    )
+
+
+def get_ort_session(model_key: str) -> ort.InferenceSession:
+    if model_key in _SESSIONS:
+        return _SESSIONS[model_key]
+
+    spec = MODEL_PRESETS[model_key]
+    onnx_path = resolve_model_path(spec.onnx_path)
+
+    options = ort.SessionOptions()
+    options.intra_op_num_threads = 1
+    options.inter_op_num_threads = 1
+
+    sess = ort.InferenceSession(
+        onnx_path,
+        sess_options=options,
+        providers=["CPUExecutionProvider"],
+    )
+    _SESSIONS[model_key] = sess
+    return sess
+
+
+def _load_initial_state(model_key: str, session: ort.InferenceSession) -> np.ndarray:
+    if model_key in _INIT_STATES:
+        return _INIT_STATES[model_key]
+
+    if len(session.get_inputs()) < 2:
+        raise gr.Error("Expected streaming ONNX model with two inputs: (spec, state).")
+
+    meta = session.get_modelmeta().custom_metadata_map
+    try:
+        state_size = int(meta["state_size"])
+        erb_norm_state_size = int(meta["erb_norm_state_size"])
+        spec_norm_state_size = int(meta["spec_norm_state_size"])
+        erb_norm_init = np.array(
+            [float(x) for x in meta["erb_norm_init"].split(",")], dtype=np.float32
+        )
+        spec_norm_init = np.array(
+            [float(x) for x in meta["spec_norm_init"].split(",")], dtype=np.float32
+        )
+    except KeyError as exc:
+        raise gr.Error(
+            f"ONNX model is missing required metadata key: {exc}. "
+            "Re-export the model to embed state initialisation metadata."
+        )
+
+    init_state = np.zeros(state_size, dtype=np.float32)
+    init_state[0:erb_norm_state_size] = erb_norm_init
+    init_state[erb_norm_state_size:erb_norm_state_size + spec_norm_state_size] = spec_norm_init
+    init_state = np.ascontiguousarray(init_state)
+
+    _INIT_STATES[model_key] = init_state
+    return init_state
+
+
+# -----------------------------
+# STFT/iSTFT (module-free)
+# -----------------------------
+def vorbis_window(window_len: int) -> np.ndarray:
+    window_size_h = window_len / 2
+    indices = np.arange(window_len)
+    sin = np.sin(0.5 * np.pi * (indices + 0.5) / window_size_h)
+    window = np.sin(0.5 * np.pi * sin * sin)
+    return window.astype(np.float32)
+
+
+def _infer_stft_params(model_key: str, session: ort.InferenceSession) -> Tuple[int, int, np.ndarray]:
+    # ONNX spec input is [B, T, F, 2] (or dynamic variants).
+    spec_shape = session.get_inputs()[0].shape
+    freq_bins = spec_shape[-2] if len(spec_shape) >= 2 else None
+
+    if isinstance(freq_bins, int) and freq_bins > 1:
+        win_len = int((freq_bins - 1) * 2)
+    else:
+        # 20 ms windows for DPDFNet family.
+        sr = MODEL_PRESETS[model_key].sr
+        win_len = int(round(sr * 0.02))
+
+    hop = win_len // 2
+    win = vorbis_window(win_len)
+    return win_len, hop, win
+
+
+def _preprocess_waveform(waveform: np.ndarray, win_len: int, hop: int, win: np.ndarray) -> np.ndarray:
+    audio = np.asarray(waveform, dtype=np.float32).reshape(-1)
+    audio_pad = np.pad(audio, (0, win_len), mode="constant")
+
+    spec = librosa.stft(
+        y=audio_pad,
+        n_fft=win_len,
+        hop_length=hop,
+        win_length=win_len,
+        window=win,
+        center=True,
+        pad_mode="reflect",
+    )
+    spec = spec.T.astype(np.complex64, copy=False)  # [T, F]
+    spec_ri = np.stack([spec.real, spec.imag], axis=-1).astype(np.float32, copy=False)  # [T, F, 2]
+    return np.ascontiguousarray(spec_ri[None, ...], dtype=np.float32)  # [1, T, F, 2]
+
+
+def _postprocess_spec(spec_e: np.ndarray, win_len: int, hop: int, win: np.ndarray) -> np.ndarray:
+    spec_c = np.asarray(spec_e[0], dtype=np.float32)  # [T, F, 2]
+    spec = (spec_c[..., 0] + 1j * spec_c[..., 1]).T.astype(np.complex64, copy=False)  # [F, T]
+
+    waveform_e = librosa.istft(
+        spec,
+        hop_length=hop,
+        win_length=win_len,
+        window=win,
+        center=True,
+        length=None,
+    ).astype(np.float32, copy=False)
+
+    return np.concatenate(
+        [waveform_e[win_len * 2 :], np.zeros(win_len * 2, dtype=np.float32)],
+        axis=0,
+    )
+
+
+# -----------------------------
+# ONNX inference (non-streaming pre/post, streaming ONNX state loop)
+# -----------------------------
+def enhance_audio_onnx(
+    audio_mono: np.ndarray,
+    model_key: str,
+) -> np.ndarray:
+    sess = get_ort_session(model_key)
+
+    inputs = sess.get_inputs()
+    outputs = sess.get_outputs()
+    if len(inputs) < 2 or len(outputs) < 2:
+        raise gr.Error(
+            "Expected streaming ONNX signature with 2 inputs (spec, state) and 2 outputs (spec_e, state_out)."
+        )
+
+    in_spec_name = inputs[0].name
+    in_state_name = inputs[1].name
+    out_spec_name = outputs[0].name
+    out_state_name = outputs[1].name
+
+    waveform = np.asarray(audio_mono, dtype=np.float32).reshape(-1)
+    win_len, hop, win = _infer_stft_params(model_key, sess)
+    spec_r_np = _preprocess_waveform(waveform, win_len=win_len, hop=hop, win=win)
+
+    state = _load_initial_state(model_key, sess).copy()
+    spec_e_frames = []
+    num_frames = int(spec_r_np.shape[1])
+
+    for t in range(num_frames):
+        spec_t = np.ascontiguousarray(spec_r_np[:, t : t + 1, :, :], dtype=np.float32)
+        spec_e_t, state = sess.run(
+            [out_spec_name, out_state_name],
+            {in_spec_name: spec_t, in_state_name: state},
+        )
+        spec_e_frames.append(np.ascontiguousarray(spec_e_t, dtype=np.float32))
+
+    if not spec_e_frames:
+        return waveform
+
+    spec_e_np = np.concatenate(spec_e_frames, axis=1)
+    waveform_e = _postprocess_spec(spec_e_np, win_len=win_len, hop=hop, win=win)
+    return np.asarray(waveform_e, dtype=np.float32).reshape(-1)
+
+
+# -----------------------------
+# Audio utilities
+# -----------------------------
+def _load_wav_from_gradio_path(path: str) -> Tuple[np.ndarray, int]:
+    data, sr = sf.read(path, always_2d=True)
+    data = data.astype(np.float32, copy=False)
+    return data, int(sr)
+
+
+def _to_mono(x: np.ndarray) -> Tuple[np.ndarray, int]:
+    if x.ndim == 1:
+        return x.astype(np.float32, copy=False), 1
+    if x.shape[1] == 1:
+        return x[:, 0], 1
+    return x.mean(axis=1), int(x.shape[1])
+
+
+def _resample(y: np.ndarray, sr_in: int, sr_out: int) -> np.ndarray:
+    if sr_in == sr_out:
+        return y
+    return librosa.resample(y, orig_sr=sr_in, target_sr=sr_out).astype(np.float32, copy=False)
+
+
+def _match_length(y: np.ndarray, target_len: int) -> np.ndarray:
+    if len(y) == target_len:
+        return y
+    if len(y) > target_len:
+        return y[:target_len]
+    out = np.zeros((target_len,), dtype=y.dtype)
+    out[: len(y)] = y
+    return out
+
+
+def _save_wav(y: np.ndarray, sr: int, prefix: str) -> str:
+    tmp = tempfile.NamedTemporaryFile(prefix=prefix, suffix=".wav", delete=False)
+    tmp.close()
+    sf.write(tmp.name, y, sr)
+    return tmp.name
+
+
+def _spectrogram_image(y: np.ndarray, sr: int) -> Image.Image:
+    win_length = max(256, int(0.032 * sr))
+    hop_length = max(64, int(0.008 * sr))
+    n_fft = 1 << (int(math.ceil(math.log2(win_length))))
+
+    S = librosa.stft(y, n_fft=n_fft, hop_length=hop_length, win_length=win_length, center=False)
+    S_db = librosa.amplitude_to_db(np.abs(S) + 1e-10, ref=np.max)
+
+    fig, ax = plt.subplots(figsize=(8.4, 3.2))
+    ax.imshow(S_db, origin="lower", aspect="auto")
+    ax.set_axis_off()
+    fig.subplots_adjust(left=0, right=1, top=1, bottom=0)
+
+    buf = io.BytesIO()
+    fig.savefig(buf, format="png", dpi=160)
+    plt.close(fig)
+    buf.seek(0)
+    return Image.open(buf)
+
+
+# -----------------------------
+# Main pipeline
+# -----------------------------
+def run_enhancement(
+    source: str,
+    mic_path: Optional[str],
+    file_path: Optional[str],
+    model_key: str,
+):
+    if not MODEL_PRESETS:
+        raise gr.Error("No ONNX models found under ./onnx/. Add models and retry.")
+
+    chosen_path = mic_path if source == "Microphone" else file_path
+    if not chosen_path:
+        raise gr.Error("Please provide audio either from the microphone or by uploading a file.")
+
+    x, sr_orig = _load_wav_from_gradio_path(chosen_path)
+    y_mono, n_ch = _to_mono(x)
+
+    max_samples = int(MAX_SECONDS * sr_orig)
+    was_trimmed = len(y_mono) > max_samples
+    if was_trimmed:
+        y_mono = y_mono[:max_samples]
+    dur = len(y_mono) / float(sr_orig)
+
+    spec = MODEL_PRESETS[model_key]
+    sr_model = spec.sr
+
+    y_model = _resample(y_mono, sr_orig, sr_model)
+    y_enh_model = enhance_audio_onnx(y_model, model_key)
+
+    y_enh = _resample(y_enh_model, sr_model, sr_orig)
+    y_enh = _match_length(y_enh, len(y_mono))
+
+    noisy_out = _save_wav(y_mono, sr_orig, prefix="noisy_mono_")
+    enh_out = _save_wav(y_enh, sr_orig, prefix="enhanced_")
+
+    noisy_img = _spectrogram_image(y_mono, sr_orig)
+    enh_img = _spectrogram_image(y_enh, sr_orig)
+
+    status = (
+        f"**Input:** {sr_orig} Hz, {dur:.2f}s, channels={n_ch} ⭢ mono\n\n"
+        f"**Model:** {spec.name} (runs at {sr_model} Hz)\n\n"
+        + (
+            f"**Resampling:** {sr_orig} ⭢ {sr_model} ⭢ {sr_orig}\n\n"
+            if sr_orig != sr_model
+            else "**Resampling:** none\n\n"
+        )
+        + (f"**Trimmed:** first {MAX_SECONDS:.0f}s used\n" if was_trimmed else "")
+        + "\n✅ Done."
+    )
+    return noisy_out, enh_out, noisy_img, enh_img, status
+
+
+def set_source_visibility(source: str):
+    return (
+        gr.update(visible=(source == "Microphone")),
+        gr.update(visible=(source == "Upload")),
+    )
+
+
+# -----------------------------
+# UI (light polish)
+# -----------------------------
+THEME = gr.themes.Soft(
+    primary_hue="orange",
+    neutral_hue="slate",
+    font=[
+        "Arial",
+        "ui-sans-serif",
+        "system-ui",
+        "Segoe UI",
+        "Roboto",
+        "Helvetica Neue",
+        "Noto Sans",
+        "Liberation Sans",
+        "sans-serif",
+    ],
+)
+
+CSS = """
+.gradio-container{
+  max-width: 1040px !important;
+  margin: 0 auto !important;
+  font-family: Arial, ui-sans-serif, system-ui, -apple-system, Segoe UI, Roboto, Helvetica Neue, Noto Sans, Liberation Sans, sans-serif !important;
+}
+
+#header {
+  padding: 14px 16px;
+  border-radius: 16px;
+  border: 1px solid rgba(0,0,0,0.08);
+  background: linear-gradient(135deg, rgba(255,152,0,0.14), rgba(255,152,0,0.04));
+  text-align: center;
+}
+#header h1{
+  margin: 0 0 6px 0;
+  font-size: 24px;
+  font-weight: 800;
+  letter-spacing: -0.2px;
+}
+#header p{
+  margin: 6px auto 0 auto;
+  max-width: 720px;
+  color: var(--body-text-color-subdued);
+  font-size: 14px;
+  line-height: 1.6;
+}
+#header hr{
+  margin-top: 18px;
+  border: none;
+  height: 1px;
+  background: linear-gradient(to right, transparent, #ddd, transparent);
+}
+
+.spec img { border-radius: 14px; }
+.audio { border-radius: 14px !important; overflow: hidden; }
+
+#run_btn{
+  border-radius: 12px !important;
+  font-weight: 800 !important;
+}
+
+#status_md p{ margin: 0.35rem 0; }
+"""
+
+with gr.Blocks(theme=THEME, css=CSS, title="DPDFNet Speech Enhancement") as demo:
+    gr.Markdown(
+        "# DPDFNet Speech Enhancement\n\n"
+        "Causal · Real-Time · Edge-Ready\n\n"
+        "DPDFNet extends DeepFilterNet2 with Dual-Path RNN blocks to improve "
+        "long-range temporal and cross-band modeling while preserving low latency. "
+        "Designed for single-channel streaming speech enhancement under challenging noise conditions.\n\n"
+        "---",
+        elem_id="header",
+    )
+
+    with gr.Row():
+        model_key = gr.Dropdown(
+            choices=list(MODEL_PRESETS.keys()),
+            value=DEFAULT_MODEL_KEY,
+            label="Model",
+            # info="Audio is resampled to model SR, enhanced with ONNX, then resampled back.",
+            interactive=True,
+        )
+
+        source = gr.Radio(
+            choices=["Microphone", "Upload"],
+            value="Upload",
+            label="Input source",
+        )
+
+    with gr.Row():
+        mic_audio = gr.Audio(
+            sources=["microphone"],
+            type="filepath",
+            format="wav",
+            label="Microphone (max 10s)",
+            visible=False,
+            buttons=["download"],
+            elem_classes=["audio"],
+        )
+        file_audio = gr.Audio(
+            sources=["upload"],
+            type="filepath",
+            format="wav",
+            label="Upload file (WAV/MP3/FLAC etc., max 10s)",
+            visible=True,
+            buttons=["download"],
+            elem_classes=["audio"],
+        )
+
+    run_btn = gr.Button("Enhance", variant="primary", elem_id="run_btn")
+    status = gr.Markdown(elem_id="status_md")
+
+    gr.Markdown("## Results")
+
+    with gr.Row():
+        out_noisy = gr.Audio(label="Before (mono)", interactive=False, format="wav", buttons=["download"], elem_classes=["audio"])
+        out_enh = gr.Audio(label="After (enhanced)", interactive=False, format="wav", buttons=["download"], elem_classes=["audio"])
+
+    with gr.Row():
+        img_noisy = gr.Image(label="Noisy spectrogram", elem_classes=["spec"])
+        img_enh = gr.Image(label="Enhanced spectrogram", elem_classes=["spec"])
+
+    source.change(fn=set_source_visibility, inputs=source, outputs=[mic_audio, file_audio])
+    run_btn.click(
+        fn=run_enhancement,
+        inputs=[source, mic_audio, file_audio, model_key],
+        outputs=[out_noisy, out_enh, img_noisy, img_enh, status],
+        api_name="enhance",
+    )
+
+if __name__ == "__main__":
+    demo.queue(max_size=32).launch()