AndroidGuy/Real_Time_diarization
0
1import gradio as gr2import numpy as np3import queue4import torch5import time6import threading7import os8import urllib.request9import torchaudio10from scipy.spatial.distance import cosine11from scipy.signal import resample12from RealtimeSTT import AudioToTextRecorder13from fastapi import FastAPI, APIRouter14from fastrtc import Stream, AsyncStreamHandler15import json16import asyncio17import uvicorn18from queue import Queue19import logging20 21# Set up logging22logging.basicConfig(level=logging.INFO)23logger = logging.getLogger(__name__)24 25# Simplified configuration parameters26SILENCE_THRESHS = [0, 0.4]27FINAL_TRANSCRIPTION_MODEL = "distil-large-v3"28FINAL_BEAM_SIZE = 529REALTIME_TRANSCRIPTION_MODEL = "distil-small.en"30REALTIME_BEAM_SIZE = 531TRANSCRIPTION_LANGUAGE = "en"32SILERO_SENSITIVITY = 0.433WEBRTC_SENSITIVITY = 334MIN_LENGTH_OF_RECORDING = 0.735PRE_RECORDING_BUFFER_DURATION = 0.3536 37# Speaker change detection parameters38DEFAULT_CHANGE_THRESHOLD = 0.6539EMBEDDING_HISTORY_SIZE = 540MIN_SEGMENT_DURATION = 1.541DEFAULT_MAX_SPEAKERS = 442ABSOLUTE_MAX_SPEAKERS = 843 44# Global variables45SAMPLE_RATE = 1600046BUFFER_SIZE = 102447CHANNELS = 148 49# Speaker colors - more distinguishable colors50SPEAKER_COLORS = [51 "#FF6B6B", # Red52 "#4ECDC4", # Teal53 "#45B7D1", # Blue54 "#96CEB4", # Green55 "#FFEAA7", # Yellow56 "#DDA0DD", # Plum57 "#98D8C8", # Mint58 "#F7DC6F", # Gold59]60 61SPEAKER_COLOR_NAMES = [62 "Red", "Teal", "Blue", "Green", "Yellow", "Plum", "Mint", "Gold"63]64 65 66class SpeechBrainEncoder:67 """ECAPA-TDNN encoder from SpeechBrain for speaker embeddings"""68 def __init__(self, device="cpu"):69 self.device = device70 self.model = None71 self.embedding_dim = 19272 self.model_loaded = False73 self.cache_dir = os.path.join(os.path.expanduser("~"), ".cache", "speechbrain")74 os.makedirs(self.cache_dir, exist_ok=True)75 76 def load_model(self):77 """Load the ECAPA-TDNN model"""78 try:79 from speechbrain.pretrained import EncoderClassifier80 81 self.model = EncoderClassifier.from_hparams(82 source="speechbrain/spkrec-ecapa-voxceleb",83 savedir=self.cache_dir,84 run_opts={"device": self.device}85 )86 87 self.model_loaded = True88 logger.info("ECAPA-TDNN model loaded successfully!")89 return True90 except Exception as e:91 logger.error(f"Error loading ECAPA-TDNN model: {e}")92 return False93 94 def embed_utterance(self, audio, sr=16000):95 """Extract speaker embedding from audio"""96 if not self.model_loaded:97 raise ValueError("Model not loaded. Call load_model() first.")98 99 try:100 if isinstance(audio, np.ndarray):101 # Ensure audio is float32 and properly normalized102 audio = audio.astype(np.float32)103 if np.max(np.abs(audio)) > 1.0:104 audio = audio / np.max(np.abs(audio))105 waveform = torch.tensor(audio).unsqueeze(0)106 else:107 waveform = audio.unsqueeze(0)108 109 # Resample if necessary110 if sr != 16000:111 waveform = torchaudio.functional.resample(waveform, orig_freq=sr, new_freq=16000)112 113 with torch.no_grad():114 embedding = self.model.encode_batch(waveform)115 116 return embedding.squeeze().cpu().numpy()117 except Exception as e:118 logger.error(f"Error extracting embedding: {e}")119 return np.zeros(self.embedding_dim)120 121 122class AudioProcessor:123 """Processes audio data to extract speaker embeddings"""124 def __init__(self, encoder):125 self.encoder = encoder126 self.audio_buffer = []127 self.min_audio_length = int(SAMPLE_RATE * 1.0) # Minimum 1 second of audio128 129 def add_audio_chunk(self, audio_chunk):130 """Add audio chunk to buffer"""131 self.audio_buffer.extend(audio_chunk)132 133 # Keep buffer from getting too large134 max_buffer_size = int(SAMPLE_RATE * 10) # 10 seconds max135 if len(self.audio_buffer) > max_buffer_size:136 self.audio_buffer = self.audio_buffer[-max_buffer_size:]137 138 def extract_embedding_from_buffer(self):139 """Extract embedding from current audio buffer"""140 if len(self.audio_buffer) < self.min_audio_length:141 return None142 143 try:144 # Use the last portion of the buffer for embedding145 audio_segment = np.array(self.audio_buffer[-self.min_audio_length:], dtype=np.float32)146 147 # Normalize audio148 if np.max(np.abs(audio_segment)) > 0:149 audio_segment = audio_segment / np.max(np.abs(audio_segment))150 else:151 return None152 153 embedding = self.encoder.embed_utterance(audio_segment)154 return embedding155 except Exception as e:156 logger.error(f"Embedding extraction error: {e}")157 return None158 159 160class SpeakerChangeDetector:161 """Improved speaker change detector"""162 def __init__(self, embedding_dim=192, change_threshold=DEFAULT_CHANGE_THRESHOLD, max_speakers=DEFAULT_MAX_SPEAKERS):163 self.embedding_dim = embedding_dim164 self.change_threshold = change_threshold165 self.max_speakers = min(max_speakers, ABSOLUTE_MAX_SPEAKERS)166 self.current_speaker = 0167 self.speaker_embeddings = [[] for _ in range(self.max_speakers)]168 self.speaker_centroids = [None] * self.max_speakers169 self.last_change_time = time.time()170 self.last_similarity = 1.0171 self.active_speakers = set([0])172 self.segment_counter = 0173 174 def set_max_speakers(self, max_speakers):175 """Update the maximum number of speakers"""176 new_max = min(max_speakers, ABSOLUTE_MAX_SPEAKERS)177 178 if new_max < self.max_speakers:179 # Remove speakers beyond the new limit180 for speaker_id in list(self.active_speakers):181 if speaker_id >= new_max:182 self.active_speakers.discard(speaker_id)183 184 if self.current_speaker >= new_max:185 self.current_speaker = 0186 187 # Resize arrays188 if new_max > self.max_speakers:189 self.speaker_embeddings.extend([[] for _ in range(new_max - self.max_speakers)])190 self.speaker_centroids.extend([None] * (new_max - self.max_speakers))191 else:192 self.speaker_embeddings = self.speaker_embeddings[:new_max]193 self.speaker_centroids = self.speaker_centroids[:new_max]194 195 self.max_speakers = new_max196 197 def set_change_threshold(self, threshold):198 """Update the threshold for detecting speaker changes"""199 self.change_threshold = max(0.1, min(threshold, 0.95))200 201 def add_embedding(self, embedding, timestamp=None):202 """Add a new embedding and detect speaker changes"""203 current_time = timestamp or time.time()204 self.segment_counter += 1205 206 # Initialize first speaker207 if not self.speaker_embeddings[0]:208 self.speaker_embeddings[0].append(embedding)209 self.speaker_centroids[0] = embedding.copy()210 self.active_speakers.add(0)211 return 0, 1.0212 213 # Calculate similarity with current speaker214 current_centroid = self.speaker_centroids[self.current_speaker]215 if current_centroid is not None:216 similarity = 1.0 - cosine(embedding, current_centroid)217 else:218 similarity = 0.5219 220 self.last_similarity = similarity221 222 # Check for speaker change223 time_since_last_change = current_time - self.last_change_time224 speaker_changed = False225 226 if time_since_last_change >= MIN_SEGMENT_DURATION and similarity < self.change_threshold:227 # Find best matching speaker228 best_speaker = self.current_speaker229 best_similarity = similarity230 231 for speaker_id in self.active_speakers:232 if speaker_id == self.current_speaker:233 continue234 235 centroid = self.speaker_centroids[speaker_id]236 if centroid is not None:237 speaker_similarity = 1.0 - cosine(embedding, centroid)238 if speaker_similarity > best_similarity and speaker_similarity > self.change_threshold:239 best_similarity = speaker_similarity240 best_speaker = speaker_id241 242 # If no good match found and we can add a new speaker243 if best_speaker == self.current_speaker and len(self.active_speakers) < self.max_speakers:244 for new_id in range(self.max_speakers):245 if new_id not in self.active_speakers:246 best_speaker = new_id247 self.active_speakers.add(new_id)248 break249 250 if best_speaker != self.current_speaker:251 self.current_speaker = best_speaker252 self.last_change_time = current_time253 speaker_changed = True254 255 # Update speaker embeddings and centroids256 self.speaker_embeddings[self.current_speaker].append(embedding)257 258 # Keep only recent embeddings (sliding window)259 max_embeddings = 20260 if len(self.speaker_embeddings[self.current_speaker]) > max_embeddings:261 self.speaker_embeddings[self.current_speaker] = self.speaker_embeddings[self.current_speaker][-max_embeddings:]262 263 # Update centroid264 if self.speaker_embeddings[self.current_speaker]:265 self.speaker_centroids[self.current_speaker] = np.mean(266 self.speaker_embeddings[self.current_speaker], axis=0267 )268 269 return self.current_speaker, similarity270 271 def get_color_for_speaker(self, speaker_id):272 """Return color for speaker ID"""273 if 0 <= speaker_id < len(SPEAKER_COLORS):274 return SPEAKER_COLORS[speaker_id]275 return "#FFFFFF"276 277 def get_status_info(self):278 """Return status information"""279 speaker_counts = [len(self.speaker_embeddings[i]) for i in range(self.max_speakers)]280 281 return {282 "current_speaker": self.current_speaker,283 "speaker_counts": speaker_counts,284 "active_speakers": len(self.active_speakers),285 "max_speakers": self.max_speakers,286 "last_similarity": self.last_similarity,287 "threshold": self.change_threshold,288 "segment_counter": self.segment_counter289 }290 291 292class RealtimeSpeakerDiarization:293 def __init__(self):294 self.encoder = None295 self.audio_processor = None296 self.speaker_detector = None297 self.recorder = None298 self.sentence_queue = queue.Queue()299 self.full_sentences = []300 self.sentence_speakers = []301 self.pending_sentences = []302 self.current_conversation = ""303 self.is_running = False304 self.change_threshold = DEFAULT_CHANGE_THRESHOLD305 self.max_speakers = DEFAULT_MAX_SPEAKERS306 self.last_transcription = ""307 self.transcription_lock = threading.Lock()308 309 def initialize_models(self):310 """Initialize the speaker encoder model"""311 try:312 device_str = "cuda" if torch.cuda.is_available() else "cpu"313 logger.info(f"Using device: {device_str}")314 315 self.encoder = SpeechBrainEncoder(device=device_str)316 success = self.encoder.load_model()317 318 if success:319 self.audio_processor = AudioProcessor(self.encoder)320 self.speaker_detector = SpeakerChangeDetector(321 embedding_dim=self.encoder.embedding_dim,322 change_threshold=self.change_threshold,323 max_speakers=self.max_speakers324 )325 logger.info("Models initialized successfully!")326 return True327 else:328 logger.error("Failed to load models")329 return False330 except Exception as e:331 logger.error(f"Model initialization error: {e}")332 return False333 334 def live_text_detected(self, text):335 """Callback for real-time transcription updates"""336 with self.transcription_lock:337 self.last_transcription = text.strip()338 339 def process_final_text(self, text):340 """Process final transcribed text with speaker embedding"""341 text = text.strip()342 if text:343 try:344 # Get audio data for this transcription345 audio_bytes = getattr(self.recorder, 'last_transcription_bytes', None)346 if audio_bytes:347 self.sentence_queue.put((text, audio_bytes))348 else:349 # If no audio bytes, use current speaker350 self.sentence_queue.put((text, None))351 352 except Exception as e:353 logger.error(f"Error processing final text: {e}")354 355 def process_sentence_queue(self):356 """Process sentences in the queue for speaker detection"""357 while self.is_running:358 try:359 text, audio_bytes = self.sentence_queue.get(timeout=1)360 361 current_speaker = self.speaker_detector.current_speaker362 363 if audio_bytes:364 # Convert audio data and extract embedding365 audio_int16 = np.frombuffer(audio_bytes, dtype=np.int16)366 audio_float = audio_int16.astype(np.float32) / 32768.0367 368 # Extract embedding369 embedding = self.audio_processor.encoder.embed_utterance(audio_float)370 if embedding is not None:371 current_speaker, similarity = self.speaker_detector.add_embedding(embedding)372 373 # Store sentence with speaker374 with self.transcription_lock:375 self.full_sentences.append((text, current_speaker))376 self.update_conversation_display()377 378 except queue.Empty:379 continue380 except Exception as e:381 logger.error(f"Error processing sentence: {e}")382 383 def update_conversation_display(self):384 """Update the conversation display"""385 try:386 sentences_with_style = []387 388 for sentence_text, speaker_id in self.full_sentences:389 color = self.speaker_detector.get_color_for_speaker(speaker_id)390 speaker_name = f"Speaker {speaker_id + 1}"391 sentences_with_style.append(392 f'<span style="color:{color}; font-weight: bold;">{speaker_name}:</span> '393 f'<span style="color:#333333;">{sentence_text}</span>'394 )395 396 # Add current transcription if available397 if self.last_transcription:398 current_color = self.speaker_detector.get_color_for_speaker(self.speaker_detector.current_speaker)399 current_speaker = f"Speaker {self.speaker_detector.current_speaker + 1}"400 sentences_with_style.append(401 f'<span style="color:{current_color}; font-weight: bold; opacity: 0.7;">{current_speaker}:</span> '402 f'<span style="color:#666666; font-style: italic;">{self.last_transcription}...</span>'403 )404 405 if sentences_with_style:406 self.current_conversation = "<br><br>".join(sentences_with_style)407 else:408 self.current_conversation = "<i>Waiting for speech input...</i>"409 410 except Exception as e:411 logger.error(f"Error updating conversation display: {e}")412 self.current_conversation = f"<i>Error: {str(e)}</i>"413 414 def start_recording(self):415 """Start the recording and transcription process"""416 if self.encoder is None:417 return "Please initialize models first!"418 419 try:420 # Setup recorder configuration421 recorder_config = {422 'spinner': False,423 'use_microphone': False, # Using FastRTC for audio input424 'model': FINAL_TRANSCRIPTION_MODEL,425 'language': TRANSCRIPTION_LANGUAGE,426 'silero_sensitivity': SILERO_SENSITIVITY,427 'webrtc_sensitivity': WEBRTC_SENSITIVITY,428 'post_speech_silence_duration': SILENCE_THRESHS[1],429 'min_length_of_recording': MIN_LENGTH_OF_RECORDING,430 'pre_recording_buffer_duration': PRE_RECORDING_BUFFER_DURATION,431 'min_gap_between_recordings': 0,432 'enable_realtime_transcription': True,433 'realtime_processing_pause': 0.1,434 'realtime_model_type': REALTIME_TRANSCRIPTION_MODEL,435 'on_realtime_transcription_update': self.live_text_detected,436 'beam_size': FINAL_BEAM_SIZE,437 'beam_size_realtime': REALTIME_BEAM_SIZE,438 'sample_rate': SAMPLE_RATE,439 }440 441 self.recorder = AudioToTextRecorder(**recorder_config)442 443 # Start processing threads444 self.is_running = True445 self.sentence_thread = threading.Thread(target=self.process_sentence_queue, daemon=True)446 self.sentence_thread.start()447 448 self.transcription_thread = threading.Thread(target=self.run_transcription, daemon=True)449 self.transcription_thread.start()450 451 return "Recording started successfully!"452 453 except Exception as e:454 logger.error(f"Error starting recording: {e}")455 return f"Error starting recording: {e}"456 457 def run_transcription(self):458 """Run the transcription loop"""459 try:460 logger.info("Starting transcription thread")461 while self.is_running:462 # Just check for final text from recorder, audio is fed externally via FastRTC463 text = self.recorder.text(self.process_final_text)464 time.sleep(0.01) # Small sleep to prevent CPU hogging465 except Exception as e:466 logger.error(f"Transcription error: {e}")467 468 def stop_recording(self):469 """Stop the recording process"""470 self.is_running = False471 if self.recorder:472 self.recorder.stop()473 return "Recording stopped!"474 475 def clear_conversation(self):476 """Clear all conversation data"""477 with self.transcription_lock:478 self.full_sentences = []479 self.last_transcription = ""480 self.current_conversation = "Conversation cleared!"481 482 if self.speaker_detector:483 self.speaker_detector = SpeakerChangeDetector(484 embedding_dim=self.encoder.embedding_dim,485 change_threshold=self.change_threshold,486 max_speakers=self.max_speakers487 )488 489 return "Conversation cleared!"490 491 def update_settings(self, threshold, max_speakers):492 """Update speaker detection settings"""493 self.change_threshold = threshold494 self.max_speakers = max_speakers495 496 if self.speaker_detector:497 self.speaker_detector.set_change_threshold(threshold)498 self.speaker_detector.set_max_speakers(max_speakers)499 500 return f"Settings updated: Threshold={threshold:.2f}, Max Speakers={max_speakers}"501 502 def get_formatted_conversation(self):503 """Get the formatted conversation"""504 return self.current_conversation505 506 def get_status_info(self):507 """Get current status information"""508 if not self.speaker_detector:509 return "Speaker detector not initialized"510 511 try:512 status = self.speaker_detector.get_status_info()513 514 status_lines = [515 f"**Current Speaker:** {status['current_speaker'] + 1}",516 f"**Active Speakers:** {status['active_speakers']} of {status['max_speakers']}",517 f"**Last Similarity:** {status['last_similarity']:.3f}",518 f"**Change Threshold:** {status['threshold']:.2f}",519 f"**Total Sentences:** {len(self.full_sentences)}",520 f"**Segments Processed:** {status['segment_counter']}",521 "",522 "**Speaker Activity:**"523 ]524 525 for i in range(status['max_speakers']):526 color_name = SPEAKER_COLOR_NAMES[i] if i < len(SPEAKER_COLOR_NAMES) else f"Speaker {i+1}"527 count = status['speaker_counts'][i]528 active = "🟢" if count > 0 else "⚫"529 status_lines.append(f"{active} Speaker {i+1} ({color_name}): {count} segments")530 531 return "\n".join(status_lines)532 533 except Exception as e:534 return f"Error getting status: {e}"535 536 def process_audio_chunk(self, audio_data, sample_rate=16000):537 """Process audio chunk from FastRTC input"""538 if not self.is_running or self.audio_processor is None:539 return540 541 try:542 # Ensure audio is float32543 if isinstance(audio_data, np.ndarray):544 if audio_data.dtype != np.float32:545 audio_data = audio_data.astype(np.float32)546 else:547 audio_data = np.array(audio_data, dtype=np.float32)548 549 # Ensure mono550 if len(audio_data.shape) > 1:551 audio_data = np.mean(audio_data, axis=1) if audio_data.shape[1] > 1 else audio_data.flatten()552 553 # Normalize if needed554 if np.max(np.abs(audio_data)) > 1.0:555 audio_data = audio_data / np.max(np.abs(audio_data))556 557 # Add to audio processor buffer for speaker detection558 self.audio_processor.add_audio_chunk(audio_data)559 560 # Periodically extract embeddings for speaker detection561 if len(self.audio_processor.audio_buffer) % (SAMPLE_RATE // 2) == 0: # Every 0.5 seconds562 embedding = self.audio_processor.extract_embedding_from_buffer()563 if embedding is not None:564 self.speaker_detector.add_embedding(embedding)565 566 # Feed audio to RealtimeSTT recorder567 if self.recorder and self.is_running:568 # Convert float32 [-1.0, 1.0] to int16 for RealtimeSTT569 int16_data = (audio_data * 32768.0).astype(np.int16).tobytes()570 if sample_rate != 16000:571 int16_data = self.resample_audio(int16_data, sample_rate, 16000)572 self.recorder.feed_audio(int16_data)573 574 except Exception as e:575 logger.error(f"Error processing audio chunk: {e}")576 577 def resample_audio(self, audio_bytes, from_rate, to_rate):578 """Resample audio to target sample rate"""579 try:580 audio_np = np.frombuffer(audio_bytes, dtype=np.int16)581 num_samples = len(audio_np)582 num_target_samples = int(num_samples * to_rate / from_rate)583 584 resampled = resample(audio_np, num_target_samples)585 586 return resampled.astype(np.int16).tobytes()587 except Exception as e:588 logger.error(f"Error resampling audio: {e}")589 return audio_bytes590 591 592# FastRTC Audio Handler593class DiarizationHandler(AsyncStreamHandler):594 def __init__(self, diarization_system):595 super().__init__()596 self.diarization_system = diarization_system597 self.audio_buffer = []598 self.buffer_size = BUFFER_SIZE599 600 def copy(self):601 """Return a fresh handler for each new stream connection"""602 return DiarizationHandler(self.diarization_system)603 604 async def emit(self):605 """Not used - we only receive audio"""606 return None607 608 async def receive(self, frame):609 """Receive audio data from FastRTC"""610 try:611 if not self.diarization_system.is_running:612 return613 614 # Extract audio data615 audio_data = getattr(frame, 'data', frame)616 617 # Convert to numpy array618 if isinstance(audio_data, bytes):619 audio_array = np.frombuffer(audio_data, dtype=np.int16).astype(np.float32) / 32768.0620 elif isinstance(audio_data, (list, tuple)):621 sample_rate, audio_array = audio_data622 if isinstance(audio_array, (list, tuple)):623 audio_array = np.array(audio_array, dtype=np.float32)624 else:625 audio_array = np.array(audio_data, dtype=np.float32)626 627 # Ensure 1D628 if len(audio_array.shape) > 1:629 audio_array = audio_array.flatten()630 631 # Buffer audio chunks632 self.audio_buffer.extend(audio_array)633 634 # Process in chunks635 while len(self.audio_buffer) >= self.buffer_size:636 chunk = np.array(self.audio_buffer[:self.buffer_size])637 self.audio_buffer = self.audio_buffer[self.buffer_size:]638 639 # Process asynchronously640 await self.process_audio_async(chunk)641 642 except Exception as e:643 logger.error(f"Error in FastRTC receive: {e}")644 645 async def process_audio_async(self, audio_data):646 """Process audio data asynchronously"""647 try:648 loop = asyncio.get_event_loop()649 await loop.run_in_executor(650 None, 651 self.diarization_system.process_audio_chunk, 652 audio_data, 653 SAMPLE_RATE654 )655 except Exception as e:656 logger.error(f"Error in async audio processing: {e}")657 658 659# Global instances660diarization_system = RealtimeSpeakerDiarization()661audio_handler = None662 663def initialize_system():664 """Initialize the diarization system"""665 global audio_handler666 try:667 success = diarization_system.initialize_models()668 if success:669 audio_handler = DiarizationHandler(diarization_system)670 return "✅ System initialized successfully!"671 else:672 return "❌ Failed to initialize system. Check logs for details."673 except Exception as e:674 logger.error(f"Initialization error: {e}")675 return f"❌ Initialization error: {str(e)}"676 677def start_recording():678 """Start recording and transcription"""679 try:680 result = diarization_system.start_recording()681 return f"🎙️ {result}"682 except Exception as e:683 return f"❌ Failed to start recording: {str(e)}"684 685def stop_recording():686 """Stop recording and transcription"""687 try:688 result = diarization_system.stop_recording()689 return f"⏹️ {result}"690 except Exception as e:691 return f"❌ Failed to stop recording: {str(e)}"692 693def clear_conversation():694 """Clear the conversation"""695 try:696 result = diarization_system.clear_conversation()697 return f"🗑️ {result}"698 except Exception as e:699 return f"❌ Failed to clear conversation: {str(e)}"700 701def update_settings(threshold, max_speakers):702 """Update system settings"""703 try:704 result = diarization_system.update_settings(threshold, max_speakers)705 return f"⚙️ {result}"706 except Exception as e:707 return f"❌ Failed to update settings: {str(e)}"708 709def get_conversation():710 """Get the current conversation"""711 try:712 return diarization_system.get_formatted_conversation()713 except Exception as e:714 return f"<i>Error getting conversation: {str(e)}</i>"715 716def get_status():717 """Get system status"""718 try:719 return diarization_system.get_status_info()720 except Exception as e:721 return f"Error getting status: {str(e)}"722 723# Create Gradio interface724def create_interface():725 with gr.Blocks(title="Real-time Speaker Diarization", theme=gr.themes.Soft()) as interface:726 gr.Markdown("# 🎤 Real-time Speech Recognition with Speaker Diarization")727 gr.Markdown("Live transcription with automatic speaker identification using FastRTC audio streaming.")728 729 with gr.Row():730 with gr.Column(scale=2):731 # Conversation display732 conversation_output = gr.HTML(733 value="<div style='padding: 20px; background: #f8f9fa; border-radius: 10px; min-height: 300px;'><i>Click 'Initialize System' to start...</i></div>",734 label="Live Conversation"735 )736 737 # Control buttons738 with gr.Row():739 init_btn = gr.Button("🔧 Initialize System", variant="secondary", size="lg")740 start_btn = gr.Button("🎙️ Start", variant="primary", size="lg", interactive=False)741 stop_btn = gr.Button("⏹️ Stop", variant="stop", size="lg", interactive=False)742 clear_btn = gr.Button("🗑️ Clear", variant="secondary", size="lg", interactive=False)743 744 # Status display745 status_output = gr.Textbox(746 label="System Status",747 value="Ready to initialize...",748 lines=8,749 interactive=False750 )751 752 with gr.Column(scale=1):753 # Settings754 gr.Markdown("## ⚙️ Settings")755 756 threshold_slider = gr.Slider(757 minimum=0.3,758 maximum=0.9,759 step=0.05,760 value=DEFAULT_CHANGE_THRESHOLD,761 label="Speaker Change Sensitivity",762 info="Lower = more sensitive"763 )764 765 max_speakers_slider = gr.Slider(766 minimum=2,767 maximum=ABSOLUTE_MAX_SPEAKERS,768 step=1,769 value=DEFAULT_MAX_SPEAKERS,770 label="Maximum Speakers"771 )772 773 update_btn = gr.Button("Update Settings", variant="secondary")774 775 # Instructions776 gr.Markdown("""777 ## 📋 Instructions778 1. **Initialize** the system (loads AI models)779 2. **Start** recording 780 3. **Speak** - system will transcribe and identify speakers781 4. **Monitor** real-time results below782 783 ## 🎨 Speaker Colors784 - 🔴 Speaker 1 (Red)785 - 🟢 Speaker 2 (Teal) 786 - 🔵 Speaker 3 (Blue)787 - 🟡 Speaker 4 (Green)788 - 🟣 Speaker 5 (Yellow)789 - 🟤 Speaker 6 (Plum)790 - 🟫 Speaker 7 (Mint)791 - 🟨 Speaker 8 (Gold)792 """)793 794 # Event handlers795 def on_initialize():796 result = initialize_system()797 if "✅" in result:798 return result, gr.update(interactive=True), gr.update(interactive=True), gr.update(interactive=True)799 else:800 return result, gr.update(interactive=False), gr.update(interactive=False), gr.update(interactive=False)801 802 def on_start():803 result = start_recording()804 return result, gr.update(interactive=False), gr.update(interactive=True)805 806 def on_stop():807 result = stop_recording()808 return result, gr.update(interactive=True), gr.update(interactive=False)809 810 def on_clear():811 result = clear_conversation()812 return result813 814 def on_update_settings(threshold, max_speakers):815 result = update_settings(threshold, int(max_speakers))816 return result817 818 def refresh_conversation():819 return get_conversation()820 821 def refresh_status():822 return get_status()823 824 # Button click handlers825 init_btn.click(826 fn=on_initialize,827 outputs=[status_output, start_btn, stop_btn, clear_btn]828 )829 830 start_btn.click(831 fn=on_start,832 outputs=[status_output, start_btn, stop_btn]833 )834 835 stop_btn.click(836 fn=on_stop,837 outputs=[status_output, start_btn, stop_btn]838 )839 840 clear_btn.click(841 fn=on_clear,842 outputs=[status_output]843 )844 845 update_btn.click(846 fn=on_update_settings,847 inputs=[threshold_slider, max_speakers_slider],848 outputs=[status_output]849 )850 851 # Auto-refresh conversation display every 1 second852 conversation_timer = gr.Timer(1)853 conversation_timer.tick(refresh_conversation, outputs=[conversation_output])854 855 # Auto-refresh status every 2 seconds 856 status_timer = gr.Timer(2)857 status_timer.tick(refresh_status, outputs=[status_output])858 859 return interface860 861 862# FastAPI setup for FastRTC integration863app = FastAPI()864 865@app.get("/")866async def root():867 return {"message": "Real-time Speaker Diarization API"}868 869@app.get("/health")870async def health_check():871 return {"status": "healthy", "system_running": diarization_system.is_running}872 873@app.post("/initialize")874async def api_initialize():875 result = initialize_system()876 return {"result": result, "success": "✅" in result}877 878@app.post("/start")879async def api_start():880 result = start_recording()881 return {"result": result, "success": "🎙️" in result}882 883@app.post("/stop")884async def api_stop():885 result = stop_recording()886 return {"result": result, "success": "⏹️" in result}887 888@app.post("/clear")889async def api_clear():890 result = clear_conversation()891 return {"result": result}892 893@app.get("/conversation")894async def api_get_conversation():895 return {"conversation": get_conversation()}896 897@app.get("/status")898async def api_get_status():899 return {"status": get_status()}900 901@app.post("/settings")902async def api_update_settings(threshold: float, max_speakers: int):903 result = update_settings(threshold, max_speakers)904 return {"result": result}905 906# FastRTC Stream setup907if audio_handler:908 stream = Stream(handler=audio_handler)909 app.include_router(stream.router, prefix="/stream")910 911 912# Main execution913if __name__ == "__main__":914 import argparse915 916 parser = argparse.ArgumentParser(description="Real-time Speaker Diarization System")917 parser.add_argument("--mode", choices=["gradio", "api", "both"], default="gradio", 918 help="Run mode: gradio interface, API only, or both")919 parser.add_argument("--host", default="0.0.0.0", help="Host to bind to")920 parser.add_argument("--port", type=int, default=7860, help="Port to bind to")921 parser.add_argument("--api-port", type=int, default=8000, help="API port (when running both)")922 923 args = parser.parse_args()924 925 if args.mode == "gradio":926 # Run Gradio interface only927 interface = create_interface()928 interface.launch(929 server_name=args.host,930 server_port=args.port,931 share=True,932 show_error=True933 )934 935 elif args.mode == "api":936 # Run FastAPI only937 uvicorn.run(938 app, 939 host=args.host, 940 port=args.port,941 log_level="info"942 )943 944 elif args.mode == "both":945 # Run both Gradio and FastAPI946 import multiprocessing947 import threading948 949 def run_gradio():950 interface = create_interface()951 interface.launch(952 server_name=args.host,953 server_port=args.port,954 share=True,955 show_error=True956 )957 958 def run_fastapi():959 uvicorn.run(960 app,961 host=args.host,962 port=args.api_port,963 log_level="info"964 )965 966 # Start FastAPI in a separate thread967 api_thread = threading.Thread(target=run_fastapi, daemon=True)968 api_thread.start()969 970 # Start Gradio in main thread971 run_gradio()