app.py

"""Gradio entrypoint for the ChatCAD Hugging Face Space demo.

Single-page, single-analysis UI around `chatcad.run_pipeline`. See
`2026-07-06-gradio-hf-demo-design.md` for the full design.
"""

from __future__ import annotations

import json
import os
import tempfile

import gradio as gr
from huggingface_hub import hf_hub_download

import chatcad

WEIGHTS_REPO_ID = os.environ.get("CHATCAD_WEIGHTS_REPO_ID", "")
# (filename in the HF model repo, local destination directory)
REMOTE_ASSETS = [
    ("JFchexpert.pth", "./weights"),
    ("r2gcmn_mimic-cxr.pth", "./weights"),
    ("annotation.json", "./r2g"),
]

PROVIDER_MAP = {"Claude": "anthropic", "OpenAI": "openai", "Gemini": "gemini"}

REPORT_FIELD_LABELS = [
    ("study_type", "Study Type"),
    ("summary", "Summary"),
    ("main_findings", "Main Findings"),
    ("detail_findings", "Detailed Findings"),
    ("impression", "Impression"),
    ("recommendations", "Recommendations"),
    ("additional_informations", "Additional Info"),
]


def _ensure_weights() -> None:
    """Download model checkpoints and annotation data from the HF model repo if not already local."""

    if not WEIGHTS_REPO_ID:
        return

    for filename, dest_dir in REMOTE_ASSETS:
        os.makedirs(dest_dir, exist_ok=True)
        local_path = os.path.join(dest_dir, filename)
        if os.path.exists(local_path):
            continue
        hf_hub_download(
            repo_id=WEIGHTS_REPO_ID,
            filename=filename,
            local_dir=dest_dir,
        )


def _format_final_report(report: dict) -> str:
    lines = []
    for key, label in REPORT_FIELD_LABELS:
        lines.append(f"### {label}\n\n{report.get(key, '')}\n")
    return "\n".join(lines)


def _format_network_a(network_a: dict) -> str:
    probs = network_a.get("probabilities", {})
    prob_lines = "\n".join(f"- **{name}**: {p:.3f}" for name, p in probs.items())
    return (
        "### Network A — Disease Classifier (JF CheXpert)\n\n"
        f"{prob_lines}\n\n"
        f"{network_a.get('severity_text', '')}\n"
    )


def _format_network_b(lesion_summary: dict) -> str:
    lines = ["### Network B — Lesion Segmentation\n"]
    for disease, info in lesion_summary.items():
        lines.append(
            f"- **{disease}**: probability={info['probability']:.3f}, "
            f"coverage={info['lesion_coverage_percent']:.1f}%, "
            f"peak_activation={info['peak_activation']:.3f}, "
            f"detected={info['detected']}"
        )
    return "\n".join(lines)


def _format_network_c(draft_report: str) -> str:
    return f"---\n\n### Network C — Initial Draft Report\n\n{draft_report}"


def _write_payload_json(payload: dict) -> str:
    fd, path = tempfile.mkstemp(suffix=".json", prefix="chatcad_report_")
    with os.fdopen(fd, "w", encoding="utf-8") as f:
        json.dump(payload, f, indent=2)
    return path


def handle_upload(image_path):
    if not image_path:
        return gr.update(visible=True), gr.update(value=None, visible=False), gr.update(visible=False)
    try:
        preview = chatcad.load_image_rgb_for_display(image_path)
    except Exception:
        preview = None
    return gr.update(visible=False), gr.update(value=preview, visible=True), gr.update(visible=True)


def clear_upload():
    return (
        gr.update(value=None, visible=True),
        gr.update(value=None, visible=False),
        gr.update(visible=False),
    )


def start_loading(image_path, provider_label, api_key):
    if not image_path:
        raise gr.Error("Please upload a chest X-ray image (raster or DICOM).")

    if not api_key or not api_key.strip():
        raise gr.Error(f"Please provide an API key for {provider_label}.")

    return (
        "⏳ Running analysis — this can take up to 60 seconds on CPU...",
        gr.update(visible=False),
        "",
        gr.update(visible=False),
    )


def analyze(image_path, provider_label, api_key):
    provider = PROVIDER_MAP[provider_label]

    try:
        _ensure_weights()
        state = chatcad.run_pipeline(image_path, provider=provider, api_key=api_key.strip())
    except (FileNotFoundError, ValueError) as e:
        raise gr.Error(str(e))
    except Exception as e:
        raise gr.Error(f"Analysis failed: {e}")

    payload = state.payload

    if state.modality != "chest_xray":
        return (
            f"**Result:** {payload['final_report']}",
            gr.update(visible=False),
            "",
            gr.update(visible=False),
        )

    report_md = "\n\n---\n\n".join(
        [
            _format_final_report(payload["final_report"]),
            _format_network_a(payload["network_a_disease_classifier"]),
            _format_network_b(payload["network_b_lesion_segmentation"]),
        ]
    )
    network_c_md = _format_network_c(payload["network_c_report_generation"])
    json_path = _write_payload_json(payload)

    return (
        report_md,
        gr.update(value=state.heatmap_fig, visible=True),
        network_c_md,
        gr.update(value=json_path, visible=True),
    )


with gr.Blocks(title="ChatCAD — Chest X-ray Report Assistant") as demo:
    gr.Markdown("# ChatCAD — Chest X-ray Report Assistant")

    with gr.Row():
        with gr.Column():
            image_input = gr.File(
                label="Chest X-ray (raster image or DICOM)",
                file_types=[".dcm", ".jpg", ".jpeg", ".png"],
                type="filepath",
            )
            image_preview = gr.Image(
                label="Chest X-ray (raster image or DICOM)", interactive=False, visible=False
            )
            change_image_button = gr.Button("Change image", size="sm", visible=False)
            gr.Examples(examples=[["imgs/examples/chest.jpg"]], inputs=[image_input])
            provider_dropdown = gr.Dropdown(
                list(PROVIDER_MAP.keys()), label="LLM Provider", value="Claude"
            )
            api_key_box = gr.Textbox(label="API Key", type="password")
            run_button = gr.Button("Run Analysis", variant="primary")

        with gr.Column():
            report_output = gr.Markdown(min_height=200)
            heatmap_output = gr.Plot(label="Lesion Segmentation Heatmaps", visible=False)
            network_c_output = gr.Markdown()
            download_button = gr.DownloadButton("Download Report (JSON)", visible=False)

    image_input.change(
        fn=handle_upload,
        inputs=[image_input],
        outputs=[image_input, image_preview, change_image_button],
    )
    change_image_button.click(
        fn=clear_upload,
        outputs=[image_input, image_preview, change_image_button],
    )

    run_button.click(
        fn=start_loading,
        inputs=[image_input, provider_dropdown, api_key_box],
        outputs=[report_output, heatmap_output, network_c_output, download_button],
    ).then(
        fn=analyze,
        inputs=[image_input, provider_dropdown, api_key_box],
        outputs=[report_output, heatmap_output, network_c_output, download_button],
    )


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