Update app.py

app.py

@@ -5,57 +5,59 @@ import re
 import random
 import librosa
 import soundfile as sf
+import pandas as pd
 from transformers import pipeline
 from datasets import load_dataset, Audio
 from gradio_client import Client
 from stats_data import get_indomain_breakdown, get_experimental_summary, SPEAKER_META
 
-# 1. Initialize Baseline ASR (Strict English, Repetition Penalty 3.0)
-print("Loading Whisper Tiny...")
+# 1. Configuration
+TORGO_INDICES = {'FC01': 0, 'FC02': 302, 'FC03': 2489, 'MC02': 4411, 'MC01': 5534, 'MC03': 7689, 'MC04': 9358, 'M05': 10978, 'M02': 11565, 'M04': 12337, 'M01': 13003, 'F01': 13746, 'M03': 13982, 'F04': 14792, 'F03': 15465}
+HF_TOKEN = os.getenv("HF_TOKEN")
+PRIVATE_BACKEND_URL = "st192011/Torgo-DSR-Private"
+
+# 2. Local Whisper Baseline (Strict English, Repetition Penalty 3.0)
 whisper_asr = pipeline(
     "automatic-speech-recognition", 
     model="openai/whisper-tiny",
-    generate_kwargs={"language": "en", "task": "transcribe", "repetition_penalty": 3.0}
+    generate_kwargs={"language": "en", "task": "transcribe", "repetition_penalty": 3.0, "max_new_tokens": 64}
 )
 
-HF_TOKEN = os.getenv("HF_TOKEN")
-PRIVATE_BACKEND_URL = "st192011/Torgo-DSR-Private"
-
-def normalize(text):
+def normalize_text(text):
     if not text: return ""
     return re.sub(r'[^\w\s]', '', text).lower().strip()
 
+def standardize_audio(input_path):
+    """Ensures audio is 16kHz, Mono, and compatible with all models."""
+    if not input_path: return None
+    audio, sr = librosa.load(input_path, sr=16000, mono=True)
+    out_path = "processed_audio.wav"
+    sf.write(out_path, audio, 16000)
+    return out_path
+
 # --- Logic: Data Loading ---
 def get_sample_logic(speaker_id):
     try:
-        if "UA" in speaker_id:
-            # UA-Speech Access (Direct pull for F02)
+        if speaker_id == "F02 (UA)":
             dataset = load_dataset("resproj007/uaspeech_female", split="train", streaming=True)
             dataset = dataset.cast_column("audio", Audio(decode=False))
-            # UA is small, skip slightly for variety
             sample = next(iter(dataset.skip(random.randint(0, 30))))
-            gt_text = sample.get('text') or sample.get('transcription') or sample.get('sentence')
+            gt_text = sample.get('text') or sample.get('transcription') or "Unknown"
         else:
-            # Torgo Access (Manual filtering as per Colab fix)
             dataset = load_dataset("abnerh/TORGO-database", split="train", streaming=True)
             dataset = dataset.cast_column("audio", Audio(decode=False))
-            
-            def filter_spk(x):
-                sid = str(x.get('speaker_id', '')).upper()
-                if not sid or sid == "NONE":
-                    sid = os.path.basename(x['audio']['path']).split('_')[0].upper()
-                return sid == speaker_id
-            
-            speaker_ds = dataset.filter(filter_spk)
-            sample = next(iter(speaker_ds.shuffle(buffer_size=10)))
-            gt_text = sample.get('transcription') or sample.get('text')
+            start_idx = TORGO_INDICES.get(speaker_id, 0)
+            sample = next(iter(dataset.skip(start_idx + random.randint(0, 15))))
+            gt_text = sample.get('transcription') or sample.get('text') or "Unknown"
 
-        # Decode Bytes manually to bypass torchcodec errors
+        # Decode Bytes manually and Standardize
         audio_bytes = sample['audio']['bytes']
-        audio_data, sr = librosa.load(io.BytesIO(audio_bytes), sr=16000)
+        audio_data, _ = librosa.load(io.BytesIO(audio_bytes), sr=16000, mono=True)
+        
         temp_path = "sample.wav"
-        sf.write(temp_path, audio_data, sr)
+        sf.write(temp_path, audio_data, 16000)
         
+        # We return the path to the gr.Audio component (which stores it in State)
         return temp_path, gt_text.lower().strip(), SPEAKER_META[speaker_id]
     except Exception as e:
         return None, f"Dataset Error: {e}", {}
@@ -63,38 +65,42 @@ def get_sample_logic(speaker_id):
 # --- Logic: Model Steps ---
 def run_whisper_step(audio_path):
     if not audio_path: return "No audio loaded", ""
-    result = whisper_asr(audio_path)
+    # Standardize format before Whisper
+    clean_audio = standardize_audio(audio_path)
+    result = whisper_asr(clean_audio)
     raw_w = result["text"]
-    norm_w = normalize(raw_w)
+    norm_w = normalize_text(raw_w)
     return raw_w, norm_w
 
 def run_model_step(audio_path, norm_whisper):
-    if not audio_path or not norm_whisper: return "Complete Steps 1 & 2 first."
+    if not audio_path or not norm_whisper: return "Complete Step 1 & 2 first."
+    # Standardize format before sending to Private Backend
+    clean_audio = standardize_audio(audio_path)
     try:
-        # Call the private space for the 5K Gemma Model prediction
         client = Client(PRIVATE_BACKEND_URL, hf_token=HF_TOKEN)
-        prediction = client.predict(audio_path, norm_whisper, api_name="/predict_dsr")
+        prediction = client.predict(clean_audio, norm_whisper, api_name="/predict_dsr")
         return prediction
     except Exception as e:
-        return f"Backend Offline. Research Details: {e}"
+        return f"Backend Offline. Details: {e}"
 
 # --- UI Construction ---
 with gr.Blocks(theme=gr.themes.Soft(), title="Torgo DSR Lab") as demo:
     gr.Markdown("# ⚗️ Torgo DSR Lab")
-    gr.Markdown("Neural Reconstruction for Severe Dysarthria benchmarked on Torgo and UA-Speech.")
+    gr.Markdown("Neural Reconstruction for Severe Dysarthria. Load samples from Torgo/UA or record your own.")
     
-    current_audio_path = gr.State("")
-
     with gr.Tab("🔬 Laboratory"):
         with gr.Row():
             with gr.Column(scale=1):
-                gr.Markdown("### Step 1: Load Sample")
-                speaker_input = gr.Dropdown(sorted(list(SPEAKER_META.keys())), label="Speaker ID", value="F01")
-                load_btn = gr.Button("Load Data")
-                meta_display = gr.JSON(label="Speaker Meta")
-                gt_box = gr.Textbox(label="Ground Truth")
-                # Added visible audio for user verification
-                audio_preview = gr.Audio(label="Audio Preview", type="filepath")
+                gr.Markdown("### Step 1: Input Audio")
+                speaker_input = gr.Dropdown(sorted(list(SPEAKER_META.keys())), label="Speaker Profile (for Dataset Samples)")
+                load_btn = gr.Button("🎲 Load Dataset Sample")
+                
+                gr.Markdown("---")
+                # Unified Input: Handles both Dataset Samples and User Input
+                audio_input = gr.Audio(sources=["microphone", "upload"], type="filepath", label="Audio Input (Record/Upload/Dataset)")
+                
+                meta_display = gr.JSON(label="Speaker Metadata")
+                gt_box = gr.Textbox(label="Ground Truth (if from dataset)")
 
             with gr.Column(scale=2):
                 gr.Markdown("### Step 2: ASR Baseline")
@@ -109,35 +115,22 @@ with gr.Blocks(theme=gr.themes.Soft(), title="Torgo DSR Lab") as demo:
 
     with gr.Tab("📊 Research Statistics"):
         gr.Markdown("# 🔬 Performance Evaluation")
-        
         with gr.Row():
             with gr.Column():
-                gr.Markdown("""
-                ### 📏 Metric: Exact Match Accuracy
-                Accuracy is the percentage of samples where the **normalized prediction** (lowercase, no punctuation) matches the **ground truth**.
-                """)
-            
+                gr.Markdown("### 📏 Metric: Exact Match Accuracy")
+                gr.Markdown("Accuracy is calculated on normalized text (lowercase, no punctuation).")
             with gr.Column():
-                gr.Markdown("""
-                ### 🧪 Model Definitions
-                * **5K Pure Model:** Trained on real articulatory distortions. Optimized for phonetic fidelity.
-                * **10K Triple-Mix Model:** Includes synthetic data and anchors; utilized for generalization testing.
-                """)
+                gr.Markdown("### 🧪 Model Definitions")
+                gr.Markdown("* **5K Pure Model:** Real data focus. \n* **10K Triple-Mix Model:** LOSO Generalization focus.")
 
         gr.Markdown("## 1. Torgo In-Domain Analysis")
         gr.DataFrame(get_indomain_breakdown())
-        
         gr.Markdown("## 2. Experimental Summary")
         gr.DataFrame(get_experimental_summary())
 
-    # Event Mapping
-    load_btn.click(
-        get_sample_logic, 
-        inputs=speaker_input, 
-        outputs=[current_audio_path, gt_box, meta_display]
-    ).then(lambda x: x, inputs=current_audio_path, outputs=audio_preview)
-
-    whisper_btn.click(run_whisper_step, inputs=current_audio_path, outputs=[w_raw, w_norm])
-    model_btn.click(run_model_step, inputs=[current_audio_path, w_norm], outputs=final_out)
+    # Logic connections
+    load_btn.click(get_sample_logic, inputs=speaker_input, outputs=[audio_input, gt_box, meta_display])
+    whisper_btn.click(run_whisper_step, inputs=audio_input, outputs=[w_raw, w_norm])
+    model_btn.click(run_model_step, inputs=[audio_input, w_norm], outputs=final_out)
 
 demo.launch()