Update app.py

app.py

@@ -1,109 +1,87 @@
+from fastapi import FastAPI, UploadFile, File, Form
+import tempfile
+import shutil
+import uvicorn
+import whisperx
+import torch
 import numpy as np
-import librosa
-import gradio as gr
-from sklearn.preprocessing import StandardScaler
+import soundfile as sf
+from speechbrain.pretrained import EncoderClassifier
 
-CHUNK_DURATION = 30
-SIMILARITY_THRESHOLD = 0.75
+app = FastAPI()
 
-def extract_features(y, sr):
-    features = []
+device = "cpu"
 
-    mfcc = librosa.feature.mfcc(y=y, sr=sr, n_mfcc=13)
-    features.extend(np.mean(mfcc, axis=1))
+# Load models (light)
+asr_model = whisperx.load_model("small", device)
 
-    features.append(np.mean(librosa.feature.spectral_centroid(y=y, sr=sr)))
-    features.append(np.mean(librosa.feature.spectral_bandwidth(y=y, sr=sr)))
-    features.append(np.mean(librosa.feature.zero_crossing_rate(y)))
-    features.append(np.mean(librosa.feature.rms(y=y)))
+speaker_model = EncoderClassifier.from_hparams(
+    source="speechbrain/spkrec-ecapa-voxceleb",
+    run_opts={"device": device}
+)
 
-    return np.array(features)
+@app.post("/transcribe")
+async def transcribe(audio: UploadFile = File(...), lang: str = Form("en")):
 
-def split_audio(y, sr, frame_sec=1.0):
-    frame_len = int(sr * frame_sec)
-    segments, times = [], []
+    temp = tempfile.NamedTemporaryFile(delete=False, suffix=".wav")
+    with temp as buffer:
+        shutil.copyfileobj(audio.file, buffer)
 
-    for i in range(0, len(y), frame_len):
-        seg = y[i:i+frame_len]
-        if len(seg) < frame_len:
-            continue
+    audio_path = temp.name
 
-        if np.mean(np.abs(seg)) > 0.01:
-            segments.append(seg)
-            times.append((i/sr, (i+frame_len)/sr))
+    # Load audio
+    audio_data = whisperx.load_audio(audio_path)
 
-    return segments, times
+    # Transcribe
+    result = asr_model.transcribe(audio_data, language=lang)
+    segments = result["segments"]
 
-def process_audio(file_path):
-    y, sr = librosa.load(file_path, sr=None)
-    total_duration = len(y) / sr
+    y, sr = sf.read(audio_path)
 
-    all_segments = []
     speaker_embeddings = []
     speaker_labels = []
-    speaker_count = 0
 
-    current_time = 0
+    final_segments = []
+
+    for i, seg in enumerate(segments):
+
+        start = int(seg["start"] * sr)
+        end = int(seg["end"] * sr)
 
-    while current_time < total_duration:
-        start = int(current_time * sr)
-        end = int(min((current_time + CHUNK_DURATION) * sr, len(y)))
         chunk = y[start:end]
 
-        segments, times = split_audio(chunk, sr)
-
-        for seg, (s, e) in zip(segments, times):
-            feat = extract_features(seg, sr)
-
-            if speaker_embeddings:
-                scaler = StandardScaler()
-                X = np.vstack(speaker_embeddings + [feat])
-                X = scaler.fit_transform(X)
-                feat_norm = X[-1]
-                existing = X[:-1]
-            else:
-                feat_norm = feat
-                existing = []
-
-            assigned = False
-            for i, emb in enumerate(existing):
-                sim = np.dot(feat_norm, emb) / (
-                    np.linalg.norm(feat_norm) * np.linalg.norm(emb)
+        if len(chunk) < sr * 0.5:  # skip very short
+            continue
+
+        chunk_tensor = torch.tensor(chunk).unsqueeze(0)
+
+        emb = speaker_model.encode_batch(chunk_tensor)
+        emb = emb.squeeze().detach().cpu().numpy()
+
+        # Assign speakers
+        if len(speaker_embeddings) < 2:
+            speaker_id = f"SPEAKER_{len(speaker_embeddings)+1}"
+            speaker_embeddings.append(emb)
+            speaker_labels.append(speaker_id)
+        else:
+            sims = []
+            for e in speaker_embeddings:
+                sim = np.dot(emb, e) / (
+                    np.linalg.norm(emb) * np.linalg.norm(e)
                 )
-                if sim > SIMILARITY_THRESHOLD:
-                    speaker_id = speaker_labels[i]
-                    assigned = True
-                    break
-
-            if not assigned:
-                speaker_count += 1
-                speaker_id = f"SPEAKER_{speaker_count}"
-                speaker_embeddings.append(feat)
-                speaker_labels.append(speaker_id)
-
-            all_segments.append({
-                "speaker": speaker_id,
-                "start": round(current_time + s, 2),
-                "end": round(current_time + e, 2)
-            })
-
-        current_time += CHUNK_DURATION
-
-    return {"segments": all_segments}
-
-# 🎯 Gradio UI
-def run(audio):
-    if audio is None:
-        return {"error": "Upload audio"}
-
-    return process_audio(audio)
-
-demo = gr.Interface(
-    fn=run,
-    inputs=gr.Audio(type="filepath"),
-    outputs=gr.JSON(),
-    title="Speaker Diarization (CPU)",
-    description="Upload audio → get speaker labels with timestamps"
-)
+                sims.append(sim)
+
+            speaker_id = speaker_labels[np.argmax(sims)]
+
+        final_segments.append({
+            "speaker": speaker_id,
+            "start": round(seg["start"], 2),
+            "end": round(seg["end"], 2),
+            "text": seg["text"]
+        })
+
+    return {"segments": final_segments}
+
 
-demo.launch()
+if __name__ == "__main__":
+    uvicorn.run(app, host="0.0.0.0", port=7860)