""" app.py ------ Gradio chatbot demo for the trained CXR-VLM. Flow ---- 1. Upload a chest X-ray image. 2. The model auto-generates a report — Findings + Impression (cascade, matching report_mode=split_cascade) — as the assistant's first message. 3. Ask any follow-up question about the same image; it's answered VQA-style, the way the model was trained. The image stays "loaded" for the whole conversation; uploading a new image resets the chat and produces a fresh report. Usage (Colab / cloud GPU — auto-download the checkpoint from the HF runs repo) ----------------------------------------------------------------------------- python -m demo.app \ --hf_repo hieu3636/cxr-vlm-runs \ --run_id MIMIC-CXR_resized_run_3 \ --share # public Gradio link (needed on Colab) Usage (local — checkpoint already on disk) ------------------------------------------ python -m demo.app \ --checkpoint checkpoints/IU-Xray_run_1/stage2_instruct # dir OR .pt Notes ----- * On Colab pass --share so you get a public URL (the in-VM port isn't reachable otherwise). * --hf_token is optional if the runs repo is public or $HF_TOKEN is set. """ import argparse import sys from pathlib import Path sys.path.insert(0, str(Path(__file__).resolve().parents[1])) from demo.cxr_demo import CXRChatbot, download_checkpoint_from_hf WELCOME = ( "Upload a chest X-ray to begin. CXR-HIEU writes a radiology report for it, " "then you can ask any follow-up question about the same image." ) def parse_args(): p = argparse.ArgumentParser(description="CXR-VLM chatbot demo (Gradio)") # Checkpoint source — either a local path OR an HF runs-repo coordinate. p.add_argument("--checkpoint", type=str, default=None, help="Local checkpoint dir (contains checkpoint_projection.pt " "+ checkpoint_lora/) or a stage2_final.pt path. " "If omitted, --hf_repo + --run_id are used to download.") p.add_argument("--hf_repo", type=str, default="hieu3636/cxr-vlm-runs", help="HF runs repo to pull the checkpoint from.") p.add_argument("--run_id", type=str, default="MIMIC-CXR_resized_run_3", help="Run id on the HF runs repo, e.g. MIMIC-CXR_resized_run_3.") p.add_argument("--stage", type=str, default="stage2", choices=["stage2", "stage1"], help="Which stage's checkpoint to load (default stage2).") p.add_argument("--which", type=str, default="best", choices=["best", "last"], help="best (lowest eval_loss) or last (default best).") p.add_argument("--hf_token", type=str, default=None, help="HF token (defaults to $HF_TOKEN). Needed for private repos.") p.add_argument("--model_config", type=str, default="configs/model_config.yaml") p.add_argument("--device", type=str, default="auto", choices=["auto", "cuda", "cpu"]) p.add_argument("--max_new_tokens", type=int, default=300) p.add_argument("--share", action="store_true", help="Create a public Gradio link (use this on Colab).") p.add_argument("--server_name", type=str, default="0.0.0.0") p.add_argument("--server_port", type=int, default=7860) return p.parse_args() def resolve_checkpoint(args) -> str: if args.checkpoint: return args.checkpoint if not args.run_id: raise SystemExit( "Provide --checkpoint , OR --run_id (with --hf_repo) " "to auto-download from the HF runs repo." ) print(f"[app] downloading {args.run_id}/{args.stage}/{args.which} " f"from {args.hf_repo} …") return download_checkpoint_from_hf( repo_id = args.hf_repo, run_id = args.run_id, stage = args.stage, which = args.which, token = args.hf_token, ) def build_ui(bot: CXRChatbot, max_new_tokens: int): import demo._gradio_compat # noqa: F401 (fix gradio_client schema crash) import gradio as gr with gr.Blocks(title="CXR-HIEU Demo", theme=gr.themes.Soft()) as demo: gr.Markdown( "# 🩻 CXR-HIEU — Chest X-ray Report & Q&A\n" "Upload a chest X-ray. CXR-HIEU generates a radiology report and " "answers questions about the image." ) # Holds the active conversation's session dict {img, pnu}. The image # (and its predicted PNU) persist across follow-up questions; a new # upload replaces it → the bot answers about the new image. sess_state = gr.State(value=None) with gr.Row(): with gr.Column(scale=1): image_in = gr.Image(type="pil", label="Chest X-ray image", height=380) analyze_btn = gr.Button("🔍 Analyze image", variant="primary") gr.Markdown( "_Upload an image and click **Analyze image** (it also " "analyzes automatically on upload). Then ask questions on " "the right._" ) with gr.Column(scale=2): chatbot = gr.Chatbot(label="Conversation", height=460, type="messages") with gr.Row(): msg = gr.Textbox( placeholder="Ask about this image, e.g. Is there pleural effusion?", scale=5, show_label=False, container=False, ) send_btn = gr.Button("Send", scale=1, variant="primary") clear_btn = gr.Button("🗑️ Clear conversation") # ── handlers ─────────────────────────────────────────────────────── def on_analyze(image): """New image → new session, reset chat, generate the report.""" if image is None: return None, [{ "role": "assistant", "content": "⚠️ Please upload a chest X-ray first.", }] session = bot.prepare(image) report = bot.generate_report(session, max_new_tokens=max_new_tokens) history = [{"role": "assistant", "content": report}] return session, history def on_send(message, history, session): history = history or [] if not message or not message.strip(): return "", history if session is None: history.append({"role": "user", "content": message}) history.append({ "role": "assistant", "content": "⚠️ Please upload an image and click " "**Analyze image** before asking questions.", }) return "", history answer = bot.answer_question(session, message.strip(), max_new_tokens=max_new_tokens) history.append({"role": "user", "content": message}) history.append({"role": "assistant", "content": answer}) return "", history def on_clear(): return None, None, "" # Auto-analyze on upload AND via the button. image_in.upload(on_analyze, inputs=image_in, outputs=[sess_state, chatbot]) analyze_btn.click(on_analyze, inputs=image_in, outputs=[sess_state, chatbot]) send_btn.click(on_send, inputs=[msg, chatbot, sess_state], outputs=[msg, chatbot]) msg.submit(on_send, inputs=[msg, chatbot, sess_state], outputs=[msg, chatbot]) clear_btn.click(on_clear, outputs=[sess_state, chatbot, msg]) gr.Markdown(f"\n\n_{WELCOME}_") return demo def main(): args = parse_args() ckpt = resolve_checkpoint(args) bot = CXRChatbot( checkpoint_dir = ckpt, model_config = args.model_config, device = args.device, ) demo = build_ui(bot, args.max_new_tokens) demo.queue().launch( share = args.share, server_name = args.server_name, server_port = args.server_port, ) if __name__ == "__main__": main()