Upload folder using huggingface_hub

README.md

@@ -9,11 +9,8 @@ pipeline_tag: automatic-speech-recognition
 
 This package provides an **Automatic Speech Recognition (ASR) + Speaker Diarization** pipeline using:
 - [OpenAI Whisper](https://huggingface.co/openai/whisper-medium)
-- [Pyannote diarization](https://huggingface.co/pyannote/speaker-diarization-3.1)
+- [SpeechBrain ECAPA-TDNN](https://huggingface.co/speechbrain/spkrec-ecapa-voxceleb)
 
 ## Install
 ```bash
-pip install git+https://huggingface.co/Capstone04/asr-diarization-pipeline
-
-## Speaker Identification
-You can now enroll known speakers by providing reference audio samples. The pipeline will match incoming speaker segments against stored embeddings and label them accordingly. Unknown speakers are dynamically tracked per session.
+pip install git+https://huggingface.co/Capstone04/Prayashi_RealTime

asr_diarization/inference.py

@@ -1,6 +1,7 @@
 import os
 from .pipeline import ASR_Diarization
 
+
 import json
 import numpy as np
 
@@ -13,7 +14,7 @@ def load_known_embeddings(path="known_speakers.json"):
 
 HF_TOKEN = os.environ.get("HF_TOKEN", None)
 known_embeddings = load_known_embeddings()
-pipe = ASR_Diarization(HF_TOKEN, known_embeddings=known_embeddings)
+pipe = ASR_Diarization(HF_TOKEN)
 
 def inference(inputs):
     return pipe(inputs)

asr_diarization/pipeline.py

@@ -1,12 +1,6 @@
 import os
 import json
 import torch
-
-# Fix TF32 reproducibility warning and potential computation issues
-if torch.cuda.is_available():
-    torch.backends.cuda.matmul.allow_tf32 = True
-    torch.backends.cudnn.allow_tf32 = True
-
 import tempfile
 import torchaudio
 import threading
@@ -15,21 +9,18 @@ import soundfile as sf
 import noisereduce as nr
 from scipy import signal
 from numpy.linalg import norm
-from pyannote.audio import Pipeline
-from speechbrain.pretrained import EncoderClassifier
-from pyannote.core import Annotation, Segment
+from speechbrain.pretrained import SpeakerRecognition, EncoderClassifier
+from speechbrain.pretrained import SpectralMaskEnhancement
 from transformers import pipeline as hf_pipeline
-from pyannote.metrics.diarization import DiarizationErrorRate
 from jiwer import wer, Compose, ToLowerCase, RemovePunctuation, RemoveMultipleSpaces, Strip
 
 class ASR_Diarization:
-    def __init__(self, HF_TOKEN,
-                 diar_model="pyannote/speaker-diarization-3.1",
-                 asr_model="openai/whisper-medium"):
+    def __init__(self, HF_TOKEN, asr_model="openai/whisper-medium"):
         self.HF_TOKEN = HF_TOKEN
         self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
         self._unknown_lock = threading.Lock()
 
+        # Load SpeechBrain models
         try:
             self.embedding_model = EncoderClassifier.from_hparams(
                 source="speechbrain/spkrec-ecapa-voxceleb",
@@ -40,9 +31,17 @@ class ASR_Diarization:
             self.embedding_model = None
             print(f"[ERROR] Failed to load ECAPA: {e}")
 
-        self.diar_pipeline = Pipeline.from_pretrained(diar_model, use_auth_token=None)
-diar_model="pyannote/speaker-diarization-3.1"
+        try:
+            self.speaker_diarization = SpeakerRecognition.from_hparams(
+                source="speechbrain/spkrec-ecapa-voxceleb",
+                savedir="pretrained_models/spkrec-ecapa-voxceleb"
+            )
+            print("[Speaker Recognition] Model loaded successfully.")
+        except Exception as e:
+            self.speaker_diarization = None
+            print(f"[ERROR] Failed to load Speaker Recognition: {e}")
 
+        # Load ASR pipeline
         device_index = 0 if torch.cuda.is_available() else -1
         self.asr_pipeline = hf_pipeline(
             "automatic-speech-recognition",
@@ -52,11 +51,70 @@ diar_model="pyannote/speaker-diarization-3.1"
         )
 
     def run_diarization(self, audio_path):
-        diarization = self.diar_pipeline(audio_path)
-        return [
-            {"start": t.start, "end": t.end, "speaker": spk}
-            for t, _, spk in diarization.itertracks(yield_label=True)
-        ]
+        """Simple diarization using SpeechBrain embedding clustering"""
+        audio, sr = torchaudio.load(audio_path)
+        audio_np = audio[0].numpy() if audio.shape[0] == 1 else audio.mean(dim=0).numpy()
+        
+        # Segment audio into chunks for diarization
+        chunk_duration = 2.0  # 2-second chunks
+        chunk_size = int(chunk_duration * sr)
+        segments = []
+        
+        for i in range(0, len(audio_np), chunk_size):
+            start_time = i / sr
+            end_time = min((i + chunk_size) / sr, len(audio_np) / sr)
+            chunk = audio_np[i:i+chunk_size]
+            
+            if len(chunk) < 8000:  # Skip very short chunks
+                continue
+                
+            # Get speaker embedding for this chunk
+            if self.embedding_model:
+                try:
+                    chunk_tensor = torch.from_numpy(chunk).unsqueeze(0).to(self.device)
+                    with torch.no_grad():
+                        embedding = self.embedding_model.encode_batch(chunk_tensor).squeeze().cpu().numpy()
+                    
+                    # Simple speaker assignment based on embedding similarity
+                    speaker_id = self._assign_speaker(embedding, segments)
+                    
+                    segments.append({
+                        "start": start_time,
+                        "end": end_time,
+                        "speaker": speaker_id,
+                        "embedding": embedding
+                    })
+                except Exception as e:
+                    print(f"Error processing chunk: {e}")
+                    continue
+        
+        return segments
+
+    def _assign_speaker(self, embedding, existing_segments, threshold=0.7):
+        """Assign speaker based on embedding similarity"""
+        if not existing_segments:
+            return "speaker_1"
+        
+        # Calculate similarity with existing speakers
+        similarities = []
+        for seg in existing_segments[-10:]:  # Check last 10 segments
+            if "embedding" in seg:
+                sim = np.dot(embedding.flatten(), seg["embedding"].flatten()) / (
+                    norm(embedding.flatten()) * norm(seg["embedding"].flatten())
+                )
+                similarities.append((seg["speaker"], sim))
+        
+        if similarities:
+            best_speaker, best_sim = max(similarities, key=lambda x: x[1])
+            if best_sim > threshold:
+                return best_speaker
+        
+        # Create new speaker
+        existing_speakers = set(seg["speaker"] for seg in existing_segments)
+        speaker_num = 1
+        while f"speaker_{speaker_num}" in existing_speakers:
+            speaker_num += 1
+        return f"speaker_{speaker_num}"
 
     def load_unknown_speakers(self, unknown_speakers_path):
         if os.path.exists(unknown_speakers_path):
@@ -123,7 +181,6 @@ diar_model="pyannote/speaker-diarization-3.1"
 
         return best_name, best_score, is_enrolled
 
-
     def run_transcription(self, audio_path, diar_json, enrolled_speakers=None, unknown_speakers_path=None):
         unknown_speakers_path = unknown_speakers_path or os.path.join(os.path.dirname(audio_path), "unknown_speakers.json")
 
@@ -137,6 +194,8 @@ diar_model="pyannote/speaker-diarization-3.1"
         enrolled_speakers_np = {n: v/norm(v) for n,v in (enrolled_speakers or {}).items() if norm(v) > 0}
 
         target_sr = 16000
+        
+        # Group segments by speaker for clustering
         clusters = {}
         for seg in diar_json:
             clusters.setdefault(seg["speaker"], []).append(seg)
@@ -163,7 +222,7 @@ diar_model="pyannote/speaker-diarization-3.1"
                 cluster_emb = np.mean(np.stack(seg_embs), axis=0)
                 cluster_embeddings[cluster_label] = cluster_emb / norm(cluster_emb)
 
-        speaker_map, speakers_updated = {}, False
+        speaker_map, speakers_updated = {}, {}
         threshold = 0.5
 
         # Thread-safe unknown speaker update
@@ -191,7 +250,7 @@ diar_model="pyannote/speaker-diarization-3.1"
             if speakers_updated:
                 self.save_unknown_speakers(unknown_speakers, unknown_speakers_path)
 
-        # ASR transcription (same as before)
+        # ASR transcription
         for seg in diar_json:
             start, end, spk = seg["start"], seg["end"], seg["speaker"]
             start_sample, end_sample = int(start*sr), int(end*sr)
@@ -276,19 +335,6 @@ diar_model="pyannote/speaker-diarization-3.1"
         hyp_rttm = os.path.join(output_dir, f"{base_name}.rttm")
         hyp_json = os.path.join(output_dir, f"{base_name}_merged_transcription.json")
 
-        if ref_rttm:
-            def load_rttm(path):
-                ann = Annotation()
-                for line in open(path):
-                    if line.startswith("SPEAKER"):
-                        p = line.split()
-                        start, dur, spk = float(p[3]), float(p[4]), p[7]
-                        ann[Segment(start, start+dur)] = spk
-                return ann
-
-            der_score = DiarizationErrorRate()(load_rttm(ref_rttm), load_rttm(hyp_rttm))
-            results["DER"] = round(der_score * 100, 2)
-
         if ref_json:
             def load_words(path):
                 data = json.load(open(path))

requirements.txt

@@ -1,10 +1,8 @@
-torch>=2.3.0
-torchaudio>=2.3.0
-pyannote.audio==3.1.1
-transformers>=4.41.0
-huggingface_hub>=0.24.0
+torch
+torchaudio
+speechbrain
+transformers
 noisereduce
 scikit-learn
 jiwer
 librosa
-speechbrain

setup.py

@@ -7,7 +7,7 @@ setup(
     install_requires=[
         "torch",
         "torchaudio",
-        "pyannote.audio",
+        "speechbrain",
         "transformers",
         "noisereduce",
         "scikit-learn",