Upload 3 files

app.py

@@ -1,48 +1,21 @@
-"""
-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 numpy as np
 import os
-import json
-import soundfile as sf
-import librosa
+import traceback
 from datetime import datetime
 from transformers import WavLMModel
 import torch.nn as nn
-import whisper
 
-# ============================================================================
-# MODEL DEFINITION (same as models/WaveLm_model.py)
-# ============================================================================
+print(f"APP STARTUP: {datetime.now()}")
 
 class WaveLmStutterClassification(nn.Module):
-    def __init__(self, num_labels=5, freeze_encoder=True, unfreeze_last_n_layers=1):
+    def __init__(self, num_labels=5):
         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
+        for param in self.wavlm.parameters():
+            param.requires_grad = False
         self.classifier = nn.Linear(self.hidden_size, num_labels)
         self.num_labels = num_labels
     
@@ -53,310 +26,173 @@ class WaveLmStutterClassification(nn.Module):
         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"'
+    'Prolongation': 'Sound stretched longer than normal',
+    'Block': 'Complete stoppage of airflow/sound',
+    'SoundRep': 'Sound/syllable repetition',
+    'WordRep': 'Whole word repetition',
+    'Interjection': 'Filler words like um, uh'
 }
 
-SEVERITY_THRESHOLDS = {'very_mild': 5, 'mild': 10, 'moderate': 20, 'severe': 30}
-
-# ============================================================================
-# GLOBAL MODEL LOADING
-# ============================================================================
-
 device = "cuda" if torch.cuda.is_available() else "cpu"
+print(f"Device: {device}")
+
 wavlm_model = None
 whisper_model = None
+models_loaded = False
 
 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."""
+    global wavlm_model, whisper_model, models_loaded
+    if models_loaded:
+        return True
     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)
-            
+        print("Loading WavLM...")
+        wavlm_model = WaveLmStutterClassification(num_labels=5)
+        checkpoint_path = "wavlm_stutter_classification_best.pth"
+        if os.path.exists(checkpoint_path):
+            checkpoint = torch.load(checkpoint_path, map_location=device, weights_only=False)
+            if isinstance(checkpoint, dict) and 'model_state_dict' in checkpoint:
+                wavlm_model.load_state_dict(checkpoint['model_state_dict'])
+            else:
+                wavlm_model.load_state_dict(checkpoint)
+            print("Checkpoint loaded!")
+        wavlm_model.to(device)
+        wavlm_model.eval()
+        
+        print("Loading Whisper...")
+        import whisper
+        whisper_model = whisper.load_model("base", device=device)
+        
+        models_loaded = True
+        print("Models loaded!")
+        return True
     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
+        print(f"Model loading error: {e}")
+        traceback.print_exc()
+        return False
 
-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 load_audio(audio_path):
+    print(f"Loading: {audio_path}")
+    try:
+        import librosa
+        waveform, sr = librosa.load(audio_path, sr=16000, mono=True)
+        return torch.from_numpy(waveform).float(), 16000
+    except Exception as e:
+        print(f"librosa error: {e}")
+    try:
+        import soundfile as sf
+        waveform, sr = sf.read(audio_path, dtype='float32')
+        if len(waveform.shape) > 1:
+            waveform = waveform.mean(axis=1)
+        waveform = torch.from_numpy(waveform).float()
+        if sr != 16000:
+            import torchaudio
+            waveform = torchaudio.transforms.Resample(sr, 16000)(waveform.unsqueeze(0)).squeeze(0)
+        return waveform, 16000
+    except Exception as e:
+        print(f"soundfile error: {e}")
+    raise Exception("Could not load audio")
 
-def analyze_chunk(chunk_waveform, threshold=0.5):
-    """Run WavLM on a single chunk"""
+def analyze_chunk(chunk_tensor, threshold=0.5):
     with torch.no_grad():
-        input_tensor = chunk_waveform.unsqueeze(0).to(device)
-        logits = wavlm_model(input_tensor)
+        logits = wavlm_model(chunk_tensor.unsqueeze(0).to(device))
         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}
+    return detected, dict(zip(STUTTER_LABELS, probs.tolist()))
 
-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"""
+def analyze_audio(audio_input, threshold):
+    print(f"\n=== ANALYZE CLICKED ===")
+    print(f"Input: {audio_input}, Type: {type(audio_input)}, Threshold: {threshold}")
+    
+    if audio_input is None:
+        return "Please upload an audio file first!", "", "", ""
     
-    if wavlm_model is None:
-        load_models()
+    audio_path = audio_input
+    if isinstance(audio_input, tuple):
+        import tempfile, soundfile as sf
+        sr, data = audio_input
+        f = tempfile.NamedTemporaryFile(suffix=".wav", delete=False)
+        sf.write(f.name, data, sr)
+        audio_path = f.name
     
-    if audio_file is None:
-        return "⚠️ Please upload an audio file", "", "", ""
+    if not os.path.exists(audio_path):
+        return f"File not found: {audio_path}", "", "", ""
+    
+    print(f"File: {audio_path}, Size: {os.path.getsize(audio_path)}")
     
     try:
-        print(f"Starting analysis of: {audio_file}")
+        if not models_loaded and not load_models():
+            return "Failed to load models", "", "", ""
         
-        # Preprocess
-        waveform, sr = preprocess_audio(audio_file)
+        waveform, sr = load_audio(audio_path)
         duration = len(waveform) / sr
-        print(f"Audio preprocessed: {duration:.1f}s, {sr}Hz")
-        
-        # Chunk and analyze with WavLM
-        chunks = chunk_audio(waveform, sr)
+        print(f"Duration: {duration:.1f}s")
         
-        stutter_counts = {label: 0 for label in STUTTER_LABELS}
+        chunk_samples = int(3.0 * sr)
+        stutter_counts = {l: 0 for l 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
+        for start in range(0, len(waveform), chunk_samples):
+            end = min(start + chunk_samples, len(waveform))
+            chunk = waveform[start:end]
+            if len(chunk) < chunk_samples:
+                chunk = torch.nn.functional.pad(chunk, (0, chunk_samples - len(chunk)))
             
-            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']
-            })
+            detected, _ = analyze_chunk(chunk, threshold)
+            for l in detected:
+                stutter_counts[l] += 1
+            timeline.append({"time": f"{start/sr:.1f}-{end/sr:.1f}s", "detected": detected or ["Clear"]})
         
-        # Transcribe with Whisper
-        whisper_result = whisper_model.transcribe(audio_file, word_timestamps=True)
-        transcription = whisper_result['text']
+        print("Running Whisper...")
+        transcription = whisper_model.transcribe(audio_path).get('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'])
+        total = sum(stutter_counts.values())
+        summary = f"## Analysis Complete\n\n**Duration:** {duration:.1f}s\n**Stutters:** {total}\n\n"
+        for l, c in stutter_counts.items():
+            summary += f"- {l}: {c}\n"
         
-        # Map stutters to words
-        words_with_stutter = 0
-        annotated_words = []
+        timeline_md = "| Time | Detected |\n|---|---|\n"
+        for t in timeline[:15]:
+            timeline_md += f"| {t['time']} | {', '.join(t['detected'])} |\n"
         
-        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)
+        defs = "\n".join([f"**{k}:** {v}" for k, v in STUTTER_DEFINITIONS.items()])
         
-        # 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
+        print("Done!")
+        return summary, transcription, timeline_md, defs
         
     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```", "", "", ""
+        print(f"Error: {e}")
+        traceback.print_exc()
+        return f"Error: {e}\n\n{traceback.format_exc()}", "", "", ""
 
-# ============================================================================
-# GRADIO INTERFACE
-# ============================================================================
+print("Building UI...")
 
-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.Blocks(title="Stutter Analysis") as demo:
+    gr.Markdown("# Speech Fluency Analysis\nUpload audio to analyze stuttering.")
     
     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.Column():
+            audio = gr.Audio(label="Upload Audio", type="filepath")
+            threshold = gr.Slider(0.3, 0.7, 0.5, label="Threshold")
+            btn = gr.Button("Analyze", variant="primary")
+        with gr.Column():
+            summary = gr.Markdown(value="Upload audio and click Analyze")
     
     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.
+        with gr.TabItem("Transcription"):
+            trans = gr.Markdown()
+        with gr.TabItem("Timeline"):
+            timeline = gr.Markdown()
+        with gr.TabItem("Definitions"):
+            defs = gr.Markdown()
     
-    Built with WavLM + Whisper | [GitHub](https://github.com/abhicodes-here2001/Multimodal-stuttering-analysis)
-    """)
+    btn.click(analyze_audio, [audio, threshold], [summary, trans, timeline, defs])
 
-# Load models on startup
+print("Loading models...")
 load_models()
 
-if __name__ == "__main__":
-    demo.launch(theme=gr.themes.Soft())
+print("Launching...")
+demo.queue()
+demo.launch()

requirements.txt

@@ -1,6 +1,3 @@
-# Hugging Face Spaces Requirements
-# For Gradio deployment
-
 torch>=2.0.0
 torchaudio>=2.0.0
 transformers>=4.30.0