Changed gradio approx real-stream to FastRTC

app.py

@@ -8,6 +8,7 @@ import os
 import urllib.request
 import torchaudio
 from scipy.spatial.distance import cosine
+from RealtimeSTT import AudioToTextRecorder
 import json
 import io
 import wave
@@ -57,9 +58,6 @@ SPEAKER_COLOR_NAMES = [
 ]
 
 
-
-
-
 class SpeechBrainEncoder:
     """ECAPA-TDNN encoder from SpeechBrain for speaker embeddings"""
     def __init__(self, device="cpu"):
@@ -70,11 +68,24 @@ class SpeechBrainEncoder:
         self.cache_dir = os.path.join(os.path.expanduser("~"), ".cache", "speechbrain")
         os.makedirs(self.cache_dir, exist_ok=True)
     
+    def _download_model(self):
+        """Download pre-trained SpeechBrain ECAPA-TDNN model if not present"""
+        model_url = "https://huggingface.co/speechbrain/spkrec-ecapa-voxceleb/resolve/main/embedding_model.ckpt"
+        model_path = os.path.join(self.cache_dir, "embedding_model.ckpt")
+        
+        if not os.path.exists(model_path):
+            print(f"Downloading ECAPA-TDNN model to {model_path}...")
+            urllib.request.urlretrieve(model_url, model_path)
+        
+        return model_path
+    
     def load_model(self):
         """Load the ECAPA-TDNN model"""
         try:
             from speechbrain.pretrained import EncoderClassifier
             
+            model_path = self._download_model()
+            
             self.model = EncoderClassifier.from_hparams(
                 source="speechbrain/spkrec-ecapa-voxceleb",
                 savedir=self.cache_dir,
@@ -82,10 +93,9 @@ class SpeechBrainEncoder:
             )
             
             self.model_loaded = True
-            print("ECAPA-TDNN model loaded successfully!")
             return True
         except Exception as e:
-            print(f"SpeechBrain not available: {e}")
+            print(f"Error loading ECAPA-TDNN model: {e}")
             return False
     
     def embed_utterance(self, audio, sr=16000):
@@ -116,21 +126,16 @@ class AudioProcessor:
     def __init__(self, encoder):
         self.encoder = encoder
     
-    def extract_embedding(self, audio_data, sample_rate=16000):
+    def extract_embedding(self, audio_int16):
         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)
+            float_audio = audio_int16.astype(np.float32) / 32768.0
             
-            # 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
+            embedding = self.encoder.embed_utterance(float_audio)
             
+            return embedding
         except Exception as e:
             print(f"Embedding extraction error: {e}")
             return np.zeros(self.encoder.embedding_dim)
@@ -266,14 +271,68 @@ class SpeakerChangeDetector:
         }
 
 
-class GradioSpeakerDiarization:
+class WebRTCAudioProcessor:
+    """Processes WebRTC audio streams for speaker diarization"""
+    def __init__(self, diarization_system):
+        self.diarization_system = diarization_system
+        self.audio_buffer = []
+        self.buffer_lock = threading.Lock()
+        self.processing_thread = None
+        self.is_processing = False
+        
+    def process_audio(self, audio_data, sample_rate):
+        """Process incoming audio data from WebRTC"""
+        try:
+            # Convert audio data to numpy array if needed
+            if isinstance(audio_data, bytes):
+                audio_array = np.frombuffer(audio_data, dtype=np.int16)
+            elif isinstance(audio_data, tuple):
+                # Handle tuple format (sample_rate, audio_array)
+                sample_rate, audio_array = audio_data
+                if isinstance(audio_array, np.ndarray):
+                    if audio_array.dtype != np.int16:
+                        audio_array = (audio_array * 32767).astype(np.int16)
+                else:
+                    audio_array = np.array(audio_array, dtype=np.int16)
+            else:
+                audio_array = np.array(audio_data, dtype=np.int16)
+            
+            # Ensure mono audio
+            if len(audio_array.shape) > 1:
+                audio_array = audio_array[:, 0]
+            
+            # Add to buffer
+            with self.buffer_lock:
+                self.audio_buffer.extend(audio_array)
+                
+                # Process buffer when it's large enough (1 second of audio)
+                if len(self.audio_buffer) >= sample_rate:
+                    buffer_to_process = np.array(self.audio_buffer[:sample_rate])
+                    self.audio_buffer = self.audio_buffer[sample_rate//2:]  # Keep 50% overlap
+                    
+                    # Feed to recorder in separate thread
+                    if self.diarization_system.recorder:
+                        audio_bytes = buffer_to_process.tobytes()
+                        self.diarization_system.recorder.feed_audio(audio_bytes)
+                        
+        except Exception as e:
+            print(f"Error processing WebRTC audio: {e}")
+
+
+class RealtimeSpeakerDiarization:
     def __init__(self):
         self.encoder = None
         self.audio_processor = None
         self.speaker_detector = None
+        self.recorder = None
+        self.webrtc_processor = None
+        self.sentence_queue = queue.Queue()
         self.full_sentences = []
         self.sentence_speakers = []
-        self.is_initialized = False
+        self.pending_sentences = []
+        self.displayed_text = ""
+        self.last_realtime_text = ""
+        self.is_running = False
         self.change_threshold = DEFAULT_CHANGE_THRESHOLD
         self.max_speakers = DEFAULT_MAX_SPEAKERS
         
@@ -283,7 +342,6 @@ class GradioSpeakerDiarization:
             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()
             
@@ -294,62 +352,145 @@ class GradioSpeakerDiarization:
                     change_threshold=self.change_threshold,
                     max_speakers=self.max_speakers
                 )
-                self.is_initialized = True
+                self.webrtc_processor = WebRTCAudioProcessor(self)
+                print("ECAPA-TDNN model loaded successfully!")
                 return True
             else:
+                print("Failed to load ECAPA-TDNN model")
                 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()
+    def live_text_detected(self, text):
+        """Callback for real-time transcription updates"""
+        text = text.strip()
+        if text:
+            sentence_delimiters = '.?!。'
+            prob_sentence_end = (
+                len(self.last_realtime_text) > 0
+                and text[-1] in sentence_delimiters
+                and self.last_realtime_text[-1] in sentence_delimiters
+            )
+
+            self.last_realtime_text = text
+
+            if prob_sentence_end and FAST_SENTENCE_END:
+                self.recorder.stop()
+            elif prob_sentence_end:
+                self.recorder.post_speech_silence_duration = SILENCE_THRESHS[0]
+            else:
+                self.recorder.post_speech_silence_duration = SILENCE_THRESHS[1]
+    
+    def process_final_text(self, text):
+        """Process final transcribed text with speaker embedding"""
+        text = text.strip()
+        if text:
+            try:
+                bytes_data = self.recorder.last_transcription_bytes
+                self.sentence_queue.put((text, bytes_data))
+                self.pending_sentences.append(text)
+            except Exception as e:
+                print(f"Error processing final text: {e}")
+    
+    def process_sentence_queue(self):
+        """Process sentences in the queue for speaker detection"""
+        while self.is_running:
+            try:
+                text, bytes_data = self.sentence_queue.get(timeout=1)
+                
+                # Convert audio data to int16
+                audio_int16 = np.int16(bytes_data * 32767)
+                
+                # Extract speaker embedding
+                speaker_embedding = self.audio_processor.extract_embedding(audio_int16)
+                
+                # Store sentence and embedding
+                self.full_sentences.append((text, speaker_embedding))
+                
+                # Fill in missing speaker assignments
+                while len(self.sentence_speakers) < len(self.full_sentences) - 1:
+                    self.sentence_speakers.append(0)
+                
+                # Detect speaker changes
+                speaker_id, similarity = self.speaker_detector.add_embedding(speaker_embedding)
+                self.sentence_speakers.append(speaker_id)
+                
+                # Remove from pending
+                if text in self.pending_sentences:
+                    self.pending_sentences.remove(text)
+                    
+            except queue.Empty:
+                continue
+            except Exception as e:
+                print(f"Error processing sentence: {e}")
+    
+    def start_recording(self):
+        """Start the recording and transcription process"""
+        if self.encoder is None:
+            return "Please initialize models first!"
         
         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))
+            # Setup recorder configuration for WebRTC input
+            recorder_config = {
+                'spinner': False,
+                'use_microphone': False,  # We'll feed audio manually
+                'model': FINAL_TRANSCRIPTION_MODEL,
+                'language': TRANSCRIPTION_LANGUAGE,
+                'silero_sensitivity': SILERO_SENSITIVITY,
+                'webrtc_sensitivity': WEBRTC_SENSITIVITY,
+                'post_speech_silence_duration': SILENCE_THRESHS[1],
+                'min_length_of_recording': MIN_LENGTH_OF_RECORDING,
+                'pre_recording_buffer_duration': PRE_RECORDING_BUFFER_DURATION,
+                'min_gap_between_recordings': 0,
+                'enable_realtime_transcription': True,
+                'realtime_processing_pause': 0,
+                'realtime_model_type': REALTIME_TRANSCRIPTION_MODEL,
+                'on_realtime_transcription_update': self.live_text_detected,
+                'beam_size': FINAL_BEAM_SIZE,
+                'beam_size_realtime': REALTIME_BEAM_SIZE,
+                'buffer_size': BUFFER_SIZE,
+                'sample_rate': SAMPLE_RATE,
+            }
+
+            self.recorder = AudioToTextRecorder(**recorder_config)
             
-            # Detect speaker changes
-            speaker_id, similarity = self.speaker_detector.add_embedding(speaker_embedding)
-            self.sentence_speakers.append(speaker_id)
+            # Start sentence processing thread
+            self.is_running = True
+            self.sentence_thread = threading.Thread(target=self.process_sentence_queue, daemon=True)
+            self.sentence_thread.start()
             
-            status_msg = f"✅ Processed audio segment. Detected as Speaker {speaker_id + 1} (similarity: {similarity:.3f})"
+            # Start transcription thread
+            self.transcription_thread = threading.Thread(target=self.run_transcription, daemon=True)
+            self.transcription_thread.start()
             
-            return status_msg, self.get_formatted_conversation(), self.get_status_info()
+            return "Recording started successfully! WebRTC audio input ready."
             
         except Exception as e:
-            error_msg = f"❌ Error processing audio: {str(e)}"
-            return error_msg, self.get_formatted_conversation(), self.get_status_info()
+            return f"Error starting recording: {e}"
+    
+    def run_transcription(self):
+        """Run the transcription loop"""
+        try:
+            while self.is_running:
+                self.recorder.text(self.process_final_text)
+        except Exception as e:
+            print(f"Transcription error: {e}")
+    
+    def stop_recording(self):
+        """Stop the recording process"""
+        self.is_running = False
+        if self.recorder:
+            self.recorder.stop()
+        return "Recording stopped!"
     
     def clear_conversation(self):
         """Clear all conversation data"""
         self.full_sentences = []
         self.sentence_speakers = []
+        self.pending_sentences = []
+        self.displayed_text = ""
+        self.last_realtime_text = ""
         
         if self.speaker_detector:
             self.speaker_detector = SpeakerChangeDetector(
@@ -358,7 +499,7 @@ class GradioSpeakerDiarization:
                 max_speakers=self.max_speakers
             )
         
-        return "Conversation cleared!", self.get_formatted_conversation(), self.get_status_info()
+        return "Conversation cleared!"
     
     def update_settings(self, threshold, max_speakers):
         """Update speaker detection settings"""
@@ -369,22 +510,18 @@ class GradioSpeakerDiarization:
             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()
+        return f"Settings updated: Threshold={threshold:.2f}, Max Speakers={max_speakers}"
     
     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 = []
             
+            # Process completed sentences
             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)
@@ -393,7 +530,15 @@ class GradioSpeakerDiarization:
                 sentences_with_style.append(
                     f'<span style="color:{color};"><b>{speaker_name}:</b> {sentence_text}</span>')
             
-            return "<br><br>".join(sentences_with_style)
+            # Add pending sentences
+            for pending_sentence in self.pending_sentences:
+                sentences_with_style.append(
+                    f'<span style="color:#60FFFF;"><b>Processing:</b> {pending_sentence}</span>')
+            
+            if sentences_with_style:
+                return "<br><br>".join(sentences_with_style)
+            else:
+                return "Waiting for speech input..."
                 
         except Exception as e:
             return f"Error formatting conversation: {e}"
@@ -411,7 +556,7 @@ class GradioSpeakerDiarization:
                 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)}",
+                f"**Total Sentences:** {len(self.full_sentences)}",
                 "",
                 "**Speaker Segment Counts:**"
             ]
@@ -427,21 +572,26 @@ class GradioSpeakerDiarization:
 
 
 # Global instance
-diarization_system = GradioSpeakerDiarization()
+diarization_system = RealtimeSpeakerDiarization()
 
 
 def initialize_system():
     """Initialize the diarization system"""
     success = diarization_system.initialize_models()
     if success:
-        return "✅ System initialized successfully! Models loaded.", "", ""
+        return "✅ System initialized successfully! Models loaded."
     else:
-        return "❌ Failed to initialize system. Please check the logs.", "", ""
+        return "❌ Failed to initialize system. Please check the logs."
+
+
+def start_recording():
+    """Start recording and transcription"""
+    return diarization_system.start_recording()
 
 
-def process_audio(audio):
-    """Process uploaded or recorded audio"""
-    return diarization_system.transcribe_audio(audio)
+def stop_recording():
+    """Stop recording and transcription"""
+    return diarization_system.stop_recording()
 
 
 def clear_conversation():
@@ -454,52 +604,59 @@ def update_settings(threshold, max_speakers):
     return diarization_system.update_settings(threshold, max_speakers)
 
 
+def get_conversation():
+    """Get the current conversation"""
+    return diarization_system.get_formatted_conversation()
+
+
+def get_status():
+    """Get system status"""
+    return diarization_system.get_status_info()
+
+
+def process_audio_stream(audio):
+    """Process audio stream from WebRTC"""
+    if diarization_system.webrtc_processor and diarization_system.is_running:
+        diarization_system.webrtc_processor.process_audio(audio, SAMPLE_RATE)
+    return None
+
+
 # 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.Blocks(title="Real-time Speaker Diarization", theme=gr.themes.Dark()) as app:
+        gr.Markdown("# 🎤 Real-time Speech Recognition with Speaker Diarization")
+        gr.Markdown("This app performs real-time speech recognition with automatic speaker identification and color-coding using WebRTC.")
         
         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
+                # WebRTC Audio Input
+                audio_input = gr.Audio(
+                    sources=["microphone"],
+                    streaming=True,
+                    label="🎙️ Microphone Input",
+                    type="numpy"
                 )
                 
+                # Main conversation display
                 conversation_output = gr.HTML(
-                    value="<i>System not initialized...</i>",
-                    label="Speaker Analysis Results"
+                    value="<i>Click 'Initialize System' to start...</i>",
+                    label="Live Conversation"
                 )
                 
                 # Control buttons
                 with gr.Row():
-                    clear_btn = gr.Button("🗑️ Clear Results", variant="stop")
+                    init_btn = gr.Button("🔧 Initialize System", variant="secondary")
+                    start_btn = gr.Button("🎙️ Start Recording", variant="primary", interactive=False)
+                    stop_btn = gr.Button("⏹️ Stop Recording", variant="stop", interactive=False)
+                    clear_btn = gr.Button("🗑️ Clear Conversation", interactive=False)
+                
+                # Status display
+                status_output = gr.Textbox(
+                    label="System Status",
+                    value="System not initialized",
+                    lines=8,
+                    interactive=False
+                )
             
             with gr.Column(scale=1):
                 # Settings panel
@@ -511,7 +668,7 @@ def create_interface():
                     step=0.05,
                     value=DEFAULT_CHANGE_THRESHOLD,
                     label="Speaker Change Sensitivity",
-                    info="Lower = more sensitive to speaker changes"
+                    info="Lower values = more sensitive to speaker changes"
                 )
                 
                 max_speakers_slider = gr.Slider(
@@ -522,51 +679,106 @@ def create_interface():
                     label="Maximum Number of Speakers"
                 )
                 
-                update_settings_btn = gr.Button("Update Settings", variant="secondary")
+                update_settings_btn = gr.Button("Update Settings")
                 
-                # System status
-                system_status = gr.Textbox(
-                    label="System Status",
-                    value="System not initialized",
-                    lines=12,
-                    interactive=False
-                )
+                # Instructions
+                gr.Markdown("## 📝 Instructions")
+                gr.Markdown("""
+                1. Click **Initialize System** to load models
+                2. Click **Start Recording** to begin processing
+                3. Allow microphone access when prompted
+                4. Speak into your microphone
+                5. Watch real-time transcription with speaker labels
+                6. Adjust settings as needed
+                """)
                 
                 # 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})')
+                for i, (color, name) in enumerate(zip(SPEAKER_COLORS, SPEAKER_COLOR_NAMES)):
+                    color_info.append(f'<span style="color:{color};">■</span> Speaker {i+1} ({name})')
                 
-                gr.HTML("<br>".join(color_info))
+                gr.HTML("<br>".join(color_info[:DEFAULT_MAX_SPEAKERS]))
+        
+        # Auto-refresh conversation and status
+        def refresh_display():
+            return get_conversation(), get_status()
         
         # Event handlers
+        def on_initialize():
+            result = initialize_system()
+            if "successfully" in result:
+                return (
+                    result,
+                    gr.update(interactive=True),   # start_btn
+                    gr.update(interactive=True),   # clear_btn
+                    get_conversation(),
+                    get_status()
+                )
+            else:
+                return (
+                    result,
+                    gr.update(interactive=False),  # start_btn
+                    gr.update(interactive=False),  # clear_btn
+                    get_conversation(),
+                    get_status()
+                )
+        
+        def on_start():
+            result = start_recording()
+            return (
+                result,
+                gr.update(interactive=False),  # start_btn
+                gr.update(interactive=True),   # stop_btn
+            )
+        
+        def on_stop():
+            result = stop_recording()
+            return (
+                result,
+                gr.update(interactive=True),   # start_btn
+                gr.update(interactive=False),  # stop_btn
+            )
+        
+        # Connect event handlers
         init_btn.click(
-            initialize_system,
-            outputs=[status_output, conversation_output, system_status]
+            on_initialize,
+            outputs=[status_output, start_btn, clear_btn, conversation_output, status_output]
         )
         
-        process_file_btn.click(
-            process_audio,
-            inputs=[audio_file],
-            outputs=[status_output, conversation_output, system_status]
+        start_btn.click(
+            on_start,
+            outputs=[status_output, start_btn, stop_btn]
         )
         
-        process_mic_btn.click(
-            process_audio,
-            inputs=[audio_mic],
-            outputs=[status_output, conversation_output, system_status]
+        stop_btn.click(
+            on_stop,
+            outputs=[status_output, start_btn, stop_btn]
         )
         
         clear_btn.click(
             clear_conversation,
-            outputs=[status_output, conversation_output, system_status]
+            outputs=[status_output]
         )
         
         update_settings_btn.click(
             update_settings,
             inputs=[threshold_slider, max_speakers_slider],
-            outputs=[status_output, conversation_output, system_status]
+            outputs=[status_output]
+        )
+        
+        # Connect WebRTC audio stream to processing
+        audio_input.stream(
+            process_audio_stream,
+            inputs=[audio_input],
+            outputs=[]
+        )
+        
+        # Auto-refresh every 2 seconds when recording
+        refresh_timer = gr.Timer(2.0)
+        refresh_timer.tick(
+            refresh_display,
+            outputs=[conversation_output, status_output]
         )
     
     return app
@@ -576,5 +788,6 @@ if __name__ == "__main__":
     app = create_interface()
     app.launch(
         server_name="0.0.0.0",
-        server_port=7860
+        server_port=7860,
+        share=True
     )