""" Hugging Face Spaces - Gradio App for Stutter Analysis ===================================================== This is a standalone Gradio app for deployment on Hugging Face Spaces. To deploy: 1. Create a new Space on huggingface.co/spaces 2. Choose "Gradio" as SDK 3. Upload this folder's contents 4. Add your model checkpoint to the Space """ import gradio as gr import torch import torchaudio import tempfile import os import json import soundfile as sf import librosa from datetime import datetime from transformers import WavLMModel import torch.nn as nn import whisper # ============================================================================ # MODEL DEFINITION (same as models/WaveLm_model.py) # ============================================================================ class WaveLmStutterClassification(nn.Module): def __init__(self, num_labels=5, freeze_encoder=True, unfreeze_last_n_layers=1): super().__init__() self.wavlm = WavLMModel.from_pretrained("microsoft/wavlm-base") self.hidden_size = self.wavlm.config.hidden_size if freeze_encoder: for param in self.wavlm.parameters(): param.requires_grad = False if unfreeze_last_n_layers > 0: for layer in self.wavlm.encoder.layers[-unfreeze_last_n_layers:]: for param in layer.parameters(): param.requires_grad = True # Single linear layer to match the trained checkpoint self.classifier = nn.Linear(self.hidden_size, num_labels) self.num_labels = num_labels def forward(self, input_values, attention_mask=None): outputs = self.wavlm(input_values, attention_mask=attention_mask) hidden_states = outputs.last_hidden_state pooled = hidden_states.mean(dim=1) logits = self.classifier(pooled) return logits # ============================================================================ # STUTTER LABELS & DEFINITIONS # ============================================================================ STUTTER_LABELS = ['Prolongation', 'Block', 'SoundRep', 'WordRep', 'Interjection'] STUTTER_DEFINITIONS = { 'Prolongation': 'Sound stretched longer than normal (e.g., "Ssssssnake")', 'Block': 'Complete stoppage of airflow/sound with tension', 'SoundRep': 'Sound/syllable repetition (e.g., "B-b-b-ball")', 'WordRep': 'Whole word repetition (e.g., "I-I-I want")', 'Interjection': 'Filler words like "um", "uh", "like"' } SEVERITY_THRESHOLDS = {'very_mild': 5, 'mild': 10, 'moderate': 20, 'severe': 30} # ============================================================================ # GLOBAL MODEL LOADING # ============================================================================ device = "cuda" if torch.cuda.is_available() else "cpu" wavlm_model = None whisper_model = None def load_models(): global wavlm_model, whisper_model # Load WavLM print("Loading WavLM model...") wavlm_model = WaveLmStutterClassification(num_labels=5) # Try to load checkpoint checkpoint_path = "wavlm_stutter_classification_best.pth" if os.path.exists(checkpoint_path): checkpoint = torch.load(checkpoint_path, map_location=device) # Handle both formats: direct state_dict OR wrapped in 'model_state_dict' if isinstance(checkpoint, dict) and 'model_state_dict' in checkpoint: wavlm_model.load_state_dict(checkpoint['model_state_dict']) print(f"Loaded checkpoint with {checkpoint.get('val_accuracy', 'N/A')} accuracy") else: # Direct state_dict (how train_waveLM.py saves it) wavlm_model.load_state_dict(checkpoint) print("Loaded checkpoint (direct state_dict format)") else: print("WARNING: No checkpoint found, using random weights") wavlm_model.to(device) wavlm_model.eval() # Load Whisper print("Loading Whisper model...") whisper_model = whisper.load_model("base", device=device) print("Models loaded!") # ============================================================================ # ANALYSIS FUNCTIONS # ============================================================================ def preprocess_audio(audio_path): """Convert audio to 16kHz mono using soundfile or librosa.""" try: # Try loading with soundfile first (faster) waveform_np, sr = sf.read(audio_path, dtype='float32') # Handle multi-channel (soundfile returns (samples, channels)) if len(waveform_np.shape) > 1: waveform_np = waveform_np.mean(axis=1) except Exception as e: print(f"Soundfile load failed, trying librosa: {e}") # Fallback to librosa (handles mp3/m4a better via ffmpeg) # librosa loads as mono by default, and we can force sr=16000 here waveform_np, sr = librosa.load(audio_path, sr=16000, mono=True) # Convert to tensor waveform = torch.from_numpy(waveform_np).float() # Resample if needed (only if soundfile was used and sr != 16000) # If librosa was used, it's already 16000 if sr != 16000: resampler = torchaudio.transforms.Resample(sr, 16000) waveform = resampler(waveform.unsqueeze(0)).squeeze(0) return waveform, 16000 def chunk_audio(waveform, sr, chunk_sec=3.0): """Split audio into chunks""" chunk_samples = int(chunk_sec * sr) chunks = [] for start in range(0, len(waveform), chunk_samples): end = min(start + chunk_samples, len(waveform)) chunk = waveform[start:end] # Pad if needed if len(chunk) < chunk_samples: chunk = torch.nn.functional.pad(chunk, (0, chunk_samples - len(chunk))) chunks.append({ 'chunk': chunk, 'start': start / sr, 'end': end / sr }) return chunks def analyze_chunk(chunk_waveform, threshold=0.5): """Run WavLM on a single chunk""" with torch.no_grad(): input_tensor = chunk_waveform.unsqueeze(0).to(device) logits = wavlm_model(input_tensor) probs = torch.sigmoid(logits).cpu().numpy()[0] detected = [STUTTER_LABELS[i] for i, p in enumerate(probs) if p > threshold] probabilities = {STUTTER_LABELS[i]: float(probs[i]) for i in range(len(STUTTER_LABELS))} return {'detected': detected, 'probabilities': probabilities} def get_severity(word_stutter_rate): """Calculate severity from word stutter rate""" if word_stutter_rate < SEVERITY_THRESHOLDS['very_mild']: return 'Very Mild', 1 elif word_stutter_rate < SEVERITY_THRESHOLDS['mild']: return 'Mild', 2 elif word_stutter_rate < SEVERITY_THRESHOLDS['moderate']: return 'Moderate', 3 elif word_stutter_rate < SEVERITY_THRESHOLDS['severe']: return 'Severe', 4 else: return 'Very Severe', 5 # ============================================================================ # MAIN ANALYSIS FUNCTION # ============================================================================ def analyze_audio(audio_file, threshold=0.5): """Main analysis function for Gradio""" if wavlm_model is None: load_models() if audio_file is None: return "⚠️ Please upload an audio file", "", "", "" try: print(f"Starting analysis of: {audio_file}") # Preprocess waveform, sr = preprocess_audio(audio_file) duration = len(waveform) / sr print(f"Audio preprocessed: {duration:.1f}s, {sr}Hz") # Chunk and analyze with WavLM chunks = chunk_audio(waveform, sr) stutter_counts = {label: 0 for label in STUTTER_LABELS} timeline = [] for chunk_info in chunks: result = analyze_chunk(chunk_info['chunk'], threshold) for label in result['detected']: stutter_counts[label] += 1 timeline.append({ 'time': f"{chunk_info['start']:.1f}s - {chunk_info['end']:.1f}s", 'detected': ', '.join(result['detected']) if result['detected'] else 'Clear', 'probs': result['probabilities'] }) # Transcribe with Whisper whisper_result = whisper_model.transcribe(audio_file, word_timestamps=True) transcription = whisper_result['text'] # Get word-level info words = [] if 'segments' in whisper_result: for seg in whisper_result['segments']: if 'words' in seg: words.extend(seg['words']) # Map stutters to words words_with_stutter = 0 annotated_words = [] for word_info in words: word_start = word_info.get('start', 0) word_end = word_info.get('end', 0) word_text = word_info.get('word', '') word_stutters = [] for chunk_info in chunks: if word_start < chunk_info['end'] and word_end > chunk_info['start']: result = analyze_chunk(chunk_info['chunk'], threshold) word_stutters.extend(result['detected']) word_stutters = list(set(word_stutters)) if word_stutters: words_with_stutter += 1 annotated_words.append(f"**[{word_text}]**({', '.join(word_stutters)})") else: annotated_words.append(word_text) # Calculate metrics total_words = len(words) if words else 1 word_stutter_rate = (words_with_stutter / total_words) * 100 severity_label, severity_score = get_severity(word_stutter_rate) # Format outputs summary = f""" ## 📊 Analysis Summary **Duration:** {duration:.1f} seconds **Total Words:** {total_words} **Words with Stutters:** {words_with_stutter} ({word_stutter_rate:.1f}%) ### Severity: {severity_label} ({severity_score}/5) ### Stutter Type Counts: """ for label, count in stutter_counts.items(): if count > 0: summary += f"- **{label}**: {count} occurrences\n" # Annotated transcription annotated_text = " ".join(annotated_words) if annotated_words else transcription # Timeline timeline_text = "| Time | Detected Stutters |\n|------|-------------------|\n" for t in timeline[:15]: # Limit to 15 rows timeline_text += f"| {t['time']} | {t['detected']} |\n" # Definitions definitions = "## 📖 Stutter Type Definitions\n\n" for label, desc in STUTTER_DEFINITIONS.items(): definitions += f"**{label}:** {desc}\n\n" return summary, annotated_text, timeline_text, definitions except Exception as e: import traceback error_trace = traceback.format_exc() print(f"Error in analyze_audio: {error_trace}") return f"❌ Error: {str(e)}\n\n```\n{error_trace}\n```", "", "", "" # ============================================================================ # GRADIO INTERFACE # ============================================================================ with gr.Blocks(title="🎙️ Stutter Analysis") as demo: gr.Markdown(""" # 🎙️ Speech Fluency Analysis System Upload an audio file to analyze stuttering patterns using AI. **Supported formats:** WAV, MP3, M4A, FLAC """) with gr.Row(): with gr.Column(scale=1): audio_input = gr.Audio( label="Upload Audio", type="filepath", sources=["upload", "microphone"] ) threshold_slider = gr.Slider( minimum=0.3, maximum=0.7, value=0.5, step=0.05, label="Detection Threshold", info="Lower = more sensitive, Higher = more conservative" ) analyze_btn = gr.Button("🔍 Analyze Speech", variant="primary") with gr.Column(scale=2): summary_output = gr.Markdown(label="Summary") with gr.Tabs(): with gr.Tab("📝 Transcription"): transcription_output = gr.Markdown(label="Annotated Transcription") with gr.Tab("📈 Timeline"): timeline_output = gr.Markdown(label="Timeline Analysis") with gr.Tab("📖 Definitions"): definitions_output = gr.Markdown(label="Stutter Definitions") analyze_btn.click( fn=analyze_audio, inputs=[audio_input, threshold_slider], outputs=[summary_output, transcription_output, timeline_output, definitions_output] ) gr.Markdown(""" --- **Disclaimer:** This tool is for educational/research purposes. Consult a qualified speech-language pathologist for clinical diagnosis. Built with WavLM + Whisper | [GitHub](https://github.com/abhicodes-here2001/Multimodal-stuttering-analysis) """) # Load models on startup load_models() if __name__ == "__main__": demo.launch(theme=gr.themes.Soft())