Code Update

app.py

@@ -0,0 +1,453 @@
+import gradio as gr
+import numpy as np
+import torch
+import torchaudio
+from scipy.spatial.distance import cosine
+import tempfile
+import os
+import warnings
+warnings.filterwarnings("ignore", category=UserWarning)
+
+try:
+    from transformers import pipeline
+except ImportError:
+    print("transformers not found. Install with: pip install transformers")
+
+# Configuration
+class Config:
+    # Audio settings
+    SAMPLE_RATE = 16000
+    
+    # Speaker detection
+    CHANGE_THRESHOLD = 0.65
+    MAX_SPEAKERS = 4
+    MIN_SEGMENT_DURATION = 1.0
+    EMBEDDING_HISTORY_SIZE = 3
+    SPEAKER_MEMORY_SIZE = 20
+
+# Console colors for speakers (HTML version)
+SPEAKER_COLORS = [
+    "#FFD700",  # Gold
+    "#FF6B6B",  # Red
+    "#4ECDC4",  # Teal
+    "#45B7D1",  # Blue
+    "#96CEB4",  # Mint
+    "#FFEAA7",  # Light Yellow
+    "#DDA0DD",  # Plum
+    "#98D8C8",  # Mint Green
+]
+
+class SpeakerEncoder:
+    """Simplified speaker encoder using torchaudio transforms"""
+    
+    def __init__(self, device="cpu"):
+        self.device = device
+        self.embedding_dim = 128
+        self.model_loaded = False
+        self._setup_model()
+    
+    def _setup_model(self):
+        """Setup a simple MFCC-based feature extractor"""
+        try:
+            self.mfcc_transform = torchaudio.transforms.MFCC(
+                sample_rate=Config.SAMPLE_RATE,
+                n_mfcc=13,
+                melkwargs={"n_fft": 400, "hop_length": 160, "n_mels": 23}
+            ).to(self.device)
+            self.model_loaded = True
+            print("Simple MFCC-based encoder initialized")
+        except Exception as e:
+            print(f"Error setting up encoder: {e}")
+            self.model_loaded = False
+    
+    def extract_embedding(self, audio):
+        """Extract speaker embedding from audio"""
+        if not self.model_loaded:
+            return np.zeros(self.embedding_dim)
+        
+        try:
+            # Ensure audio is float32 and normalized
+            if isinstance(audio, np.ndarray):
+                audio = torch.from_numpy(audio).float()
+            
+            # Normalize audio
+            if audio.abs().max() > 0:
+                audio = audio / audio.abs().max()
+            
+            # Add batch dimension if needed
+            if audio.dim() == 1:
+                audio = audio.unsqueeze(0)
+            
+            # Extract MFCC features
+            with torch.no_grad():
+                mfcc = self.mfcc_transform(audio)
+                # Simple statistics-based embedding
+                embedding = torch.cat([
+                    mfcc.mean(dim=2).flatten(),
+                    mfcc.std(dim=2).flatten(),
+                    mfcc.max(dim=2)[0].flatten(),
+                    mfcc.min(dim=2)[0].flatten()
+                ])
+                
+                # Pad or truncate to fixed size
+                if embedding.size(0) > self.embedding_dim:
+                    embedding = embedding[:self.embedding_dim]
+                elif embedding.size(0) < self.embedding_dim:
+                    padding = torch.zeros(self.embedding_dim - embedding.size(0))
+                    embedding = torch.cat([embedding, padding])
+            
+            return embedding.cpu().numpy()
+            
+        except Exception as e:
+            print(f"Error extracting embedding: {e}")
+            return np.zeros(self.embedding_dim)
+
+class SpeakerDetector:
+    """Speaker change detection using embeddings"""
+    
+    def __init__(self, threshold=Config.CHANGE_THRESHOLD, max_speakers=Config.MAX_SPEAKERS):
+        self.threshold = threshold
+        self.max_speakers = max_speakers
+        self.current_speaker = 0
+        self.speaker_embeddings = [[] for _ in range(max_speakers)]
+        self.speaker_centroids = [None] * max_speakers
+        self.active_speakers = {0}
+        
+    def reset(self):
+        """Reset speaker detection state"""
+        self.current_speaker = 0
+        self.speaker_embeddings = [[] for _ in range(self.max_speakers)]
+        self.speaker_centroids = [None] * self.max_speakers
+        self.active_speakers = {0}
+        
+    def detect_speaker(self, embedding):
+        """Detect current speaker from embedding"""
+        # Initialize first speaker
+        if not self.speaker_embeddings[0]:
+            self.speaker_embeddings[0].append(embedding)
+            self.speaker_centroids[0] = embedding.copy()
+            return 0, 1.0
+        
+        # Calculate similarity with current speaker
+        current_centroid = self.speaker_centroids[self.current_speaker]
+        if current_centroid is not None:
+            similarity = 1.0 - cosine(embedding, current_centroid)
+        else:
+            similarity = 0.0
+        
+        # Check for speaker change
+        if similarity < self.threshold:
+            # Find best matching existing speaker
+            best_speaker = self.current_speaker
+            best_similarity = similarity
+            
+            for speaker_id in self.active_speakers:
+                if speaker_id == self.current_speaker:
+                    continue
+                    
+                centroid = self.speaker_centroids[speaker_id]
+                if centroid is not None:
+                    sim = 1.0 - cosine(embedding, centroid)
+                    if sim > best_similarity and sim > self.threshold:
+                        best_similarity = sim
+                        best_speaker = speaker_id
+            
+            # Create new speaker if no good match and slots available
+            if (best_speaker == self.current_speaker and 
+                len(self.active_speakers) < self.max_speakers):
+                for new_id in range(self.max_speakers):
+                    if new_id not in self.active_speakers:
+                        best_speaker = new_id
+                        best_similarity = 0.0
+                        self.active_speakers.add(new_id)
+                        break
+            
+            # Update current speaker if changed
+            if best_speaker != self.current_speaker:
+                self.current_speaker = best_speaker
+                similarity = best_similarity
+        
+        # Update speaker model
+        self._update_speaker_model(self.current_speaker, embedding)
+        return self.current_speaker, similarity
+    
+    def _update_speaker_model(self, speaker_id, embedding):
+        """Update speaker model with new embedding"""
+        self.speaker_embeddings[speaker_id].append(embedding)
+        
+        # Keep only recent embeddings
+        if len(self.speaker_embeddings[speaker_id]) > Config.SPEAKER_MEMORY_SIZE:
+            self.speaker_embeddings[speaker_id] = \
+                self.speaker_embeddings[speaker_id][-Config.SPEAKER_MEMORY_SIZE:]
+        
+        # Update centroid
+        if self.speaker_embeddings[speaker_id]:
+            self.speaker_centroids[speaker_id] = np.mean(
+                self.speaker_embeddings[speaker_id], axis=0
+            )
+
+class AudioProcessor:
+    """Handles audio processing and transcription"""
+    
+    def __init__(self):
+        self.encoder = SpeakerEncoder()
+        self.detector = SpeakerDetector()
+        
+        # Initialize Whisper model for transcription
+        try:
+            self.transcriber = pipeline(
+                "automatic-speech-recognition",
+                model="openai/whisper-base",
+                chunk_length_s=30,
+                device=0 if torch.cuda.is_available() else -1
+            )
+            print("Whisper model loaded successfully")
+        except Exception as e:
+            print(f"Error loading Whisper model: {e}")
+            self.transcriber = None
+    
+    def process_audio_file(self, audio_file):
+        """Process uploaded audio file"""
+        if audio_file is None:
+            return "Please upload an audio file.", ""
+        
+        try:
+            # Reset speaker detection for new file
+            self.detector.reset()
+            
+            # Load audio file
+            waveform, sample_rate = torchaudio.load(audio_file)
+            
+            # Convert to mono if stereo
+            if waveform.shape[0] > 1:
+                waveform = waveform.mean(dim=0, keepdim=True)
+            
+            # Resample to 16kHz if needed
+            if sample_rate != Config.SAMPLE_RATE:
+                resampler = torchaudio.transforms.Resample(sample_rate, Config.SAMPLE_RATE)
+                waveform = resampler(waveform)
+            
+            # Convert to numpy
+            audio_data = waveform.squeeze().numpy()
+            
+            # Transcribe entire audio
+            if self.transcriber:
+                transcription_result = self.transcriber(audio_file)
+                full_transcription = transcription_result['text']
+            else:
+                full_transcription = "Transcription service unavailable"
+            
+            # Process audio in chunks for speaker detection
+            chunk_duration = 3.0  # 3 second chunks
+            chunk_samples = int(chunk_duration * Config.SAMPLE_RATE)
+            results = []
+            
+            for i in range(0, len(audio_data), chunk_samples // 2):  # 50% overlap
+                chunk = audio_data[i:i + chunk_samples]
+                
+                if len(chunk) < Config.SAMPLE_RATE:  # Skip chunks less than 1 second
+                    continue
+                
+                # Extract speaker embedding
+                embedding = self.encoder.extract_embedding(chunk)
+                speaker_id, similarity = self.detector.detect_speaker(embedding)
+                
+                # Get timestamp
+                start_time = i / Config.SAMPLE_RATE
+                end_time = (i + len(chunk)) / Config.SAMPLE_RATE
+                
+                # Transcribe chunk
+                if self.transcriber and len(chunk) > Config.SAMPLE_RATE:
+                    # Save chunk temporarily for transcription
+                    with tempfile.NamedTemporaryFile(suffix='.wav', delete=False) as tmp_file:
+                        torchaudio.save(tmp_file.name, torch.tensor(chunk).unsqueeze(0), Config.SAMPLE_RATE)
+                        chunk_result = self.transcriber(tmp_file.name)
+                        chunk_text = chunk_result['text'].strip()
+                        os.unlink(tmp_file.name)  # Clean up temp file
+                else:
+                    chunk_text = ""
+                
+                if chunk_text:  # Only add if there's actual text
+                    results.append({
+                        'speaker_id': speaker_id,
+                        'start_time': start_time,
+                        'end_time': end_time,
+                        'text': chunk_text,
+                        'similarity': similarity
+                    })
+            
+            # Format results
+            formatted_output = self._format_results(results)
+            return formatted_output, full_transcription
+            
+        except Exception as e:
+            return f"Error processing audio: {str(e)}", ""
+    
+    def _format_results(self, results):
+        """Format results with speaker colors"""
+        if not results:
+            return "No speech detected in the audio file."
+        
+        formatted_lines = []
+        formatted_lines.append("🎤 **Speaker Diarization Results**\n")
+        
+        for result in results:
+            speaker_id = result['speaker_id']
+            start_time = result['start_time']
+            end_time = result['end_time']
+            text = result['text']
+            similarity = result['similarity']
+            
+            color = SPEAKER_COLORS[speaker_id % len(SPEAKER_COLORS)]
+            
+            # Format timestamp
+            start_min, start_sec = divmod(int(start_time), 60)
+            end_min, end_sec = divmod(int(end_time), 60)
+            timestamp = f"[{start_min:02d}:{start_sec:02d} - {end_min:02d}:{end_sec:02d}]"
+            
+            # Create colored HTML output
+            formatted_lines.append(
+                f'<div style="margin-bottom: 10px; padding: 8px; border-left: 4px solid {color}; background-color: {color}20;">'
+                f'<strong style="color: {color};">Speaker {speaker_id + 1}</strong> '
+                f'<span style="color: #666; font-size: 0.9em;">{timestamp}</span><br>'
+                f'<span style="color: #333;">{text}</span>'
+                f'</div>'
+            )
+        
+        return "".join(formatted_lines)
+
+# Global processor instance
+processor = AudioProcessor()
+
+def process_audio(audio_file, sensitivity):
+    """Process audio file with speaker detection"""
+    if audio_file is None:
+        return "Please upload an audio file.", ""
+    
+    # Update sensitivity
+    processor.detector.threshold = sensitivity
+    
+    # Process the audio
+    diarized_output, full_transcription = processor.process_audio_file(audio_file)
+    
+    return diarized_output, full_transcription
+
+# Create Gradio interface
+def create_interface():
+    """Create Gradio interface"""
+    
+    with gr.Blocks(
+        theme=gr.themes.Soft(),
+        title="Speaker Diarization & Transcription",
+        css="""
+        .gradio-container {
+            max-width: 1200px !important;
+        }
+        .speaker-output {
+            font-family: 'Segoe UI', Tahoma, Geneva, Verdana, sans-serif;
+        }
+        """
+    ) as demo:
+        
+        gr.Markdown(
+            """
+            # 🎙️ Speaker Diarization & Transcription
+            
+            Upload an audio file to automatically detect different speakers and transcribe their speech.
+            The system will identify speaker changes and display each speaker's text in different colors.
+            """
+        )
+        
+        with gr.Row():
+            with gr.Column(scale=1):
+                audio_input = gr.Audio(
+                    label="Upload Audio File",
+                    type="filepath",
+                    sources=["upload", "microphone"]
+                )
+                
+                sensitivity_slider = gr.Slider(
+                    minimum=0.1,
+                    maximum=1.0,
+                    value=0.65,
+                    step=0.05,
+                    label="Speaker Change Sensitivity",
+                    info="Lower values = more sensitive to speaker changes"
+                )
+                
+                process_btn = gr.Button("🎯 Process Audio", variant="primary", size="lg")
+                
+                gr.Markdown(
+                    """
+                    ### Instructions:
+                    1. Upload an audio file (WAV, MP3, etc.)
+                    2. Adjust sensitivity if needed
+                    3. Click "Process Audio"
+                    4. View results with speaker colors
+                    
+                    ### Tips:
+                    - Works best with clear speech
+                    - Supports multiple file formats
+                    - Different speakers shown in different colors
+                    - Processing may take a moment for longer files
+                    """
+                )
+            
+            with gr.Column(scale=2):
+                with gr.Tabs():
+                    with gr.TabItem("🎨 Speaker Diarization"):
+                        diarized_output = gr.HTML(
+                            label="Speaker Diarization Results",
+                            elem_classes=["speaker-output"]
+                        )
+                    
+                    with gr.TabItem("📝 Full Transcription"):
+                        full_transcription = gr.Textbox(
+                            label="Complete Transcription",
+                            lines=15,
+                            max_lines=20,
+                            show_copy_button=True
+                        )
+        
+        # Event handlers
+        process_btn.click(
+            fn=process_audio,
+            inputs=[audio_input, sensitivity_slider],
+            outputs=[diarized_output, full_transcription],
+            show_progress=True
+        )
+        
+        # Auto-process when audio is uploaded
+        audio_input.change(
+            fn=process_audio,
+            inputs=[audio_input, sensitivity_slider],
+            outputs=[diarized_output, full_transcription],
+            show_progress=True
+        )
+        
+        gr.Markdown(
+            """
+            ---
+            ### About
+            This application uses:
+            - **MFCC features** for speaker embedding extraction
+            - **Cosine similarity** for speaker change detection
+            - **OpenAI Whisper** for speech-to-text transcription
+            - **Gradio** for the web interface
+            
+            **Note**: This is a simplified speaker diarization system. For production use, 
+            consider more advanced speaker embedding models like speechbrain or pyannote.audio.
+            """
+        )
+    
+    return demo
+
+# Create and launch the interface
+if __name__ == "__main__":
+    demo = create_interface()
+    demo.launch(
+        server_name="0.0.0.0",
+        server_port=7860,
+        share=False,
+        show_error=True
+    )