Fixing a lethal bug

app.py

@@ -9,7 +9,7 @@ from fastapi.middleware.cors import CORSMiddleware
 from tensorflow.keras import layers
 import keras
 
-# ── Custom SpecAugment layer (must match your training definition exactly) ─────
+# ── Custom SpecAugment layer ───────────────────────────────────────────────────
 @keras.saving.register_keras_serializable()
 class SpecAugment(layers.Layer):
     def __init__(self, time_mask=20, freq_mask=15, **kwargs):
@@ -46,17 +46,17 @@ class SpecAugment(layers.Layer):
         cfg = super().get_config()
         cfg.update({"time_mask": self.time_mask, "freq_mask": self.freq_mask})
         return cfg
-        
-# ── Constants (from model_config.json) ────────────────────────────────────────
+
+# ── Constants ──────────────────────────────────────────────────────────────────
 SR          = 22050
 DURATION    = 3
 HOP_LENGTH  = 512
 TIME_FRAMES = 130
 N_BINS      = 84
-N_MFCC      = 120
+N_MFCC      = 40        # 40 × 3 (+ delta + delta2) = 120 features
 CLASSES     = ["Muwa", "east", "loyal", "poems", "rai"]
 
-# ── App ───────────────────────────────────────────────────────────────────────
+# ── App ────────────────────────────────────────────────────────────────────────
 app = FastAPI(title="Audio Classifier API")
 
 app.add_middleware(
@@ -71,8 +71,12 @@ app.add_middleware(
 async def load_artifacts():
     global model, scaler_cqt, scaler_mfcc, label_encoder
 
-    model = tf.keras.models.load_model("model_artifacts/best_cross_attn.keras", custom_objects={"SpecAugment": SpecAugment}, compile=False)
-    print(" Model loaded")
+    model = tf.keras.models.load_model(
+        "model_artifacts/best_cross_attn.keras",
+        custom_objects={"SpecAugment": SpecAugment},
+        compile=False
+    )
+    print("Model loaded")
 
     with open("model_artifacts/scaler_cqt.pkl", "rb") as f:
         scaler_cqt = pickle.load(f)
@@ -85,33 +89,42 @@ async def load_artifacts():
 
     print(" Scalers and label encoder loaded")
 
-# ── Feature extraction ─────────────────────────────────────────────────────────
-def extract_features(audio_path: str):
-    y, sr = librosa.load(audio_path, sr=SR, duration=DURATION)
-
-    # Pad if audio is shorter than expected
-    expected_samples = SR * DURATION
-    if len(y) < expected_samples:
-        y = np.pad(y, (0, expected_samples - len(y)))
-
-    # CQT → (TIME_FRAMES, N_BINS)
-    cqt = np.abs(librosa.cqt(y, sr=sr, hop_length=HOP_LENGTH, n_bins=N_BINS))
-    cqt = librosa.util.fix_length(cqt, size=TIME_FRAMES, axis=1).T  # (130, 84)
-
-    # Apply scaler (fit was on 2D, so flatten → scale → reshape)
-    cqt_shape = cqt.shape
-    cqt = scaler_cqt.transform(cqt.reshape(-1, cqt_shape[-1])).reshape(cqt_shape)
+# ── Feature extraction (mirrors training pipeline exactly) ─────────────────────
+def extract_features_from_array(y):
+    # Pad or trim to fixed duration
+    expected = SR * DURATION
+    if len(y) < expected:
+        y = np.pad(y, (0, expected - len(y)))
+    else:
+        y = y[:expected]
+
+    # CQT → amplitude_to_db → (T, 84)
+    cqt    = librosa.cqt(y=y, sr=SR, hop_length=HOP_LENGTH,
+                         n_bins=N_BINS, bins_per_octave=12)
+    cqt    = librosa.amplitude_to_db(np.abs(cqt)).T          # (T, 84)
+
+    # MFCC + delta + delta2 → (T, 120)
+    mfcc   = librosa.feature.mfcc(y=y, sr=SR, n_mfcc=N_MFCC, hop_length=HOP_LENGTH)
+    delta  = librosa.feature.delta(mfcc)
+    delta2 = librosa.feature.delta(mfcc, order=2)
+    mfcc   = np.vstack([mfcc, delta, delta2]).T               # (T, 120)
+
+    # Align & pad/trim to TIME_FRAMES
+    min_t = min(cqt.shape[0], mfcc.shape[0])
+    cqt, mfcc = cqt[:min_t], mfcc[:min_t]
+    if min_t < TIME_FRAMES:
+        cqt  = np.pad(cqt,  ((0, TIME_FRAMES - min_t), (0, 0)))
+        mfcc = np.pad(mfcc, ((0, TIME_FRAMES - min_t), (0, 0)))
+    else:
+        cqt  = cqt[:TIME_FRAMES]
+        mfcc = mfcc[:TIME_FRAMES]
+
+    return cqt, mfcc    # (130, 84), (130, 120)
 
-    # MFCC → (TIME_FRAMES, N_MFCC)
-    mfcc = librosa.feature.mfcc(y=y, sr=sr, n_mfcc=N_MFCC, hop_length=HOP_LENGTH)
-    mfcc = librosa.util.fix_length(mfcc, size=TIME_FRAMES, axis=1).T  # (130, 40)
 
-    # Apply scaler
-    mfcc_shape = mfcc.shape
-    mfcc = scaler_mfcc.transform(mfcc.reshape(-1, mfcc_shape[-1])).reshape(mfcc_shape)
-
-    # Add batch dimension → (1, 130, 84) and (1, 130, 40)
-    return cqt[np.newaxis], mfcc[np.newaxis]
+def extract_features(audio_path: str):
+    y, _ = librosa.load(audio_path, sr=SR, duration=DURATION)
+    return extract_features_from_array(y)
 
 # ── Predict endpoint ───────────────────────────────────────────────────────────
 @app.post("/predict")
@@ -125,10 +138,17 @@ async def predict(file: UploadFile = File(...)):
 
     try:
         cqt, mfcc = extract_features(tmp_path)
-        preds     = model.predict([cqt, mfcc])[0]          # (5,)
-        top_idx   = int(np.argmax(preds))
 
-        # Use label_encoder if sklearn LabelEncoder, else fall back to CLASSES list
+        # Apply scalers exactly as in training
+        cqt_shape  = cqt.shape
+        mfcc_shape = mfcc.shape
+        cqt  = scaler_cqt.transform(cqt.reshape(-1, 84)).reshape(cqt_shape)
+        mfcc = scaler_mfcc.transform(mfcc.reshape(-1, 120)).reshape(mfcc_shape)
+
+        # Add batch dim
+        preds   = model.predict([cqt[np.newaxis], mfcc[np.newaxis]])[0]
+        top_idx = int(np.argmax(preds))
+
         try:
             label = label_encoder.inverse_transform([top_idx])[0]
         except Exception: