app.py

import gradio as gr
import time
from src.pipeline import generate_report

# ------------------------------------------------------------------
# 1. Pre-load models
# ------------------------------------------------------------------
from src.tools_loader import get_tools
_ = get_tools()

# ------------------------------------------------------------------
# 2. Helper: streaming generator with progress
# ------------------------------------------------------------------
def process_upload(image_path: str):
    """
    Streamed generator: yields loading states then final report.
    Gradio shows spinner automatically during execution.
    """
    if image_path is None:
        yield "**Please upload a chest X-ray image to begin analysis.**"
        return

    start = time.time()
    
    # Show loading state immediately
    yield "**Analyzing X-ray image...**\n\nThis will take a few seconds..."
    
    # Generate the actual report
    report = generate_report(image_path)
    
    elapsed = time.time() - start
    
    # Return final formatted report
    yield f"""### Radiology Report

{report}

---
*Report generated in {elapsed:.1f} seconds*"""

# ------------------------------------------------------------------
# 3. Gradio UI - Vertical Layout
# ------------------------------------------------------------------
with gr.Blocks(
    theme=gr.themes.Soft(),
    title="Multi-Agent Radiology Assistant",
    css="""
    .image-container { max-width: 600px; margin: 0 auto; }
    .report-container { margin-top: 40px; padding-top: 20px; }
    .generate-btn { margin: 30px auto; display: block; }
    .progress-bar { z-index: 1000 !important; position: relative; }
    .gradio-container .wrap { z-index: auto; }
    """
) as demo:
    
    # Header
    gr.Markdown(
        "# Multi-Agent Radiology Assistant\n"
        "Upload a chest X-ray image to receive an AI-powered radiology report"
    )
    
    # Image upload section (centered, top)
    with gr.Column():
        input_image = gr.Image(
            type="filepath",
            label="Upload Chest X-ray Image",
            height=400,
            elem_classes=["image-container"]
        )
        
        # Generate button (centered with more spacing)
        generate_btn = gr.Button(
            "Generate Report",
            variant="primary",
            size="lg",
            elem_classes=["generate-btn"]
        )
        
        # Add some spacing to prevent overlap
        gr.HTML("<div style='height: 20px;'></div>")
    
    # Report output section (bottom)
    with gr.Column(elem_classes=["report-container"]):
        output_report = gr.Markdown(
            value="**Ready to analyze**\n\nUpload an X-ray image above and click 'Generate Report' to begin.",
            label="Analysis Results"
        )
    
    # Event handler with progress settings
    generate_btn.click(
        fn=process_upload,
        inputs=input_image,
        outputs=output_report,
        show_progress="full",  # Shows Gradio's built-in progress bar
        concurrency_limit=1   # Prevent multiple simultaneous requests
    )
    
    # Footer with example hint
    gr.Markdown(
        "### Need an example?\n"
        "Use any frontal chest X-ray PNG/JPG file and click **Generate Report**."
    )

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