app.py

"""
app.py — gradio.Server backend for Village Health Screener.
Uses @app.api() for queued ML endpoints and standard FastAPI routes for serving the UI.
"""

import os
import base64
import tempfile

from gradio import Server
from fastapi.responses import HTMLResponse, FileResponse

import inference

# Conditional import for @spaces.GPU (only available on HF Spaces)
try:
    import spaces
except ImportError:

    class spaces:
        GPU = staticmethod(lambda f: f)


app = Server()


# ---------------------------------------------------------------------------
# Static file serving
# ---------------------------------------------------------------------------
@app.get("/")
async def homepage():
    """Serve the custom index.html frontend."""
    html_path = os.path.join(os.path.dirname(os.path.abspath(__file__)), "index.html")
    with open(html_path, "r", encoding="utf-8") as f:
        return HTMLResponse(content=f.read())


@app.get("/static/{filename}")
async def serve_static(filename: str):
    """Serve CSS/JS files from the static/ directory."""
    file_path = os.path.join(
        os.path.dirname(os.path.abspath(__file__)), "static", filename
    )
    return FileResponse(file_path)


# ---------------------------------------------------------------------------
# Health check
# ---------------------------------------------------------------------------
@app.get("/health")
async def health():
    return {"status": "ok", "model": "nemotron-mini-4b"}


# ---------------------------------------------------------------------------
# API Endpoints (queued via Gradio, GPU-managed via @spaces.GPU)
# ---------------------------------------------------------------------------
@app.api()
@spaces.GPU
def transcribe(audio_b64: str) -> dict:
    """Transcribe Hindi speech from base64-encoded audio."""
    try:
        audio_bytes = base64.b64decode(audio_b64)
        with tempfile.NamedTemporaryFile(suffix=".wav", delete=False) as f:
            f.write(audio_bytes)
            temp_path = f.name

        text = inference.transcribe_audio(temp_path)

        try:
            os.remove(temp_path)
        except OSError:
            pass

        return {"text": text}
    except Exception as e:
        print(f"[app] Transcription error: {e}")
        return {"text": "", "error": str(e)}


@app.api()
@spaces.GPU
def describe_image(image_b64: str) -> dict:
    """Describe visible symptoms from a base64-encoded image."""
    try:
        image_bytes = base64.b64decode(image_b64)
        with tempfile.NamedTemporaryFile(suffix=".jpg", delete=False) as f:
            f.write(image_bytes)
            temp_path = f.name

        description = inference.describe_image_file(temp_path)

        try:
            os.remove(temp_path)
        except OSError:
            pass

        return {"description": description}
    except Exception as e:
        print(f"[app] Vision error: {e}")
        return {"description": "", "error": str(e)}


@app.api()
@spaces.GPU
def run_triage(symptoms_text: str, visual_description: str) -> dict:
    """Produce a structured clinical triage report."""
    try:
        report = inference.run_nemotron_triage(symptoms_text, visual_description)
        return report
    except Exception as e:
        print(f"[app] Triage error: {e}")
        return {
            "likely_condition": "Assessment Error",
            "severity": 3,
            "immediate_actions": ["Consult a doctor immediately"],
            "refer_to_doctor": True,
            "refer_reason": f"Triage failed: {str(e)}",
            "followup_days": 1,
            "confidence": "low",
        }


@app.api()
@spaces.GPU
def synthesize_audio(text: str, language: str) -> dict:
    """Synthesize TTS audio and return as base64."""
    try:
        if language == "hindi":
            audio_bytes = inference.synthesize_hindi(text)
        else:
            audio_bytes = inference.synthesize_english(text)

        audio_b64 = base64.b64encode(audio_bytes).decode("utf-8") if audio_bytes else ""
        return {"audio_b64": audio_b64, "mime_type": "audio/wav"}
    except Exception as e:
        print(f"[app] TTS error: {e}")
        return {"audio_b64": "", "mime_type": "audio/wav", "error": str(e)}


# ---------------------------------------------------------------------------
# Launch
# ---------------------------------------------------------------------------
if __name__ == "__main__":
    app.launch(show_error=True, server_name="0.0.0.0")