Update app/model.py

app/model.py

@@ -1,53 +1,61 @@
-import torch
-from transformers import Wav2Vec2FeatureExtractor, Wav2Vec2Model
+# app/model.py
+import joblib
 import numpy as np
+import librosa
+import os
 
-from app.features import extract_features
-from app.classifier import simple_rule_classifier
-
-device = "cpu"
-
+# Global placeholders
 model = None
-feature_extractor = None
-
-def load_models():
-    global model, feature_extractor
-
-    if model is None:
-        feature_extractor = Wav2Vec2FeatureExtractor.from_pretrained(
-            "superb/wav2vec2-base-superb-er"
-        )
-
-        model = Wav2Vec2Model.from_pretrained(
-            "superb/wav2vec2-base-superb-er"
-        ).to(device)
-
-
-def predict(audio):
-    # ---- Tone features ----
-    tone_features = extract_features(audio)
-
-    # ---- Deep embeddings ----
-    inputs = feature_extractor(
-        audio,
-        sampling_rate=16000,
-        return_tensors="pt",
-        padding=True
-    )
-
-    with torch.no_grad():
-        outputs = model(**inputs)
-
-    embeddings = outputs.last_hidden_state.mean(dim=1).numpy()[0]
-
-    # ---- Combine ----
-    combined = np.hstack([tone_features, embeddings])
-
-    # ---- Classify ----
-    emotion, confidence = simple_rule_classifier(tone_features)
+scaler = None
+
+# ==========================
+# Load Model and Scaler
+# ==========================
+def load_models(model_path="emotion_model.pkl", scaler_path="scaler.pkl"):
+    global model, scaler
+    if not os.path.exists(model_path) or not os.path.exists(scaler_path):
+        raise FileNotFoundError("Model or scaler .pkl files not found. Upload them to the app directory.")
+    model = joblib.load(model_path)
+    scaler = joblib.load(scaler_path)
+    print("✅ Tone-based emotion model loaded successfully.")
+
+# ==========================
+# Feature extraction
+# ==========================
+def extract_features(audio_path):
+    """
+    Extract tone-based features from audio:
+    - MFCC
+    - Pitch
+    - Energy
+    """
+    audio, sr = librosa.load(audio_path, sr=16000)
+    mfcc = librosa.feature.mfcc(y=audio, sr=sr, n_mfcc=13)
+    mfcc_mean = np.mean(mfcc, axis=1)
+    pitches, magnitudes = librosa.piptrack(y=audio, sr=sr)
+    pitch = np.mean(pitches[pitches > 0]) if np.any(pitches > 0) else 0
+    energy = np.mean(librosa.feature.rms(y=audio))
+    return np.hstack([mfcc_mean, pitch, energy])
+
+# ==========================
+# Predict Emotion
+# ==========================
+def predict_tone(audio_path):
+    global model, scaler
+    if model is None or scaler is None:
+        raise RuntimeError("Model and scaler must be loaded first.")
+
+    features = extract_features(audio_path).reshape(1, -1)
+    features_scaled = scaler.transform(features)
+    pred_label = model.predict(features_scaled)[0]
+    pred_proba = model.predict_proba(features_scaled)[0]
+
+    # Confidence of the predicted class
+    pred_index = np.argmax(pred_proba)
+    confidence = float(pred_proba[pred_index])
 
     return {
-        "emotion_label": emotion,
+        "emotion_label": pred_label,
         "confidence": confidence,
         "note": "Tone-based prediction (less text bias)"
     }