Spaces:
Sleeping
Sleeping
| import gradio as gr | |
| import numpy as np | |
| import queue | |
| import torch | |
| import time | |
| import threading | |
| import os | |
| import urllib.request | |
| import torchaudio | |
| from scipy.spatial.distance import cosine | |
| import json | |
| import io | |
| import wave | |
| # Simplified configuration parameters | |
| SILENCE_THRESHS = [0, 0.4] | |
| FINAL_TRANSCRIPTION_MODEL = "distil-large-v3" | |
| FINAL_BEAM_SIZE = 5 | |
| REALTIME_TRANSCRIPTION_MODEL = "distil-small.en" | |
| REALTIME_BEAM_SIZE = 5 | |
| TRANSCRIPTION_LANGUAGE = "en" | |
| SILERO_SENSITIVITY = 0.4 | |
| WEBRTC_SENSITIVITY = 3 | |
| MIN_LENGTH_OF_RECORDING = 0.7 | |
| PRE_RECORDING_BUFFER_DURATION = 0.35 | |
| # Speaker change detection parameters | |
| DEFAULT_CHANGE_THRESHOLD = 0.7 | |
| EMBEDDING_HISTORY_SIZE = 5 | |
| MIN_SEGMENT_DURATION = 1.0 | |
| DEFAULT_MAX_SPEAKERS = 4 | |
| ABSOLUTE_MAX_SPEAKERS = 10 | |
| # Global variables | |
| FAST_SENTENCE_END = True | |
| SAMPLE_RATE = 16000 | |
| BUFFER_SIZE = 512 | |
| CHANNELS = 1 | |
| # Speaker colors | |
| SPEAKER_COLORS = [ | |
| "#FFFF00", # Yellow | |
| "#FF0000", # Red | |
| "#00FF00", # Green | |
| "#00FFFF", # Cyan | |
| "#FF00FF", # Magenta | |
| "#0000FF", # Blue | |
| "#FF8000", # Orange | |
| "#00FF80", # Spring Green | |
| "#8000FF", # Purple | |
| "#FFFFFF", # White | |
| ] | |
| SPEAKER_COLOR_NAMES = [ | |
| "Yellow", "Red", "Green", "Cyan", "Magenta", | |
| "Blue", "Orange", "Spring Green", "Purple", "White" | |
| ] | |
| class SpeechBrainEncoder: | |
| """ECAPA-TDNN encoder from SpeechBrain for speaker embeddings""" | |
| def __init__(self, device="cpu"): | |
| self.device = device | |
| self.model = None | |
| self.embedding_dim = 192 | |
| self.model_loaded = False | |
| self.cache_dir = os.path.join(os.path.expanduser("~"), ".cache", "speechbrain") | |
| os.makedirs(self.cache_dir, exist_ok=True) | |
| def load_model(self): | |
| """Load the ECAPA-TDNN model""" | |
| try: | |
| from speechbrain.pretrained import EncoderClassifier | |
| self.model = EncoderClassifier.from_hparams( | |
| source="speechbrain/spkrec-ecapa-voxceleb", | |
| savedir=self.cache_dir, | |
| run_opts={"device": self.device} | |
| ) | |
| self.model_loaded = True | |
| print("ECAPA-TDNN model loaded successfully!") | |
| return True | |
| except Exception as e: | |
| print(f"SpeechBrain not available: {e}") | |
| return False | |
| def embed_utterance(self, audio, sr=16000): | |
| """Extract speaker embedding from audio""" | |
| if not self.model_loaded: | |
| raise ValueError("Model not loaded. Call load_model() first.") | |
| try: | |
| if isinstance(audio, np.ndarray): | |
| waveform = torch.tensor(audio, dtype=torch.float32).unsqueeze(0) | |
| else: | |
| waveform = audio.unsqueeze(0) | |
| if sr != 16000: | |
| waveform = torchaudio.functional.resample(waveform, orig_freq=sr, new_freq=16000) | |
| with torch.no_grad(): | |
| embedding = self.model.encode_batch(waveform) | |
| return embedding.squeeze().cpu().numpy() | |
| except Exception as e: | |
| print(f"Error extracting embedding: {e}") | |
| return np.zeros(self.embedding_dim) | |
| class AudioProcessor: | |
| """Processes audio data to extract speaker embeddings""" | |
| def __init__(self, encoder): | |
| self.encoder = encoder | |
| def extract_embedding(self, audio_data, sample_rate=16000): | |
| try: | |
| # Ensure audio is float32 and normalized | |
| if audio_data.dtype == np.int16: | |
| float_audio = audio_data.astype(np.float32) / 32768.0 | |
| else: | |
| float_audio = audio_data.astype(np.float32) | |
| # Normalize if needed | |
| if np.abs(float_audio).max() > 1.0: | |
| float_audio = float_audio / np.abs(float_audio).max() | |
| embedding = self.encoder.embed_utterance(float_audio, sample_rate) | |
| return embedding | |
| except Exception as e: | |
| print(f"Embedding extraction error: {e}") | |
| return np.zeros(self.encoder.embedding_dim) | |
| class SpeakerChangeDetector: | |
| """Speaker change detector that supports a configurable number of speakers""" | |
| def __init__(self, embedding_dim=192, change_threshold=DEFAULT_CHANGE_THRESHOLD, max_speakers=DEFAULT_MAX_SPEAKERS): | |
| self.embedding_dim = embedding_dim | |
| self.change_threshold = change_threshold | |
| self.max_speakers = min(max_speakers, ABSOLUTE_MAX_SPEAKERS) | |
| self.current_speaker = 0 | |
| self.previous_embeddings = [] | |
| self.last_change_time = time.time() | |
| self.mean_embeddings = [None] * self.max_speakers | |
| self.speaker_embeddings = [[] for _ in range(self.max_speakers)] | |
| self.last_similarity = 0.0 | |
| self.active_speakers = set([0]) | |
| def set_max_speakers(self, max_speakers): | |
| """Update the maximum number of speakers""" | |
| new_max = min(max_speakers, ABSOLUTE_MAX_SPEAKERS) | |
| if new_max < self.max_speakers: | |
| for speaker_id in list(self.active_speakers): | |
| if speaker_id >= new_max: | |
| self.active_speakers.discard(speaker_id) | |
| if self.current_speaker >= new_max: | |
| self.current_speaker = 0 | |
| if new_max > self.max_speakers: | |
| self.mean_embeddings.extend([None] * (new_max - self.max_speakers)) | |
| self.speaker_embeddings.extend([[] for _ in range(new_max - self.max_speakers)]) | |
| else: | |
| self.mean_embeddings = self.mean_embeddings[:new_max] | |
| self.speaker_embeddings = self.speaker_embeddings[:new_max] | |
| self.max_speakers = new_max | |
| def set_change_threshold(self, threshold): | |
| """Update the threshold for detecting speaker changes""" | |
| self.change_threshold = max(0.1, min(threshold, 0.99)) | |
| def add_embedding(self, embedding, timestamp=None): | |
| """Add a new embedding and check if there's a speaker change""" | |
| current_time = timestamp or time.time() | |
| if not self.previous_embeddings: | |
| self.previous_embeddings.append(embedding) | |
| self.speaker_embeddings[self.current_speaker].append(embedding) | |
| if self.mean_embeddings[self.current_speaker] is None: | |
| self.mean_embeddings[self.current_speaker] = embedding.copy() | |
| return self.current_speaker, 1.0 | |
| current_mean = self.mean_embeddings[self.current_speaker] | |
| if current_mean is not None: | |
| similarity = 1.0 - cosine(embedding, current_mean) | |
| else: | |
| similarity = 1.0 - cosine(embedding, self.previous_embeddings[-1]) | |
| self.last_similarity = similarity | |
| time_since_last_change = current_time - self.last_change_time | |
| is_speaker_change = False | |
| if time_since_last_change >= MIN_SEGMENT_DURATION: | |
| if similarity < self.change_threshold: | |
| best_speaker = self.current_speaker | |
| best_similarity = similarity | |
| for speaker_id in range(self.max_speakers): | |
| if speaker_id == self.current_speaker: | |
| continue | |
| speaker_mean = self.mean_embeddings[speaker_id] | |
| if speaker_mean is not None: | |
| speaker_similarity = 1.0 - cosine(embedding, speaker_mean) | |
| if speaker_similarity > best_similarity: | |
| best_similarity = speaker_similarity | |
| best_speaker = speaker_id | |
| if best_speaker != self.current_speaker: | |
| is_speaker_change = True | |
| self.current_speaker = best_speaker | |
| elif len(self.active_speakers) < self.max_speakers: | |
| for new_id in range(self.max_speakers): | |
| if new_id not in self.active_speakers: | |
| is_speaker_change = True | |
| self.current_speaker = new_id | |
| self.active_speakers.add(new_id) | |
| break | |
| if is_speaker_change: | |
| self.last_change_time = current_time | |
| self.previous_embeddings.append(embedding) | |
| if len(self.previous_embeddings) > EMBEDDING_HISTORY_SIZE: | |
| self.previous_embeddings.pop(0) | |
| self.speaker_embeddings[self.current_speaker].append(embedding) | |
| self.active_speakers.add(self.current_speaker) | |
| if len(self.speaker_embeddings[self.current_speaker]) > 30: | |
| self.speaker_embeddings[self.current_speaker] = self.speaker_embeddings[self.current_speaker][-30:] | |
| if self.speaker_embeddings[self.current_speaker]: | |
| self.mean_embeddings[self.current_speaker] = np.mean( | |
| self.speaker_embeddings[self.current_speaker], axis=0 | |
| ) | |
| return self.current_speaker, similarity | |
| def get_color_for_speaker(self, speaker_id): | |
| """Return color for speaker ID""" | |
| if 0 <= speaker_id < len(SPEAKER_COLORS): | |
| return SPEAKER_COLORS[speaker_id] | |
| return "#FFFFFF" | |
| def get_status_info(self): | |
| """Return status information about the speaker change detector""" | |
| speaker_counts = [len(self.speaker_embeddings[i]) for i in range(self.max_speakers)] | |
| return { | |
| "current_speaker": self.current_speaker, | |
| "speaker_counts": speaker_counts, | |
| "active_speakers": len(self.active_speakers), | |
| "max_speakers": self.max_speakers, | |
| "last_similarity": self.last_similarity, | |
| "threshold": self.change_threshold | |
| } | |
| class GradioSpeakerDiarization: | |
| def __init__(self): | |
| self.encoder = None | |
| self.audio_processor = None | |
| self.speaker_detector = None | |
| self.full_sentences = [] | |
| self.sentence_speakers = [] | |
| self.is_initialized = False | |
| self.change_threshold = DEFAULT_CHANGE_THRESHOLD | |
| self.max_speakers = DEFAULT_MAX_SPEAKERS | |
| def initialize_models(self): | |
| """Initialize the speaker encoder model""" | |
| try: | |
| device_str = "cuda" if torch.cuda.is_available() else "cpu" | |
| print(f"Using device: {device_str}") | |
| # Load SpeechBrain encoder | |
| self.encoder = SpeechBrainEncoder(device=device_str) | |
| success = self.encoder.load_model() | |
| if success: | |
| self.audio_processor = AudioProcessor(self.encoder) | |
| self.speaker_detector = SpeakerChangeDetector( | |
| embedding_dim=self.encoder.embedding_dim, | |
| change_threshold=self.change_threshold, | |
| max_speakers=self.max_speakers | |
| ) | |
| self.is_initialized = True | |
| return True | |
| else: | |
| return False | |
| except Exception as e: | |
| print(f"Model initialization error: {e}") | |
| return False | |
| def transcribe_audio(self, audio_input): | |
| """Process audio input and perform transcription with speaker diarization""" | |
| if not self.is_initialized: | |
| return "❌ Please initialize the system first!", self.get_formatted_conversation(), self.get_status_info() | |
| if audio_input is None: | |
| return "No audio received", self.get_formatted_conversation(), self.get_status_info() | |
| try: | |
| # Handle different audio input formats | |
| if isinstance(audio_input, tuple): | |
| sample_rate, audio_data = audio_input | |
| else: | |
| # Assume it's a file path | |
| import librosa | |
| audio_data, sample_rate = librosa.load(audio_input, sr=16000) | |
| # Ensure audio is in the right format | |
| if len(audio_data.shape) > 1: | |
| audio_data = audio_data.mean(axis=1) # Convert to mono | |
| # Perform simple transcription (placeholder - you'd want to integrate with Whisper or similar) | |
| # For now, we'll just do speaker diarization | |
| transcription = f"Audio segment {len(self.full_sentences) + 1} (duration: {len(audio_data)/sample_rate:.1f}s)" | |
| # Extract speaker embedding | |
| speaker_embedding = self.audio_processor.extract_embedding(audio_data, sample_rate) | |
| # Store sentence and embedding | |
| self.full_sentences.append((transcription, speaker_embedding)) | |
| # Detect speaker changes | |
| speaker_id, similarity = self.speaker_detector.add_embedding(speaker_embedding) | |
| self.sentence_speakers.append(speaker_id) | |
| status_msg = f"✅ Processed audio segment. Detected as Speaker {speaker_id + 1} (similarity: {similarity:.3f})" | |
| return status_msg, self.get_formatted_conversation(), self.get_status_info() | |
| except Exception as e: | |
| error_msg = f"❌ Error processing audio: {str(e)}" | |
| return error_msg, self.get_formatted_conversation(), self.get_status_info() | |
| def clear_conversation(self): | |
| """Clear all conversation data""" | |
| self.full_sentences = [] | |
| self.sentence_speakers = [] | |
| if self.speaker_detector: | |
| self.speaker_detector = SpeakerChangeDetector( | |
| embedding_dim=self.encoder.embedding_dim, | |
| change_threshold=self.change_threshold, | |
| max_speakers=self.max_speakers | |
| ) | |
| return "Conversation cleared!", self.get_formatted_conversation(), self.get_status_info() | |
| def update_settings(self, threshold, max_speakers): | |
| """Update speaker detection settings""" | |
| self.change_threshold = threshold | |
| self.max_speakers = max_speakers | |
| if self.speaker_detector: | |
| self.speaker_detector.set_change_threshold(threshold) | |
| self.speaker_detector.set_max_speakers(max_speakers) | |
| status_msg = f"Settings updated: Threshold={threshold:.2f}, Max Speakers={max_speakers}" | |
| return status_msg, self.get_formatted_conversation(), self.get_status_info() | |
| def get_formatted_conversation(self): | |
| """Get the formatted conversation with speaker colors""" | |
| try: | |
| if not self.full_sentences: | |
| return "No audio processed yet. Upload an audio file or record using the microphone." | |
| sentences_with_style = [] | |
| for i, sentence in enumerate(self.full_sentences): | |
| sentence_text, _ = sentence | |
| if i >= len(self.sentence_speakers): | |
| color = "#FFFFFF" | |
| speaker_name = "Unknown" | |
| else: | |
| speaker_id = self.sentence_speakers[i] | |
| color = self.speaker_detector.get_color_for_speaker(speaker_id) | |
| speaker_name = f"Speaker {speaker_id + 1}" | |
| sentences_with_style.append( | |
| f'<span style="color:{color};"><b>{speaker_name}:</b> {sentence_text}</span>') | |
| return "<br><br>".join(sentences_with_style) | |
| except Exception as e: | |
| return f"Error formatting conversation: {e}" | |
| def get_status_info(self): | |
| """Get current status information""" | |
| if not self.speaker_detector: | |
| return "Speaker detector not initialized" | |
| try: | |
| status = self.speaker_detector.get_status_info() | |
| status_lines = [ | |
| f"**Current Speaker:** {status['current_speaker'] + 1}", | |
| f"**Active Speakers:** {status['active_speakers']} of {status['max_speakers']}", | |
| f"**Last Similarity:** {status['last_similarity']:.3f}", | |
| f"**Change Threshold:** {status['threshold']:.2f}", | |
| f"**Total Segments:** {len(self.full_sentences)}", | |
| "", | |
| "**Speaker Segment Counts:**" | |
| ] | |
| for i in range(status['max_speakers']): | |
| color_name = SPEAKER_COLOR_NAMES[i] if i < len(SPEAKER_COLOR_NAMES) else f"Speaker {i+1}" | |
| status_lines.append(f"Speaker {i+1} ({color_name}): {status['speaker_counts'][i]}") | |
| return "\n".join(status_lines) | |
| except Exception as e: | |
| return f"Error getting status: {e}" | |
| # Global instance | |
| diarization_system = GradioSpeakerDiarization() | |
| def initialize_system(): | |
| """Initialize the diarization system""" | |
| success = diarization_system.initialize_models() | |
| if success: | |
| return "✅ System initialized successfully! Models loaded.", "", "" | |
| else: | |
| return "❌ Failed to initialize system. Please check the logs.", "", "" | |
| def process_audio(audio): | |
| """Process uploaded or recorded audio""" | |
| return diarization_system.transcribe_audio(audio) | |
| def clear_conversation(): | |
| """Clear the conversation""" | |
| return diarization_system.clear_conversation() | |
| def update_settings(threshold, max_speakers): | |
| """Update system settings""" | |
| return diarization_system.update_settings(threshold, max_speakers) | |
| # Create Gradio interface | |
| def create_interface(): | |
| with gr.Blocks(title="Speaker Diarization", theme=gr.themes.Soft()) as app: | |
| gr.Markdown("# 🎤 Audio Speaker Diarization") | |
| gr.Markdown("Upload audio files or record directly to identify different speakers using voice characteristics.") | |
| with gr.Row(): | |
| with gr.Column(scale=2): | |
| # Initialize button | |
| with gr.Row(): | |
| init_btn = gr.Button("🔧 Initialize System", variant="primary", size="lg") | |
| # Audio input options | |
| gr.Markdown("### 📁 Audio Input") | |
| with gr.Tab("Upload Audio File"): | |
| audio_file = gr.Audio( | |
| label="Upload Audio File", | |
| type="filepath", | |
| sources=["upload"] | |
| ) | |
| process_file_btn = gr.Button("Process Audio File", variant="secondary") | |
| with gr.Tab("Record Audio"): | |
| audio_mic = gr.Audio( | |
| label="Record Audio", | |
| type="numpy", | |
| sources=["microphone"] | |
| ) | |
| process_mic_btn = gr.Button("Process Recording", variant="secondary") | |
| # Results display | |
| status_output = gr.Textbox( | |
| label="Status", | |
| value="Click 'Initialize System' to start...", | |
| lines=2, | |
| interactive=False | |
| ) | |
| conversation_output = gr.HTML( | |
| value="<i>System not initialized...</i>", | |
| label="Speaker Analysis Results" | |
| ) | |
| # Control buttons | |
| with gr.Row(): | |
| clear_btn = gr.Button("🗑️ Clear Results", variant="stop") | |
| with gr.Column(scale=1): | |
| # Settings panel | |
| gr.Markdown("## ⚙️ Settings") | |
| threshold_slider = gr.Slider( | |
| minimum=0.1, | |
| maximum=0.95, | |
| step=0.05, | |
| value=DEFAULT_CHANGE_THRESHOLD, | |
| label="Speaker Change Sensitivity", | |
| info="Lower = more sensitive to speaker changes" | |
| ) | |
| max_speakers_slider = gr.Slider( | |
| minimum=2, | |
| maximum=ABSOLUTE_MAX_SPEAKERS, | |
| step=1, | |
| value=DEFAULT_MAX_SPEAKERS, | |
| label="Maximum Number of Speakers" | |
| ) | |
| update_settings_btn = gr.Button("Update Settings", variant="secondary") | |
| # System status | |
| system_status = gr.Textbox( | |
| label="System Status", | |
| value="System not initialized", | |
| lines=12, | |
| interactive=False | |
| ) | |
| # Speaker color legend | |
| gr.Markdown("## 🎨 Speaker Colors") | |
| color_info = [] | |
| for i, (color, name) in enumerate(zip(SPEAKER_COLORS[:DEFAULT_MAX_SPEAKERS], SPEAKER_COLOR_NAMES[:DEFAULT_MAX_SPEAKERS])): | |
| color_info.append(f'<span style="color:{color};">●</span> Speaker {i+1} ({name})') | |
| gr.HTML("<br>".join(color_info)) | |
| # Event handlers | |
| init_btn.click( | |
| initialize_system, | |
| outputs=[status_output, conversation_output, system_status] | |
| ) | |
| process_file_btn.click( | |
| process_audio, | |
| inputs=[audio_file], | |
| outputs=[status_output, conversation_output, system_status] | |
| ) | |
| process_mic_btn.click( | |
| process_audio, | |
| inputs=[audio_mic], | |
| outputs=[status_output, conversation_output, system_status] | |
| ) | |
| clear_btn.click( | |
| clear_conversation, | |
| outputs=[status_output, conversation_output, system_status] | |
| ) | |
| update_settings_btn.click( | |
| update_settings, | |
| inputs=[threshold_slider, max_speakers_slider], | |
| outputs=[status_output, conversation_output, system_status] | |
| ) | |
| return app | |
| if __name__ == "__main__": | |
| app = create_interface() | |
| app.launch( | |
| server_name="0.0.0.0", | |
| server_port=7860 | |
| ) | |