Update server.py

server.py

@@ -1,13 +1,22 @@
-import os, io, subprocess, tempfile
+import os
+import io
+import tempfile
+import subprocess
+import numpy as np
+import torch
+import librosa
+import soundfile as sf
 from fastapi import FastAPI, UploadFile, File
 from fastapi.responses import JSONResponse
 from fastapi.middleware.cors import CORSMiddleware
 from transformers import Wav2Vec2ForSequenceClassification, AutoFeatureExtractor
-import librosa
-import torch, numpy as np, soundfile as sf
 
+# ---------------------------
+# FastAPI setup
+# ---------------------------
 app = FastAPI()
 
+# CORS config
 origins = os.environ.get('CORS_ORIGINS', '*').split(',') if os.environ.get('CORS_ORIGINS') else ['*']
 app.add_middleware(
     CORSMiddleware,
@@ -17,22 +26,26 @@ app.add_middleware(
     allow_headers=["*"],
 )
 
+# ---------------------------
+# Model setup
+# ---------------------------
 FFMPEG_BIN = os.environ.get('FFMPEG_BIN', 'ffmpeg')
+MODEL_REPO = "marshal-yash/SER_wav2vec"
 
-# ✅ FIX: MODEL IS IN THE SAME FOLDER AS server.py
-MODEL_DIR = os.path.dirname(__file__)
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
-# Load model + feature extractor
-model = Wav2Vec2ForSequenceClassification.from_pretrained(MODEL_DIR)
-fe = AutoFeatureExtractor.from_pretrained(MODEL_DIR)
-
-device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+# Load model and feature extractor from Hugging Face
+model = Wav2Vec2ForSequenceClassification.from_pretrained(MODEL_REPO)
+fe = AutoFeatureExtractor.from_pretrained(MODEL_REPO)
 model.to(device)
 model.eval()
 
-
+# ---------------------------
+# Utility: Convert audio to 16kHz mono
+# ---------------------------
 def to_wav16k_mono(data: bytes) -> np.ndarray:
     try:
+        # Use ffmpeg if available
         p = subprocess.run(
             [FFMPEG_BIN, '-hide_banner', '-loglevel', 'error',
              '-i', 'pipe:0', '-ar', str(fe.sampling_rate), '-ac', '1',
@@ -40,60 +53,49 @@ def to_wav16k_mono(data: bytes) -> np.ndarray:
             input=data, stdout=subprocess.PIPE, stderr=subprocess.PIPE, check=True
         )
         audio, sr = sf.read(io.BytesIO(p.stdout), dtype='float32', always_2d=False)
-        if isinstance(audio, np.ndarray):
-            out = audio.astype(np.float32)
-            if sr != fe.sampling_rate:
-                out = librosa.resample(out, orig_sr=sr, target_sr=fe.sampling_rate)
-            if out.size < fe.sampling_rate // 10:
-                out = np.pad(out, (0, max(0, fe.sampling_rate - out.size)), mode='constant')
-            return out
-        return np.array(audio, dtype=np.float32)
-
+        if sr != fe.sampling_rate:
+            audio = librosa.resample(audio, orig_sr=sr, target_sr=fe.sampling_rate)
+        return audio.astype(np.float32)
     except Exception:
+        # fallback: try reading directly with soundfile / librosa
         try:
             audio, sr = sf.read(io.BytesIO(data), dtype='float32', always_2d=False)
-            if isinstance(audio, np.ndarray):
-                if audio.ndim > 1:
-                    audio = np.mean(audio, axis=1)
-                out = audio.astype(np.float32)
-                if sr != fe.sampling_rate:
-                    out = librosa.resample(out, orig_sr=sr, target_sr=fe.sampling_rate)
-                if out.size < fe.sampling_rate // 10:
-                    out = np.pad(out, (0, max(0, fe.sampling_rate - out.size)), mode='constant')
-                return out
-
-            arr = np.array(audio, dtype=np.float32)
-            if sr and sr != fe.sampling_rate:
-                arr = librosa.resample(arr, orig_sr=sr, target_sr=fe.sampling_rate)
-            if arr.size < fe.sampling_rate // 10:
-                arr = np.pad(arr, (0, max(0, fe.sampling_rate - arr.size)), mode='constant')
-            return arr
-
+            if audio.ndim > 1:
+                audio = np.mean(audio, axis=1)
+            if sr != fe.sampling_rate:
+                audio = librosa.resample(audio, orig_sr=sr, target_sr=fe.sampling_rate)
+            return audio.astype(np.float32)
         except Exception:
-            try:
-                with tempfile.NamedTemporaryFile(delete=True, suffix='.audio') as tmp:
-                    tmp.write(data)
-                    tmp.flush()
-                    y, _sr = librosa.load(tmp.name, sr=fe.sampling_rate, mono=True)
-                    return y.astype(np.float32)
-            except Exception:
-                return np.zeros(fe.sampling_rate, dtype=np.float32)
-
+            # last fallback
+            with tempfile.NamedTemporaryFile(delete=True, suffix='.audio') as tmp:
+                tmp.write(data)
+                tmp.flush()
+                y, _ = librosa.load(tmp.name, sr=fe.sampling_rate, mono=True)
+                return y.astype(np.float32)
+
+# ---------------------------
+# Routes
+# ---------------------------
+@app.get("/")
+def root():
+    return {"status": "ok"}
 
-@app.post('/predict')
+@app.post("/predict")
 async def predict(file: UploadFile = File(...)):
     try:
+        # Read audio file
         data = await file.read()
         audio = to_wav16k_mono(data)
 
-        inputs = fe(audio, sampling_rate=fe.sampling_rate, return_tensors='pt')
+        # Extract features
+        inputs = fe(audio, sampling_rate=fe.sampling_rate, return_tensors="pt")
         inputs = {k: v.to(device) for k, v in inputs.items()}
 
+        # Forward pass
         with torch.no_grad():
             logits = model(**inputs).logits
 
         probs = torch.softmax(logits, dim=-1)[0].cpu().numpy()
-
         label_map = model.config.id2label
         labels = [label_map.get(str(i), f"class_{i}") for i in range(len(probs))]
 
@@ -103,23 +105,15 @@ async def predict(file: UploadFile = File(...)):
             reverse=True
         )
 
-        dominant = {
-            'label': pairs[0][0],
-            'score': pairs[0][1]
-        } if pairs else {'label': '', 'score': 0.0}
+        dominant = {"label": pairs[0][0], "score": pairs[0][1]} if pairs else {"label": "", "score": 0.0}
 
         return {
-            'results': [{'label': l, 'score': s} for l, s in pairs],
-            'dominant': dominant
+            "results": [{"label": l, "score": s} for l, s in pairs],
+            "dominant": dominant
         }
 
     except Exception as e:
         return JSONResponse(
             status_code=400,
-            content={'error': 'failed to process audio', 'message': f"{e.__class__.__name__}: {e}"}
+            content={"error": "failed to process audio", "message": f"{e.__class__.__name__}: {e}"}
         )
-
-
-@app.get('/')
-def root():
-    return {'status': 'ok'}