app.py

"""
app.py - High-Performance Hugging Face Space (Gradio + ZeroGPU)
===============================================================
Deploys Anvaya Speech Pathology & Articulation Diagnostics on Hugging Face Spaces
with support for Free ZeroGPU / 16 GB RAM CPU execution.
"""
from __future__ import annotations

# CRITICAL FOR HUGGING FACE ZEROGPU: 'import spaces' MUST occur before torch / transformers
try:
    import spaces
    HAS_SPACES = True
except ImportError:
    HAS_SPACES = False

import gc
import json
import os
import time
import traceback
from pathlib import Path
from typing import Optional, Dict, Any, List, Tuple

import gradio as gr
import numpy as np
import soundfile as sf
import torch

from ml.model.engine import SpeechDiagnosticEngine

# Lazy-loaded / Singleton Engine
_engine_instance: Optional[SpeechDiagnosticEngine] = None

def get_engine() -> SpeechDiagnosticEngine:
    global _engine_instance
    if _engine_instance is None:
        _engine_instance = SpeechDiagnosticEngine.get_instance()
    return _engine_instance

# Presets
WORD_PRESETS = [
    "rabbit",
    "red",
    "sun",
    "sweet",
    "three",
    "water",
    "kitten",
    "spot",
]

SENTENCE_PRESETS = {
    "The Red Rabbit (Rhotacism & Rhotic 'R' Evaluation)": "the red rabbit ran around the green yard",
    "The Sweet Sun (Sigmatism & Sibilant 'S' Evaluation)": "the sweet sun shines softly in the sky",
    "The Blue Spot (Phonetic Balance & Plosives)": "the blue spot is on the key",
    "The Rainbow Passage (Standard Clinical Protocol)": "the rainbow is a division of white light into many beautiful colors",
}


def _to_json_safe(obj: Any) -> Any:
    """Recursively convert numpy/torch structures to JSON-serializable primitives."""
    if isinstance(obj, dict):
        return {str(k): _to_json_safe(v) for k, v in obj.items()}
    elif isinstance(obj, (list, tuple)):
        return [_to_json_safe(x) for x in obj]
    elif isinstance(obj, (np.floating, float)):
        return float(obj)
    elif isinstance(obj, (np.integer, int)):
        return int(obj)
    elif isinstance(obj, np.ndarray):
        return _to_json_safe(obj.tolist())
    else:
        return obj


def _diagnose_speech_core(
    audio_path: Optional[str],
    target_phrase: str,
    healthy_baseline_path: Optional[str] = None,
) -> Tuple[str, str, str, str, str, str, dict]:
    """Run full diagnostic pipeline and return formatted clinical report for Gradio."""
    empty_dict = {"status": "waiting_for_input"}
    
    if not audio_path:
        return (
            "**Notice**: Please record speech or upload an audio file to evaluate.",
            "NO AUDIO",
            "N/A",
            "N/A",
            "",
            "",
            empty_dict,
        )

    if not target_phrase or not target_phrase.strip():
        return (
            "**Notice**: Please enter or select an expected Target Phrase.",
            "NO TARGET",
            "N/A",
            "N/A",
            "",
            "",
            empty_dict,
        )

    try:
        engine = get_engine()

        # Execute Diagnostic Engine
        diag_res = engine.diagnose_audio(
            audio_input=audio_path,
            target_phrase=target_phrase,
            normal_calibration_audio=healthy_baseline_path,
        )

        if diag_res.get("is_silent") or diag_res["decision"].get("is_silent"):
            return (
                "**No Speech Detected**: The audio is silent or below acoustic energy thresholds. Please speak clearly into your microphone.",
                "SILENT",
                "0 / 100",
                "0.0%",
                "",
                "- **Silence Guard Active**: No vocal signal detected in audio stream.",
                _to_json_safe(diag_res),
            )

        result = diag_res["decision"]
        pron = diag_res["pronunciation"]
        flaws = diag_res["flaws"]
        artic = diag_res["articulation"]
        p_stut = float(diag_res["stutter_probs"][1]) if (diag_res.get("stutter_probs") and len(diag_res["stutter_probs"]) > 1) else 0.0

        overall_bucket = result["buckets"]["overall"].upper()
        fluency_score = f"{int(result.get('fluency_100', 100))} / 100"
        pron_acc = f"{max(0.0, min(100.0, (1.0 - pron.get('wer', 0.0)) * 100.0)):.1f}%"

        # Build Word Alignment Chips HTML
        alignment = pron.get("alignment", [])
        chips_html = "<div style='display:flex; flex-wrap:wrap; gap:8px; padding:12px; background:rgba(15,23,42,0.6); border-radius:8px; margin:10px 0;'>"
        for item in alignment:
            status = item["status"]
            exp = item["expected"]
            spk = item["spoken"]
            if status == "correct":
                chips_html += f"<span style='padding:6px 12px; background:rgba(16,185,129,0.15); color:#34D399; border:1px solid rgba(16,185,129,0.3); border-radius:6px; font-weight:600;'>[MATCH] {exp}</span>"
            elif status == "substitution":
                chips_html += f"<span style='padding:6px 12px; background:rgba(239,68,68,0.15); color:#F87171; border:1px solid rgba(239,68,68,0.3); border-radius:6px; font-weight:600;'>[DIFF] {exp} (heard: \"{spk}\")</span>"
            elif status == "omission":
                chips_html += f"<span style='padding:6px 12px; background:rgba(245,158,11,0.15); color:#FBBF24; border:1px solid rgba(245,158,11,0.3); border-radius:6px; font-weight:600;'>[UNSPOKEN] {exp}</span>"
            elif status == "insertion":
                chips_html += f"<span style='padding:6px 12px; background:rgba(168,85,247,0.15); color:#C084FC; border:1px solid rgba(168,85,247,0.3); border-radius:6px; font-weight:600;'>[EXTRA] {spk}</span>"
        chips_html += "</div>"

        # Build Flaw Report Summary Markdown
        flaws_md = "### Specific Speech Pathology Findings:\n\n"
        if flaws["has_r_flaw"]:
            for r_err in flaws["r_sound_issues"]:
                flaws_md += f"- **Rhotacism Flaw**: {r_err['message']}\n"
        else:
            flaws_md += "- **'R' Sound Articulation**: Accurate (No R->W/L substitution detected).\n"

        if flaws["has_s_flaw"]:
            for s_err in flaws["s_sound_issues"]:
                flaws_md += f"- **Sigmatism Flaw**: {s_err['message']}\n"
        else:
            flaws_md += "- **'S' Sound Articulation**: Accurate (No sibilant lisp detected).\n"

        if p_stut >= 0.78:
            flaws_md += f"- **Disfluency Detected**: Elevated probability of repetition/block ({p_stut*100:.1f}%)\n"
        elif p_stut >= 0.60:
            flaws_md += f"- **Mild Hesitation**: Minor syllable repetition observed ({p_stut*100:.1f}%)\n"
        else:
            flaws_md += "- **Fluency Flow**: Continuous cadence (No disfluent events detected).\n"

        flaws_md += f"- **Voice Phonation Correlates**: Pitch F0={artic.get('f0_median_hz',0):.1f}Hz, HNR={artic.get('hnr_db',0):.1f}dB, Jitter={artic.get('jitter',0)*100:.2f}%\n"

        confidence_val = result.get("confidence", "high")
        heard_summary = f"**Decoded Transcription**: *\"{pron.get('asr_hypothesis','')}\"*\n\n**Confidence Rating**: `{confidence_val}` | **Latency**: `{diag_res['latency_ms']} ms`"

        return (
            heard_summary,
            overall_bucket,
            fluency_score,
            pron_acc,
            chips_html,
            flaws_md,
            _to_json_safe(diag_res),
        )
    except Exception as ex:
        err_msg = f"**Execution Error**: {ex}\n\n```\n{traceback.format_exc()}\n```"
        return (
            err_msg,
            "ERROR",
            "0 / 100",
            "0.0%",
            "",
            f"- **Internal Error**: {ex}",
            {"error": str(ex), "traceback": traceback.format_exc()},
        )


# Apply ZeroGPU acceleration decorator if running on Hugging Face ZeroGPU
if HAS_SPACES:
    @spaces.GPU
    def diagnose_speech_hf(*args, **kwargs):
        return _diagnose_speech_core(*args, **kwargs)
else:
    def diagnose_speech_hf(*args, **kwargs):
        return _diagnose_speech_core(*args, **kwargs)


# Construct Gradio Modern Interface
with gr.Blocks(title="Anvaya | Speech Screening & Phonetics") as demo:
    gr.Markdown("""
    # ANVAYA · Speech Screening & Phonetic Analysis Assistant
    ### Multi-Modal Screening: Neural Disfluency · Rhotacism ('r') · Sigmatism ('s' Lisp) · Voice Acoustics
    
    > **CLINICAL PRACTICE & RESEARCH DISCLAIMER**: Anvaya is an exploratory engineering prototype for speech screening and practice. It is **not** an FDA-cleared medical device, nor a substitute for professional clinical evaluation by a licensed Speech-Language Pathologist (SLP).
    """)

    with gr.Row():
        with gr.Column(scale=1):
            audio_input = gr.Audio(
                sources=["microphone", "upload"],
                type="filepath",
                label="Audio Ingestion (Record Microphone or Upload Audio)",
            )

            gr.Markdown("#### Single-Word Practice Presets:")
            with gr.Row():
                btn_w1 = gr.Button("rabbit", size="sm")
                btn_w2 = gr.Button("red", size="sm")
                btn_w3 = gr.Button("sun", size="sm")
                btn_w4 = gr.Button("sweet", size="sm")
            with gr.Row():
                btn_w5 = gr.Button("three", size="sm")
                btn_w6 = gr.Button("water", size="sm")
                btn_w7 = gr.Button("kitten", size="sm")
                btn_w8 = gr.Button("spot", size="sm")

            target_preset = gr.Dropdown(
                choices=list(SENTENCE_PRESETS.keys()),
                label="Standardized Clinical Protocols:",
                value="The Red Rabbit (Rhotacism & Rhotic 'R' Evaluation)",
            )

            target_text = gr.Textbox(
                label="Target Phrase (Expected Spoken Text):",
                value=SENTENCE_PRESETS["The Red Rabbit (Rhotacism & Rhotic 'R' Evaluation)"],
                lines=2,
            )

            # Bind Word Preset Buttons to set target phrase text
            btn_w1.click(fn=lambda: "rabbit", outputs=[target_text])
            btn_w2.click(fn=lambda: "red", outputs=[target_text])
            btn_w3.click(fn=lambda: "sun", outputs=[target_text])
            btn_w4.click(fn=lambda: "sweet", outputs=[target_text])
            btn_w5.click(fn=lambda: "three", outputs=[target_text])
            btn_w6.click(fn=lambda: "water", outputs=[target_text])
            btn_w7.click(fn=lambda: "kitten", outputs=[target_text])
            btn_w8.click(fn=lambda: "spot", outputs=[target_text])

            target_preset.change(
                fn=lambda k: SENTENCE_PRESETS.get(k, ""),
                inputs=[target_preset],
                outputs=[target_text],
            )

            healthy_baseline = gr.Audio(
                sources=["upload"],
                type="filepath",
                label="Healthy Baseline Sample (Optional 'My Normal' Calibration):",
            )

            diagnose_btn = gr.Button("Run Diagnostic Analysis", variant="primary", size="lg")

        with gr.Column(scale=2):
            with gr.Row():
                kpi_strat = gr.Textbox(label="Clinical Stratification", interactive=False)
                kpi_fluency = gr.Textbox(label="Fluency Index", interactive=False)
                kpi_acc = gr.Textbox(label="Pronunciation Accuracy", interactive=False)

            summary_box = gr.Markdown("### Clinical Assessment Summary\n*Results will appear here after analysis.*")
            alignment_html = gr.HTML(label="Word-Level Alignment")
            flaws_box = gr.Markdown("### Specific Speech Pathology Findings\n*Sound checks will appear here.*")

            with gr.Accordion("Auditable Telemetry & Acoustic Evidence Trace", open=False):
                raw_json = gr.JSON()

    diagnose_btn.click(
        fn=diagnose_speech_hf,
        inputs=[audio_input, target_text, healthy_baseline],
        outputs=[summary_box, kpi_strat, kpi_fluency, kpi_acc, alignment_html, flaws_box, raw_json],
    )

if __name__ == "__main__":
    demo.launch()