Fix feature extraction and models to match notebook

ml/feature_extraction.py

@@ -1,162 +1,220 @@
 """
-Feature extraction — 82-dimensional feature vector per hit.
-
-Breakdown:
-  [0:50]   Welch PSD — 50 log-spaced frequency bins (relative, sum=1)
-  [50:63]  MFCC mean — 13 coefficients (mean-subtracted per coefficient)
-  [63:76]  MFCC std  — 13 coefficients
-  [76]     Decay time constant τ (exponential fit)
-  [77]     Energy ratio (late/early)
-  [78]     Crest factor (peak/RMS)
-  [79]     Log peak amplitude
-  [80:82]  Spectral centroid mean + std
+Feature extraction — 82-dimensional physics-informed vector.
+EXACTLY matches Cell 4 of final_project_saurav_silwal.ipynb.
+
+Group 1: Relative PSD in 20 log-spaced bins (50 Hz – 8 kHz)    → 20 dims
+Group 2: MFCC mean/std + delta MFCC mean/std (13 coeffs each)  → 52 dims
+Group 3: Physics features (centroid, bandwidth, rolloff, ZCR,
+         peak freq, decay τ, energy ratio, RMS, Q-factor)       → 10 dims
+Total                                                            → 82 dims
 """
 
 import warnings
 import numpy as np
 import librosa
 from scipy.signal import welch
-from scipy.optimize import curve_fit
 
-from config import SR, N_MFCC, N_FFT, HIT_WINDOW_LEN, N_MELS, HOP_LENGTH_MEL, SPEC_TIME_FRAMES
+from config import SR, N_MELS, N_FFT, HOP_LENGTH_MEL, SPEC_TIME_FRAMES
 
-# ─── PSD ──────────────────────────────────────────────────────────────────────
+# ── Constants matching notebook ────────────────────────────────────────────
+N_PSD_BINS     = 20
+PSD_FMIN       = 50.0
+PSD_FMAX       = 8000.0
+WELCH_NPERSEG  = 2048
+WELCH_NOVERLAP = 1024
 
-def extract_psd(window: np.ndarray, n_bins: int = 50) -> np.ndarray:
-    """
-    Welch PSD normalised to sum=1 (relative PSD), log-spaced frequency bins.
-    Normalisation makes features robust to recording-level differences.
-    """
-    f, pxx = welch(window, fs=SR, nperseg=512, noverlap=256)
-    # Log-spaced bin edges from 50 Hz to Nyquist
-    edges = np.logspace(np.log10(50), np.log10(SR / 2), n_bins + 1)
-    binned = np.zeros(n_bins, dtype=np.float32)
+N_MFCC         = 13
+MFCC_NFFT      = 2048
+MFCC_HOP       = 512
+
+DECAY_FIT_MS       = 200
+EARLY_LATE_FRAC    = 0.20
+PEAK_SAMPLE_IN_WIN = int(0.020 * SR)   # 960 samples = 20 ms pre-peak
+
+FMIN_MEL = 0
+FMAX_MEL = SR // 2   # Nyquist = 24 000 Hz
+
+
+# ── Group 1: Relative PSD ─────────────────────────────────────────────────
+
+def relative_psd_log_bins(y, sr=SR, n_bins=N_PSD_BINS,
+                          f_min=PSD_FMIN, f_max=PSD_FMAX):
+    """Welch PSD → 20 log-spaced bins → normalized so sum=1."""
+    f, pxx = welch(y, fs=sr,
+                   nperseg=min(WELCH_NPERSEG, len(y)),
+                   noverlap=min(WELCH_NOVERLAP, len(y) // 2))
+    edges = np.logspace(np.log10(f_min), np.log10(f_max), n_bins + 1)
+    bins  = np.zeros(n_bins, dtype=np.float32)
     for i in range(n_bins):
-        mask = (f >= edges[i]) & (f < edges[i + 1])
-        if mask.any():
-            binned[i] = float(pxx[mask].mean())
-    total = binned.sum()
-    if total > 0:
-        binned /= total
-    return binned
-
-
-# ─── MFCC ─────────────────────────────────────────────────────────────────────
-
-def extract_mfcc(window: np.ndarray, n_mfcc: int = N_MFCC) -> tuple[np.ndarray, np.ndarray]:
-    """
-    Returns (mfcc_mean, mfcc_std) — both shape (n_mfcc,).
-    Mean subtraction (per coefficient across the window) is applied first —
-    this is the single most impactful normalisation for cross-session robustness.
-    """
+        mask    = (f >= edges[i]) & (f < edges[i + 1])
+        bins[i] = pxx[mask].sum()
+    total = bins.sum()
+    if total > 1e-20:
+        bins /= total
+    return bins, f, pxx
+
+
+# ── Group 2: MFCC + delta statistics ──────────────────────────────────────
+
+def mfcc_stats(y, sr=SR):
+    """13 MFCCs → mean+std (26) + delta mean+std (26) = 52 dims."""
     with warnings.catch_warnings():
         warnings.simplefilter("ignore")
-        mfcc = librosa.feature.mfcc(y=window, sr=SR, n_mfcc=n_mfcc, n_fft=N_FFT)
-    # Per-coefficient mean subtraction
-    mfcc -= mfcc.mean(axis=1, keepdims=True)
-    return mfcc.mean(axis=1).astype(np.float32), mfcc.std(axis=1).astype(np.float32)
-
-
-# ─── Decay ────────────────────────────────────────────────────────────────────
-
-def _exp_model(t, A, tau):
-    return A * np.exp(-t / tau)
-
-
-def extract_decay(window: np.ndarray) -> float:
-    """
-    Fit A·exp(-t/τ) to the RMS energy envelope.
-    Returns τ (seconds). Larger τ = slower decay = tighter flange.
-    Falls back to 0.05 on fit failure.
-    """
-    # RMS in 5ms frames
-    frame_len = int(0.005 * SR)
-    n_frames  = len(window) // frame_len
-    rms_env   = np.array([
-        np.sqrt(np.mean(window[i * frame_len:(i + 1) * frame_len] ** 2))
-        for i in range(n_frames)
-    ], dtype=np.float32)
-    t = np.arange(n_frames) * 0.005
-    # Trim to non-zero region
-    thresh = 0.02 * rms_env.max()
-    valid  = rms_env > thresh
-    if valid.sum() < 5:
-        return 0.05
-    try:
-        p0 = (rms_env[valid].max(), 0.15)
-        popt, _ = curve_fit(_exp_model, t[valid], rms_env[valid], p0=p0,
-                             bounds=([0, 0.001], [np.inf, 2.0]), maxfev=2000)
-        tau = float(np.clip(popt[1], 0.001, 2.0))
-    except Exception:
-        tau = 0.05
-    return tau
-
-
-# ─── Energy ratio ─────────────────────────────────────────────────────────────
-
-def extract_energy_ratio(window: np.ndarray, split_ms: float = 50.0) -> float:
-    """
-    E_late / E_early where split is at split_ms ms after the hit onset.
-    High ratio → flange still ringing → tight.
-    """
-    split_n = int(split_ms / 1000 * SR)
-    early = float(np.sum(window[:split_n] ** 2)) + 1e-12
-    late  = float(np.sum(window[split_n:] ** 2))
-    return min(late / early, 100.0)  # cap for numerical safety
-
-
-# ─── Spectral centroid ────────────────────────────────────────────────────────
-
-def extract_spectral_centroid(window: np.ndarray) -> tuple[float, float]:
-    cent = librosa.feature.spectral_centroid(y=window, sr=SR, n_fft=N_FFT)[0]
-    return float(cent.mean()), float(cent.std())
-
-
-# ─── Full 82-dim vector ───────────────────────────────────────────────────────
-
-def extract_features(window: np.ndarray) -> np.ndarray:
-    """Return 82-dim feature vector for a single hit window."""
-    psd          = extract_psd(window)                              # 50
-    mfcc_m, mfcc_s = extract_mfcc(window)                         # 13 + 13
-    tau          = extract_decay(window)                            # 1
-    energy_ratio = extract_energy_ratio(window)                    # 1
-    peak_amp     = float(np.abs(window).max())
-    rms          = float(np.sqrt(np.mean(window ** 2))) + 1e-12
-    crest        = float(peak_amp / rms)                            # 1
-    log_peak     = float(np.log1p(peak_amp))                       # 1
-    sc_mean, sc_std = extract_spectral_centroid(window)            # 2
-
-    feat = np.concatenate([
-        psd,
-        mfcc_m,
-        mfcc_s,
-        [tau, energy_ratio, crest, log_peak, sc_mean, sc_std],
+        mfcc  = librosa.feature.mfcc(y=y, sr=sr, n_mfcc=N_MFCC,
+                                      n_fft=MFCC_NFFT, hop_length=MFCC_HOP)
+        delta = librosa.feature.delta(mfcc)
+    return np.concatenate([
+        mfcc.mean(axis=1),  mfcc.std(axis=1),
+        delta.mean(axis=1), delta.std(axis=1),
     ]).astype(np.float32)
-    return feat
-
 
-FEATURE_NAMES: list[str] = (
-    [f"psd_{i}"      for i in range(50)] +
-    [f"mfcc_mean_{i}" for i in range(13)] +
-    [f"mfcc_std_{i}"  for i in range(13)] +
-    ["tau", "energy_ratio", "crest_factor", "log_peak_amp", "sc_mean", "sc_std"]
-)
 
+# ── Group 3: Physics features ──────────────────────────────────────────────
+
+def peak_frequency(f, pxx):
+    if pxx.max() <= 0:
+        return 0.0
+    return float(f[np.argmax(pxx)])
+
+
+def q_factor(f, pxx):
+    """Q = f_peak / -3 dB bandwidth. High Q = tight (rings cleanly)."""
+    if pxx.max() <= 0:
+        return 0.0
+    pdb      = 10 * np.log10(pxx + 1e-20)
+    peak_idx = int(np.argmax(pdb))
+    threshold = pdb[peak_idx] - 3.0
+    L = peak_idx
+    while L > 0 and pdb[L] >= threshold:
+        L -= 1
+    R = peak_idx
+    while R < len(pdb) - 1 and pdb[R] >= threshold:
+        R += 1
+    bw = max(f[R] - f[L], 1.0)
+    return float(f[peak_idx] / bw)
+
+
+def decay_tau(y, peak_sample=PEAK_SAMPLE_IN_WIN, sr=SR, fit_ms=DECAY_FIT_MS):
+    """Decay time constant τ. Loose → small τ. Tight → large τ."""
+    n_fit  = int(fit_ms * sr / 1000)
+    seg    = y[peak_sample:min(peak_sample + n_fit, len(y))]
+    if len(seg) < 100:
+        return np.nan
+    env_w = max(1, int(0.005 * sr))
+    env   = np.convolve(np.abs(seg), np.ones(env_w) / env_w, mode='same')
+    if env.max() < 1e-8:
+        return np.nan
+    active = np.where(env > 0.05 * env.max())[0]
+    if len(active) < 50:
+        return np.nan
+    n_active = active[-1] + 1
+    eps      = env.max() * 1e-4
+    log_env  = np.log(env[:n_active] + eps)
+    t        = np.arange(n_active) / sr
+    slope, _ = np.polyfit(t, log_env, 1)
+    if slope >= 0:
+        return np.nan
+    tau = -1.0 / slope
+    return float(tau) if 0.001 <= tau <= 10.0 else np.nan
+
+
+def energy_ratio(y, frac=EARLY_LATE_FRAC):
+    """E_late / E_early. Tight flanges still ringing → high ratio."""
+    n_chunk = int(frac * len(y))
+    e_early = np.sqrt(np.mean(y[:n_chunk] ** 2))
+    e_late  = np.sqrt(np.mean(y[-n_chunk:] ** 2))
+    return float(e_late / (e_early + 1e-12))
+
+
+# ── Master 82-dim extractor ───────────────────────────────────────────────
+
+def extract_features(y: np.ndarray, sr: int = SR) -> np.ndarray:
+    """Return 82-dim feature vector for one hit window."""
+    psd_bins, f_psd, pxx = relative_psd_log_bins(y, sr)    # 20
+
+    cepstral = mfcc_stats(y, sr)                            # 52
 
-# ─── Mel spectrogram (CNN input) ─────────────────────────────────────────────
-
-def extract_mel_spectrogram(window: np.ndarray) -> np.ndarray:
-    """
-    Returns mel spectrogram of shape (N_MELS, SPEC_TIME_FRAMES) in dB.
-    Used as CNN input (add channel dim in model).
-    """
-    mel = librosa.feature.melspectrogram(
-        y=window, sr=SR, n_mels=N_MELS, n_fft=N_FFT, hop_length=HOP_LENGTH_MEL
+    with warnings.catch_warnings():
+        warnings.simplefilter("ignore")
+        sc   = librosa.feature.spectral_centroid(y=y, sr=sr)[0]
+        sb   = librosa.feature.spectral_bandwidth(y=y, sr=sr)[0]
+        sr85 = librosa.feature.spectral_rolloff(y=y, sr=sr, roll_percent=0.85)[0]
+        zcr  = librosa.feature.zero_crossing_rate(y)[0]
+
+    physics = np.array([
+        sc.mean(),  sc.std(),           # 2: centroid mean/std
+        sb.mean(),                      # 1: bandwidth mean
+        sr85.mean(),                    # 1: rolloff 85%
+        zcr.mean(),                     # 1: zero-crossing rate
+        peak_frequency(f_psd, pxx),     # 1: dominant freq
+        decay_tau(y),                   # 1: τ (NaN → imputed after)
+        energy_ratio(y),                # 1: E_late / E_early
+        float(np.sqrt(np.mean(y ** 2))),# 1: RMS energy
+        q_factor(f_psd, pxx),           # 1: Q-factor
+    ], dtype=np.float32)               # 10 total
+
+    return np.concatenate([psd_bins, cepstral, physics])
+
+
+def impute_nans(X: np.ndarray, y_labels: np.ndarray, n_classes: int = 3) -> np.ndarray:
+    """Per-class median imputation for NaN columns (tau can be NaN)."""
+    X = X.copy()
+    nan_cols = np.where(np.isnan(X).any(axis=0))[0]
+    for c in nan_cols:
+        for cls in range(n_classes):
+            cls_mask   = (y_labels == cls)
+            median_val = float(np.nanmedian(X[cls_mask, c]))
+            if np.isnan(median_val):
+                median_val = float(np.nanmedian(X[:, c]))
+            fill_mask = cls_mask & np.isnan(X[:, c])
+            X[fill_mask, c] = median_val
+    return X
+
+
+# ── Feature name list ─────────────────────────────────────────────────────
+
+def _build_feature_names() -> list[str]:
+    names = []
+    edges = np.logspace(np.log10(PSD_FMIN), np.log10(PSD_FMAX), N_PSD_BINS + 1)
+    for i in range(N_PSD_BINS):
+        names.append(f'psd_{edges[i]:.0f}_{edges[i+1]:.0f}Hz')
+    names += [f'mfcc{i:02d}_mean'  for i in range(N_MFCC)]
+    names += [f'mfcc{i:02d}_std'   for i in range(N_MFCC)]
+    names += [f'dmfcc{i:02d}_mean' for i in range(N_MFCC)]
+    names += [f'dmfcc{i:02d}_std'  for i in range(N_MFCC)]
+    names += ['spec_centroid_mean', 'spec_centroid_std',
+              'spec_bandwidth_mean', 'spec_rolloff85_mean',
+              'zero_cross_rate_mean', 'peak_frequency',
+              'decay_tau', 'energy_ratio', 'rms_energy', 'q_factor']
+    return names
+
+FEATURE_NAMES: list[str] = _build_feature_names()
+assert len(FEATURE_NAMES) == 82
+
+
+# ── Mel spectrogram (CNN / BiLSTM input) — matches notebook Cell 5 ────────
+
+def extract_mel_spectrogram(y: np.ndarray, sr: int = SR,
+                             n_mels: int = N_MELS,
+                             n_fft: int = N_FFT,
+                             hop_length: int = HOP_LENGTH_MEL,
+                             target_frames: int = SPEC_TIME_FRAMES) -> np.ndarray:
+    """One hit → standardized log-mel spectrogram of shape (n_mels, target_frames)."""
+    mel    = librosa.feature.melspectrogram(
+        y=y, sr=sr, n_mels=n_mels, n_fft=n_fft, hop_length=hop_length,
+        fmin=FMIN_MEL, fmax=FMAX_MEL, power=2.0,
     )
-    mel_db = librosa.power_to_db(mel, ref=np.max)
-    # Pad / trim to fixed width
-    if mel_db.shape[1] < SPEC_TIME_FRAMES:
-        mel_db = np.pad(mel_db, ((0, 0), (0, SPEC_TIME_FRAMES - mel_db.shape[1])))
-    else:
-        mel_db = mel_db[:, :SPEC_TIME_FRAMES]
-    return mel_db.astype(np.float32)
+    mel_db = librosa.power_to_db(mel, ref=np.max).astype(np.float32)
+
+    n_frames = mel_db.shape[1]
+    if n_frames < target_frames:
+        pad_val = float(mel_db.min())
+        mel_db  = np.pad(mel_db, ((0, 0), (0, target_frames - n_frames)),
+                         mode='constant', constant_values=pad_val)
+    elif n_frames > target_frames:
+        start  = (n_frames - target_frames) // 2
+        mel_db = mel_db[:, start:start + target_frames]
+
+    # per-sample standardize
+    mu, sigma = mel_db.mean(), mel_db.std()
+    return (mel_db - mu) / (sigma + 1e-6)   # shape (64, 128)

routers/training.py

@@ -1,15 +1,18 @@
 """
 Router: POST /api/train   — launch training as a background task
         GET  /api/results  — fetch completed results
-        WS   /ws/train/{task_id} — stream live epoch metrics
+        WS   /ws/train/{task_id} — stream live epoch/fold metrics
 
-Models trained: SVM, LR, RF, MLP, KNN (shallow) + CNN, LSTM (deep via Keras).
-Each model runs LOIO cross-validation (Task 2) + 70/30 split (Task 1).
+Models (matching final_project_saurav_silwal.ipynb exactly):
+  Shallow (82-dim tabular features): SVM, LR, KNN
+  Deep (82-dim tabular):             MLP  (Keras, 3 hidden layers + dropout)
+  Deep (mel spectrogram):            CNN, BiLSTM (Keras)
+
+All models: Task 1 (70/30 dependent split) + Task 2 (LOIO cross-validation)
 """
 
 import asyncio
 import uuid
-import threading
 import traceback
 from concurrent.futures import ThreadPoolExecutor
 
@@ -19,123 +22,225 @@ from sklearn.preprocessing import StandardScaler
 from sklearn.metrics import accuracy_score, confusion_matrix, f1_score
 from sklearn.svm import SVC
 from sklearn.linear_model import LogisticRegression
-from sklearn.ensemble import RandomForestClassifier
-from sklearn.neural_network import MLPClassifier
 from sklearn.neighbors import KNeighborsClassifier
 from fastapi import APIRouter, Header, HTTPException, WebSocket, WebSocketDisconnect
 
 from session import session_manager
 from ws.training_ws import ws_manager, emit
+from ml.feature_extraction import (
+    extract_mel_spectrogram, impute_nans, FEATURE_NAMES
+)
 from config import SEED, TEST_SIZE, IDX_TO_CLASS, CLASS_NAMES, N_CLASSES
 
-router = APIRouter(tags=["training"])
-
-# Shared thread pool for background training
+router    = APIRouter(tags=["training"])
 _executor = ThreadPoolExecutor(max_workers=2)
 
 
-# ─── Model definitions ────────────────────────────────────────────────────────
+# ─────────────────────────────────────────────────────────────────────────────
+# Shallow model factories (sklearn)
+# ─────────────────────────────────────────────────────────────────────────────
 
-SHALLOW_MODELS = {
+SHALLOW_FACTORIES = {
     "SVM": lambda: SVC(
         kernel="rbf", C=10.0, gamma="scale",
-        probability=True, class_weight="balanced", random_state=SEED
+        probability=True, class_weight="balanced", random_state=SEED,
     ),
     "LR": lambda: LogisticRegression(
         C=1.0, max_iter=2000, class_weight="balanced",
-        multi_class="multinomial", solver="lbfgs", random_state=SEED
-    ),
-    "RF": lambda: RandomForestClassifier(
-        n_estimators=200, max_depth=None,
-        class_weight="balanced", random_state=SEED, n_jobs=-1
+        multi_class="multinomial", solver="lbfgs", random_state=SEED,
     ),
-    "MLP": lambda: MLPClassifier(
-        hidden_layer_sizes=(128, 64), activation="relu",
-        max_iter=500, early_stopping=True, random_state=SEED
+    "KNN": lambda: KNeighborsClassifier(
+        n_neighbors=5, metric="euclidean", weights="uniform",
     ),
-    "KNN": lambda: KNeighborsClassifier(n_neighbors=5, metric="euclidean"),
 }
 
 
-# ─── Training worker (runs in thread) ─────────────────────────────────────────
+# ─────────────────────────────────────────────────────────────────────────────
+# Keras model builders
+# ─────────────────────────────────────────────────────────────────────────────
+
+def build_mlp(input_dim: int, n_classes: int = 3):
+    """3-layer MLP with BatchNorm + Dropout. Matches notebook Cell 9."""
+    import tensorflow as tf
+    from tensorflow import keras
+
+    model = keras.Sequential([
+        keras.layers.Input(shape=(input_dim,)),
+        keras.layers.Dense(256),
+        keras.layers.BatchNormalization(),
+        keras.layers.Activation("relu"),
+        keras.layers.Dropout(0.4),
+
+        keras.layers.Dense(128),
+        keras.layers.BatchNormalization(),
+        keras.layers.Activation("relu"),
+        keras.layers.Dropout(0.3),
+
+        keras.layers.Dense(64),
+        keras.layers.BatchNormalization(),
+        keras.layers.Activation("relu"),
+        keras.layers.Dropout(0.2),
+
+        keras.layers.Dense(n_classes, activation="softmax"),
+    ])
+    model.compile(
+        optimizer=keras.optimizers.Adam(learning_rate=1e-3),
+        loss="sparse_categorical_crossentropy",
+        metrics=["accuracy"],
+    )
+    return model
+
+
+def build_cnn(n_mels: int = 64, n_frames: int = 128, n_classes: int = 3):
+    """CNN on log-mel spectrogram (64×128×1). Matches notebook Cell 10."""
+    import tensorflow as tf
+    from tensorflow import keras
+
+    model = keras.Sequential([
+        keras.layers.Input(shape=(n_mels, n_frames, 1)),
+
+        keras.layers.Conv2D(32, (3, 3), padding="same", activation="relu"),
+        keras.layers.BatchNormalization(),
+        keras.layers.MaxPooling2D((2, 2)),
+        keras.layers.Dropout(0.25),
+
+        keras.layers.Conv2D(64, (3, 3), padding="same", activation="relu"),
+        keras.layers.BatchNormalization(),
+        keras.layers.MaxPooling2D((2, 2)),
+        keras.layers.Dropout(0.25),
+
+        keras.layers.Conv2D(128, (3, 3), padding="same", activation="relu"),
+        keras.layers.BatchNormalization(),
+        keras.layers.GlobalAveragePooling2D(),
+        keras.layers.Dropout(0.4),
+
+        keras.layers.Dense(128, activation="relu"),
+        keras.layers.Dropout(0.3),
+        keras.layers.Dense(n_classes, activation="softmax"),
+    ])
+    model.compile(
+        optimizer="adam",
+        loss="sparse_categorical_crossentropy",
+        metrics=["accuracy"],
+    )
+    return model
+
+
+def build_bilstm(n_frames: int = 128, n_mels: int = 64, n_classes: int = 3):
+    """Bidirectional LSTM on mel sequences (128 time steps × 64 mel features).
+    Matches notebook Cell 11."""
+    import tensorflow as tf
+    from tensorflow import keras
+
+    model = keras.Sequential([
+        keras.layers.Input(shape=(n_frames, n_mels)),
+        keras.layers.Bidirectional(keras.layers.LSTM(64, return_sequences=True)),
+        keras.layers.Dropout(0.3),
+        keras.layers.Bidirectional(keras.layers.LSTM(32)),
+        keras.layers.Dropout(0.3),
+        keras.layers.Dense(64, activation="relu"),
+        keras.layers.Dropout(0.2),
+        keras.layers.Dense(n_classes, activation="softmax"),
+    ])
+    model.compile(
+        optimizer="adam",
+        loss="sparse_categorical_crossentropy",
+        metrics=["accuracy"],
+    )
+    return model
+
+
+# ─────────────────────────────────────────────────────────────────────────────
+# WebSocket epoch callback for Keras
+# ─────────────────────────────────────────────────────────────────────────────
+
+class WSCallback:
+    """Keras callback that emits epoch metrics over WebSocket."""
+
+    def __init__(self, loop, queue, model_name, total_epochs):
+        self.loop        = loop
+        self.queue       = queue
+        self.model_name  = model_name
+        self.total_epochs = total_epochs
+
+    def on_epoch_end(self, epoch, logs=None):
+        logs = logs or {}
+        emit(self.loop, self.queue, {
+            "type":       "epoch",
+            "model":      self.model_name,
+            "epoch":      epoch + 1,
+            "total":      self.total_epochs,
+            "train_acc":  round(float(logs.get("accuracy", 0)), 4),
+            "val_acc":    round(float(logs.get("val_accuracy", 0)), 4),
+            "train_loss": round(float(logs.get("loss", 0)), 4),
+            "val_loss":   round(float(logs.get("val_loss", 0)), 4),
+        })
+
+
+def _make_keras_callback(loop, queue, model_name, total_epochs):
+    """Return a tf.keras.callbacks.Callback subclass instance."""
+    import tensorflow as tf
 
-def _train_shallow(
-    task_id: str,
-    model_name: str,
-    X: np.ndarray,
-    y: np.ndarray,
-    groups: np.ndarray,
-    loop: asyncio.AbstractEventLoop,
-    queue: asyncio.Queue,
-    session,
-):
+    cb = WSCallback(loop, queue, model_name, total_epochs)
+
+    class _CB(tf.keras.callbacks.Callback):
+        def on_epoch_end(self, epoch, logs=None):
+            cb.on_epoch_end(epoch, logs)
+
+    return _CB()
+
+
+# ─────────────────────────────────────────────────────────────────────────────
+# Shallow model training (SVM / LR / KNN)
+# ─────────────────────────────────────────────────────────────────────────────
+
+def _train_shallow(model_name, X, y, groups, loop, queue, session):
     try:
         scaler = StandardScaler()
-        X_s = scaler.fit_transform(X)
+        X_s    = scaler.fit_transform(X)
 
-        # ── Task 1: Dependent 70/30 split ──
+        # Task 1
         X_tr, X_te, y_tr, y_te = train_test_split(
             X_s, y, test_size=TEST_SIZE, stratify=y, random_state=SEED
         )
-        clf_t1 = SHALLOW_MODELS[model_name]()
-        clf_t1.fit(X_tr, y_tr)
-        y_pred_t1 = clf_t1.predict(X_te)
-        acc_t1  = float(accuracy_score(y_te, y_pred_t1))
-        f1_t1   = float(f1_score(y_te, y_pred_t1, average="macro"))
-        cm_t1   = confusion_matrix(y_te, y_pred_t1, labels=[0, 1, 2]).tolist()
-
-        emit(loop, queue, {
-            "type": "task1_done", "model": model_name,
-            "acc": round(acc_t1, 4), "f1": round(f1_t1, 4),
-        })
-
-        # ── Task 2: LOIO cross-validation ──
-        logo = LeaveOneGroupOut()
-        fold_accs: list[float] = []
-        fold_records: list[dict] = []
+        clf = SHALLOW_FACTORIES[model_name]()
+        clf.fit(X_tr, y_tr)
+        y_p1   = clf.predict(X_te)
+        acc_t1 = float(accuracy_score(y_te, y_p1))
+        f1_t1  = float(f1_score(y_te, y_p1, average="macro"))
+        cm_t1  = confusion_matrix(y_te, y_p1, labels=[0, 1, 2]).tolist()
+
+        emit(loop, queue, {"type": "task1_done", "model": model_name,
+                           "acc": round(acc_t1, 4), "f1": round(f1_t1, 4)})
+
+        # Task 2 — LOIO
+        logo         = LeaveOneGroupOut()
+        fold_accs    = []
+        fold_records = []
+        all_yt, all_yp = [], []
 
         for fold_i, (tr_idx, te_idx) in enumerate(logo.split(X_s, y, groups)):
             flange_out = int(groups[te_idx[0]])
-            clf = SHALLOW_MODELS[model_name]()
-            clf.fit(X_s[tr_idx], y[tr_idx])
-            y_p   = clf.predict(X_s[te_idx])
-            y_pr  = clf.predict_proba(X_s[te_idx]) if hasattr(clf, "predict_proba") else None
-            acc_f = float(accuracy_score(y[te_idx], y_p))
+            clf2       = SHALLOW_FACTORIES[model_name]()
+            clf2.fit(X_s[tr_idx], y[tr_idx])
+            yp   = clf2.predict(X_s[te_idx])
+            acc_f = float(accuracy_score(y[te_idx], yp))
             fold_accs.append(acc_f)
-            fold_records.append({
-                "fold":       fold_i + 1,
-                "flange_out": flange_out,
-                "acc":        round(acc_f, 4),
-                "n_test":     len(te_idx),
-            })
-            emit(loop, queue, {
-                "type":       "fold_done",
-                "model":      model_name,
-                "fold":       fold_i + 1,
-                "flange_out": flange_out,
-                "acc":        round(acc_f, 4),
-            })
-
-        # Final model on all data (for ensemble / CORAL)
-        clf_final = SHALLOW_MODELS[model_name]()
+            fold_records.append({"fold": fold_i + 1, "flange_out": flange_out,
+                                  "acc": round(acc_f, 4), "n_test": len(te_idx)})
+            all_yt.extend(y[te_idx].tolist())
+            all_yp.extend(yp.tolist())
+            emit(loop, queue, {"type": "fold_done", "model": model_name,
+                                "fold": fold_i + 1, "flange_out": flange_out,
+                                "acc": round(acc_f, 4)})
+
+        cm_t2  = confusion_matrix(all_yt, all_yp, labels=[0, 1, 2]).tolist()
+        f1_t2  = float(f1_score(all_yt, all_yp, average="macro"))
+
+        # Final model on all data
+        clf_final = SHALLOW_FACTORIES[model_name]()
         clf_final.fit(X_s, y)
-        y_pred_all  = clf_final.predict(X_s)
-        train_acc   = float(accuracy_score(y, y_pred_all))
-        cm_loio_pooled = confusion_matrix(
-            [f["flange_out"] for f in fold_records],  # dummy — use actual pooled
-            [f["flange_out"] for f in fold_records],
-        ).tolist()
-
-        # Pooled LOIO confusion matrix
-        all_y_true: list[int] = []
-        all_y_pred: list[int] = []
-        for tr_idx, te_idx in logo.split(X_s, y, groups):
-            clf = SHALLOW_MODELS[model_name]()
-            clf.fit(X_s[tr_idx], y[tr_idx])
-            all_y_true.extend(y[te_idx].tolist())
-            all_y_pred.extend(clf.predict(X_s[te_idx]).tolist())
-        cm_t2 = confusion_matrix(all_y_true, all_y_pred, labels=[0, 1, 2]).tolist()
-        f1_t2 = float(f1_score(all_y_true, all_y_pred, average="macro"))
+        train_acc = float(accuracy_score(y, clf_final.predict(X_s)))
 
         result = {
             "model":        model_name,
@@ -143,17 +248,90 @@ def _train_shallow(
             "task1_f1":     round(f1_t1, 4),
             "task1_cm":     cm_t1,
             "task2_mean":   round(float(np.mean(fold_accs)), 4),
-            "task2_std":    round(float(np.std(fold_accs)),  4),
+            "task2_std":    round(float(np.std(fold_accs)), 4),
             "task2_f1":     round(f1_t2, 4),
             "task2_cm":     cm_t2,
             "folds":        fold_records,
             "train_acc":    round(train_acc, 4),
-            "scaler_mean":  scaler.mean_.tolist(),
-            "scaler_scale": scaler.scale_.tolist(),
         }
         session.training_results[model_name] = result
         session.touch()
+        emit(loop, queue, {"type": "model_done", "model": model_name, **result})
+
+    except Exception as e:
+        emit(loop, queue, {"type": "error", "model": model_name, "message": str(e)})
+        traceback.print_exc()
 
+
+# ─────────────────────────────────────────────────────────────────────────────
+# MLP training (Keras, tabular features)
+# ─────────────────────────────────────────────────────────────────────────────
+
+def _train_mlp(X, y, groups, loop, queue, session, epochs=60):
+    model_name = "MLP"
+    try:
+        import tensorflow as tf
+        tf.random.set_seed(SEED)
+
+        scaler = StandardScaler()
+        X_s    = scaler.fit_transform(X)
+
+        # Task 1
+        X_tr, X_te, y_tr, y_te = train_test_split(
+            X_s, y, test_size=TEST_SIZE, stratify=y, random_state=SEED
+        )
+        model  = build_mlp(X_s.shape[1])
+        cb     = _make_keras_callback(loop, queue, model_name, epochs)
+        es     = tf.keras.callbacks.EarlyStopping(patience=10, restore_best_weights=True)
+        model.fit(X_tr, y_tr, epochs=epochs, batch_size=32,
+                  validation_split=0.15,
+                  callbacks=[cb, es], verbose=0)
+
+        y_p1   = np.argmax(model.predict(X_te, verbose=0), axis=1)
+        acc_t1 = float(accuracy_score(y_te, y_p1))
+        f1_t1  = float(f1_score(y_te, y_p1, average="macro"))
+        cm_t1  = confusion_matrix(y_te, y_p1, labels=[0, 1, 2]).tolist()
+
+        emit(loop, queue, {"type": "task1_done", "model": model_name,
+                           "acc": round(acc_t1, 4), "f1": round(f1_t1, 4)})
+
+        # Task 2 — LOIO
+        logo         = LeaveOneGroupOut()
+        fold_accs    = []
+        fold_records = []
+        all_yt, all_yp = [], []
+
+        for fold_i, (tr_idx, te_idx) in enumerate(logo.split(X_s, y, groups)):
+            flange_out = int(groups[te_idx[0]])
+            m2 = build_mlp(X_s.shape[1])
+            es2 = tf.keras.callbacks.EarlyStopping(patience=8, restore_best_weights=True)
+            m2.fit(X_s[tr_idx], y[tr_idx], epochs=epochs, batch_size=32,
+                   validation_split=0.15, callbacks=[es2], verbose=0)
+            yp    = np.argmax(m2.predict(X_s[te_idx], verbose=0), axis=1)
+            acc_f = float(accuracy_score(y[te_idx], yp))
+            fold_accs.append(acc_f)
+            fold_records.append({"fold": fold_i + 1, "flange_out": flange_out,
+                                  "acc": round(acc_f, 4), "n_test": len(te_idx)})
+            all_yt.extend(y[te_idx].tolist())
+            all_yp.extend(yp.tolist())
+            emit(loop, queue, {"type": "fold_done", "model": model_name,
+                                "fold": fold_i + 1, "flange_out": flange_out,
+                                "acc": round(acc_f, 4)})
+
+        cm_t2 = confusion_matrix(all_yt, all_yp, labels=[0, 1, 2]).tolist()
+        f1_t2 = float(f1_score(all_yt, all_yp, average="macro"))
+        train_acc = float(accuracy_score(y, np.argmax(model.predict(X_s, verbose=0), axis=1)))
+
+        result = {
+            "model": model_name, "task1_acc": round(acc_t1, 4),
+            "task1_f1": round(f1_t1, 4), "task1_cm": cm_t1,
+            "task2_mean": round(float(np.mean(fold_accs)), 4),
+            "task2_std":  round(float(np.std(fold_accs)), 4),
+            "task2_f1":   round(f1_t2, 4), "task2_cm": cm_t2,
+            "folds": fold_records, "train_acc": round(train_acc, 4),
+        }
+        session.training_results[model_name] = result
+        session.touch()
         emit(loop, queue, {"type": "model_done", "model": model_name, **result})
 
     except Exception as e:
@@ -161,15 +339,167 @@ def _train_shallow(
         traceback.print_exc()
 
 
-def _train_all_models(task_id: str, session_id: str, models: list[str]):
-    """Entry point for background thread: trains all requested models sequentially."""
+# ─────────────────────────────────────────────────────────────────────────────
+# CNN training (Keras, mel spectrograms)
+# ─────────────────────────────────────────────────────────────────────────────
+
+def _train_cnn(waveforms, y, groups, loop, queue, session, epochs=50):
+    model_name = "CNN"
+    try:
+        import tensorflow as tf
+        tf.random.set_seed(SEED)
+
+        # Build spectrogram tensor (N, 64, 128, 1)
+        emit(loop, queue, {"type": "task1_done", "model": model_name,
+                           "acc": 0.0, "f1": 0.0, "message": "Extracting spectrograms..."})
+        X_spec = np.stack([
+            extract_mel_spectrogram(np.array(w, dtype=np.float32))
+            for w in waveforms
+        ], axis=0)[..., np.newaxis]   # (N, 64, 128, 1)
+
+        # Task 1
+        X_tr, X_te, y_tr, y_te = train_test_split(
+            X_spec, y, test_size=TEST_SIZE, stratify=y, random_state=SEED
+        )
+        model = build_cnn()
+        cb    = _make_keras_callback(loop, queue, model_name, epochs)
+        es    = tf.keras.callbacks.EarlyStopping(patience=10, restore_best_weights=True)
+        model.fit(X_tr, y_tr, epochs=epochs, batch_size=32,
+                  validation_split=0.15, callbacks=[cb, es], verbose=0)
+
+        y_p1   = np.argmax(model.predict(X_te, verbose=0), axis=1)
+        acc_t1 = float(accuracy_score(y_te, y_p1))
+        f1_t1  = float(f1_score(y_te, y_p1, average="macro"))
+        cm_t1  = confusion_matrix(y_te, y_p1, labels=[0, 1, 2]).tolist()
+
+        # Task 2 — LOIO
+        logo = LeaveOneGroupOut()
+        fold_accs, fold_records, all_yt, all_yp = [], [], [], []
+
+        for fold_i, (tr_idx, te_idx) in enumerate(logo.split(X_spec, y, groups)):
+            flange_out = int(groups[te_idx[0]])
+            m2 = build_cnn()
+            es2 = tf.keras.callbacks.EarlyStopping(patience=8, restore_best_weights=True)
+            m2.fit(X_spec[tr_idx], y[tr_idx], epochs=epochs, batch_size=32,
+                   validation_split=0.15, callbacks=[es2], verbose=0)
+            yp    = np.argmax(m2.predict(X_spec[te_idx], verbose=0), axis=1)
+            acc_f = float(accuracy_score(y[te_idx], yp))
+            fold_accs.append(acc_f)
+            fold_records.append({"fold": fold_i + 1, "flange_out": flange_out,
+                                  "acc": round(acc_f, 4), "n_test": len(te_idx)})
+            all_yt.extend(y[te_idx].tolist())
+            all_yp.extend(yp.tolist())
+            emit(loop, queue, {"type": "fold_done", "model": model_name,
+                                "fold": fold_i + 1, "flange_out": flange_out,
+                                "acc": round(acc_f, 4)})
+
+        cm_t2     = confusion_matrix(all_yt, all_yp, labels=[0, 1, 2]).tolist()
+        f1_t2     = float(f1_score(all_yt, all_yp, average="macro"))
+        train_acc = float(accuracy_score(y, np.argmax(model.predict(X_spec, verbose=0), axis=1)))
+
+        result = {
+            "model": model_name, "task1_acc": round(acc_t1, 4),
+            "task1_f1": round(f1_t1, 4), "task1_cm": cm_t1,
+            "task2_mean": round(float(np.mean(fold_accs)), 4),
+            "task2_std":  round(float(np.std(fold_accs)), 4),
+            "task2_f1":   round(f1_t2, 4), "task2_cm": cm_t2,
+            "folds": fold_records, "train_acc": round(train_acc, 4),
+        }
+        session.training_results[model_name] = result
+        session.touch()
+        emit(loop, queue, {"type": "model_done", "model": model_name, **result})
+
+    except Exception as e:
+        emit(loop, queue, {"type": "error", "model": model_name, "message": str(e)})
+        traceback.print_exc()
+
+
+# ─────────────────────────────────────────────────────────────────────────────
+# BiLSTM training (Keras, mel sequences)
+# ─────────────────────────────────────────────────────────────────────────────
+
+def _train_bilstm(waveforms, y, groups, loop, queue, session, epochs=50):
+    model_name = "BiLSTM"
+    try:
+        import tensorflow as tf
+        tf.random.set_seed(SEED)
+
+        # Reshape to (N, 128, 64) — time steps × mel features
+        X_seq = np.stack([
+            extract_mel_spectrogram(np.array(w, dtype=np.float32)).T   # (128, 64)
+            for w in waveforms
+        ], axis=0)
+
+        # Task 1
+        X_tr, X_te, y_tr, y_te = train_test_split(
+            X_seq, y, test_size=TEST_SIZE, stratify=y, random_state=SEED
+        )
+        model = build_bilstm()
+        cb    = _make_keras_callback(loop, queue, model_name, epochs)
+        es    = tf.keras.callbacks.EarlyStopping(patience=10, restore_best_weights=True)
+        model.fit(X_tr, y_tr, epochs=epochs, batch_size=32,
+                  validation_split=0.15, callbacks=[cb, es], verbose=0)
+
+        y_p1   = np.argmax(model.predict(X_te, verbose=0), axis=1)
+        acc_t1 = float(accuracy_score(y_te, y_p1))
+        f1_t1  = float(f1_score(y_te, y_p1, average="macro"))
+        cm_t1  = confusion_matrix(y_te, y_p1, labels=[0, 1, 2]).tolist()
+
+        # Task 2 — LOIO
+        logo = LeaveOneGroupOut()
+        fold_accs, fold_records, all_yt, all_yp = [], [], [], []
+
+        for fold_i, (tr_idx, te_idx) in enumerate(logo.split(X_seq, y, groups)):
+            flange_out = int(groups[te_idx[0]])
+            m2 = build_bilstm()
+            es2 = tf.keras.callbacks.EarlyStopping(patience=8, restore_best_weights=True)
+            m2.fit(X_seq[tr_idx], y[tr_idx], epochs=epochs, batch_size=32,
+                   validation_split=0.15, callbacks=[es2], verbose=0)
+            yp    = np.argmax(m2.predict(X_seq[te_idx], verbose=0), axis=1)
+            acc_f = float(accuracy_score(y[te_idx], yp))
+            fold_accs.append(acc_f)
+            fold_records.append({"fold": fold_i + 1, "flange_out": flange_out,
+                                  "acc": round(acc_f, 4), "n_test": len(te_idx)})
+            all_yt.extend(y[te_idx].tolist())
+            all_yp.extend(yp.tolist())
+            emit(loop, queue, {"type": "fold_done", "model": model_name,
+                                "fold": fold_i + 1, "flange_out": flange_out,
+                                "acc": round(acc_f, 4)})
+
+        cm_t2     = confusion_matrix(all_yt, all_yp, labels=[0, 1, 2]).tolist()
+        f1_t2     = float(f1_score(all_yt, all_yp, average="macro"))
+        train_acc = float(accuracy_score(y, np.argmax(model.predict(X_seq, verbose=0), axis=1)))
+
+        result = {
+            "model": model_name, "task1_acc": round(acc_t1, 4),
+            "task1_f1": round(f1_t1, 4), "task1_cm": cm_t1,
+            "task2_mean": round(float(np.mean(fold_accs)), 4),
+            "task2_std":  round(float(np.std(fold_accs)), 4),
+            "task2_f1":   round(f1_t2, 4), "task2_cm": cm_t2,
+            "folds": fold_records, "train_acc": round(train_acc, 4),
+        }
+        session.training_results[model_name] = result
+        session.touch()
+        emit(loop, queue, {"type": "model_done", "model": model_name, **result})
+
+    except Exception as e:
+        emit(loop, queue, {"type": "error", "model": model_name, "message": str(e)})
+        traceback.print_exc()
+
+
+# ─────────────────────────────────────────────────────────────────────────────
+# Master training thread
+# ─────────────────────────────────────────────────────────────────────────────
+
+ALL_MODELS = ["SVM", "LR", "KNN", "MLP", "CNN", "BiLSTM"]
+
+def _train_all(task_id: str, session_id: str, models: list[str]):
     loop = asyncio.new_event_loop()
     asyncio.set_event_loop(loop)
 
     session = session_manager.get(session_id)
     if session is None:
         return
-
     queue = ws_manager.get_queue(task_id)
     if queue is None:
         return
@@ -179,57 +509,57 @@ def _train_all_models(task_id: str, session_id: str, models: list[str]):
         emit(loop, queue, {"type": "error", "message": "Features not extracted yet"})
         return
 
-    X      = np.array(feats["X_feat"],        dtype=np.float32)
-    y      = np.array(feats["labels"],         dtype=np.int64)
-    groups = np.array(feats["flange_groups"],  dtype=np.int64)
+    X      = np.array(feats["X_feat"],       dtype=np.float32)
+    y      = np.array(feats["labels"],        dtype=np.int64)
+    groups = np.array(feats["flange_groups"], dtype=np.int64)
 
-    for model_name in models:
-        if model_name in SHALLOW_MODELS:
-            _train_shallow(task_id, model_name, X, y, groups, loop, queue, session)
+    # NaN imputation (tau column can have NaNs)
+    X = impute_nans(X, y)
+
+    waveforms = session.hits.get("waveforms", [])
+
+    for m in models:
+        if m in SHALLOW_FACTORIES:
+            _train_shallow(m, X, y, groups, loop, queue, session)
+        elif m == "MLP":
+            _train_mlp(X, y, groups, loop, queue, session)
+        elif m == "CNN":
+            _train_cnn(waveforms, y, groups, loop, queue, session)
+        elif m == "BiLSTM":
+            _train_bilstm(waveforms, y, groups, loop, queue, session)
 
     emit(loop, queue, {"type": "all_done", "task_id": task_id})
 
 
-# ─── Routes ──────────────────────────────────────────────────────────────────
+# ─────────────────────────────────────────────────────────────────────────────
+# Routes
+# ─────────────────────────────────────────────────────────────────────────────
 
 @router.post("/api/train")
 async def start_training(
     session_id: str = Header(..., alias="X-Session-Id"),
     body: dict = None,
 ):
-    """
-    Launch background training. Returns task_id for WebSocket connection.
-    Body: {"models": ["SVM", "LR", "RF", "MLP", "KNN"]}
-    """
     session = session_manager.get(session_id)
     if session is None:
         raise HTTPException(status_code=404, detail="Session not found")
     if not session.features:
         raise HTTPException(status_code=400, detail="Extract features first: POST /api/features")
 
-    models = (body or {}).get("models", list(SHALLOW_MODELS.keys()))
-
+    models  = (body or {}).get("models", ALL_MODELS)
     task_id = str(uuid.uuid4())
-    # Create queue before starting thread (thread will use it immediately)
     ws_manager.create_queue(task_id)
     session.training_tasks[task_id] = models
     session.touch()
 
     loop = asyncio.get_event_loop()
-    loop.run_in_executor(
-        _executor,
-        _train_all_models,
-        task_id,
-        session_id,
-        models,
-    )
+    loop.run_in_executor(_executor, _train_all, task_id, session_id, models)
 
     return {"task_id": task_id, "models": models}
 
 
 @router.get("/api/results")
 async def get_results(session_id: str = Header(..., alias="X-Session-Id")):
-    """Return all completed training results for this session."""
     session = session_manager.get(session_id)
     if session is None:
         raise HTTPException(status_code=404, detail="Session not found")
@@ -239,11 +569,8 @@ async def get_results(session_id: str = Header(..., alias="X-Session-Id")):
     }
 
 
-# ─── WebSocket endpoint ───────────────────────────────────────────────────────
-
 @router.websocket("/ws/train/{task_id}")
 async def training_websocket(websocket: WebSocket, task_id: str):
-    """Stream live training events for a given task_id."""
     await ws_manager.connect(task_id, websocket)
     try:
         await ws_manager.stream(task_id, websocket)