Update app.py

app.py

@@ -1,20 +1,24 @@
 import numpy as np
 import librosa
 import gradio as gr
-import torch
+from sklearn.preprocessing import StandardScaler
 
-from speechbrain.inference.speaker import EncoderClassifier
+CHUNK_DURATION = 30
+SIMILARITY_THRESHOLD = 0.75
 
-# Load SpeechBrain model (CPU)
-classifier = EncoderClassifier.from_hparams(
-    source="speechbrain/spkrec-ecapa-voxceleb",
-    run_opts={"device": "cpu"}
-)
+def extract_features(y, sr):
+    features = []
 
-CHUNK_DURATION = 30
-SIMILARITY_THRESHOLD = 0.65
+    mfcc = librosa.feature.mfcc(y=y, sr=sr, n_mfcc=13)
+    features.extend(np.mean(mfcc, axis=1))
+
+    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)))
+
+    return np.array(features)
 
-# 🔹 Your SAME segmentation (unchanged)
 def split_audio(y, sr, frame_sec=1.0):
     frame_len = int(sr * frame_sec)
     segments, times = [], []
@@ -30,58 +34,14 @@ def split_audio(y, sr, frame_sec=1.0):
 
     return segments, times
 
-# 🔹 NEW: SpeechBrain embedding
-def get_embedding(seg, sr):
-    wav = torch.tensor(seg).unsqueeze(0)
-    emb = classifier.encode_batch(wav)
-    emb = emb.squeeze().detach().cpu().numpy()
-
-    # normalize
-    emb = emb / (np.linalg.norm(emb) + 1e-6)
-    return emb
-
-# 🔹 Speaker memory (strong)
-class SpeakerMemory:
-    def __init__(self):
-        self.db = {}
-        self.count = 0
-
-    def match(self, emb):
-        if not self.db:
-            return self._new(emb)
-
-        best_spk = None
-        best_score = -1
-
-        for spk, embs in self.db.items():
-            centroid = np.mean(embs, axis=0)
-            centroid = centroid / (np.linalg.norm(centroid) + 1e-6)
-
-            score = np.dot(emb, centroid)
-
-            if score > best_score:
-                best_score = score
-                best_spk = spk
-
-        if best_score > SIMILARITY_THRESHOLD:
-            self.db[best_spk].append(emb)
-            return best_spk
-        else:
-            return self._new(emb)
-
-    def _new(self, emb):
-        self.count += 1
-        spk = f"SPEAKER_{self.count}"
-        self.db[spk] = [emb]
-        return spk
-
-# 🔹 Main processing
 def process_audio(file_path):
     y, sr = librosa.load(file_path, sr=None)
     total_duration = len(y) / sr
 
-    memory = SpeakerMemory()
     all_segments = []
+    speaker_embeddings = []
+    speaker_labels = []
+    speaker_count = 0
 
     current_time = 0
 
@@ -93,8 +53,33 @@ def process_audio(file_path):
         segments, times = split_audio(chunk, sr)
 
         for seg, (s, e) in zip(segments, times):
-            emb = get_embedding(seg, sr)
-            speaker_id = memory.match(emb)
+            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 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,
@@ -110,14 +95,15 @@ def process_audio(file_path):
 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="High Accuracy Speaker Diarization (SpeechBrain)",
-    description="Hybrid: segmentation + ECAPA embeddings (~90% accuracy)"
+    title="Speaker Diarization (CPU)",
+    description="Upload audio → get speaker labels with timestamps"
 )
 
 demo.launch()