Fixing a bug

app.py

@@ -6,7 +6,52 @@ import librosa
 import tensorflow as tf
 from fastapi import FastAPI, UploadFile, File, HTTPException
 from fastapi.middleware.cors import CORSMiddleware
-
+from tensorflow.keras import layers
+import keras
+
+# ── Custom SpecAugment layer (must match your training definition exactly) ─────
+@keras.saving.register_keras_serializable()
+class SpecAugment(layers.Layer):
+    def __init__(self, time_mask=20, freq_mask=10, **kwargs):
+        super().__init__(**kwargs)
+        self.time_mask = time_mask
+        self.freq_mask = freq_mask
+
+    def call(self, x, training=None):
+        if not training:
+            return x
+
+        shape = tf.shape(x)
+        time_steps = shape[1]
+        freq_steps = shape[2]
+
+        # Time masking
+        t = tf.random.uniform((), 0, self.time_mask, dtype=tf.int32)
+        t0 = tf.random.uniform((), 0, tf.maximum(1, time_steps - t), dtype=tf.int32)
+        time_mask_tensor = tf.concat([
+            tf.ones([shape[0], t0, freq_steps]),
+            tf.zeros([shape[0], t, freq_steps]),
+            tf.ones([shape[0], time_steps - t0 - t, freq_steps])
+        ], axis=1)
+        x = x * time_mask_tensor
+
+        # Frequency masking
+        f = tf.random.uniform((), 0, self.freq_mask, dtype=tf.int32)
+        f0 = tf.random.uniform((), 0, tf.maximum(1, freq_steps - f), dtype=tf.int32)
+        freq_mask_tensor = tf.concat([
+            tf.ones([shape[0], time_steps, f0]),
+            tf.zeros([shape[0], time_steps, f]),
+            tf.ones([shape[0], time_steps, freq_steps - f0 - f])
+        ], axis=2)
+        x = x * freq_mask_tensor
+
+        return x
+
+    def get_config(self):
+        config = super().get_config()
+        config.update({"time_mask": self.time_mask, "freq_mask": self.freq_mask})
+        return config
+        
 # ── Constants (from model_config.json) ────────────────────────────────────────
 SR          = 22050
 DURATION    = 3
@@ -31,8 +76,8 @@ 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", 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)
@@ -43,7 +88,7 @@ async def load_artifacts():
     with open("model_artifacts/label_encoder.pkl", "rb") as f:
         label_encoder = pickle.load(f)
 
-    print("✅ Scalers and label encoder loaded")
+    print(" Scalers and label encoder loaded")
 
 # ── Feature extraction ─────────────────────────────────────────────────────────
 def extract_features(audio_path: str):