Changed UI

app.py

@@ -1,54 +1,63 @@
 import gradio as gr
 from src.pipeline import generate_report
 
-def process_upload(image):
-    """Process uploaded X-ray image and return radiology report"""
-    if image is None:
-        return None, "Please upload a chest X-ray image."
-    
+# ------------------------------------------------------------------
+# 1. Pre-load models on Space start-up
+# ------------------------------------------------------------------
+print("Pre-loading models for fast inference …")
+try:
+    from src.tools_loader import get_tools          # downloads BiomedCLIP + SPECTER-2
+    _ = get_tools()
+    print("Models pre-loaded successfully!")
+except Exception as e:
+    print(f"Model pre-loading failed: {e}")
+
+# ------------------------------------------------------------------
+# 2. Inference wrapper
+# ------------------------------------------------------------------
+def process_upload(image_path: str):
+    """Run the multi-agent pipeline on an uploaded chest X-ray."""
+    if image_path is None:
+        return "Please upload a chest X-ray image."
+
     try:
-        report = generate_report(image)
-        return image, report  # Return both image and report
+        report = generate_report(image_path)
+        return report
     except Exception as e:
-        return image, f"Error processing image: {str(e)}"
+        return f"Error processing image: {e}"
 
-# Create enhanced Gradio interface with custom layout
+# ------------------------------------------------------------------
+# 3. Gradio UI
+# ------------------------------------------------------------------
 with gr.Blocks(title="Multi-Agent Radiology Assistant") as demo:
-    gr.Markdown("# Multi-Agent Radiology Assistant")
-    gr.Markdown("Upload a chest X-ray image to receive an AI-generated radiology report using cutting-edge, multi-agent analysis")
-    
-    with gr.Row():
-        with gr.Column(scale=1):
-            input_image = gr.Image(
-                type="filepath", 
-                label="Upload Chest X-ray",
-                height=400
-            )
-            process_btn = gr.Button("Generate Report", variant="primary")
-        
-        with gr.Column(scale=1):
-            output_image = gr.Image(
-                label="Analyzed Image",
-                height=400,
-                interactive=False
-            )
-    
-    with gr.Row():
-        output_report = gr.Markdown(
-            label="📄 Radiology Report",
-            height=300
+    gr.Markdown(
+        """
+        # Multi-Agent Radiology Assistant  
+        Upload a chest X-ray and receive an AI-generated report produced by a multi-agent pipeline.
+        """
+    )
+
+    # --- Upload widget + button ------------------------------------------------
+    with gr.Column():
+        input_image = gr.Image(
+            type="filepath",
+            label="Upload Chest X-ray",
+            height=400
         )
-    
-    # Connect the button to processing function
+        process_btn = gr.Button("Generate Report", variant="primary")
+
+    # --- Report output ---------------------------------------------------------
+    output_report = gr.Markdown(label="Radiology Report", show_label=True)
+
+    # --- Wire everything together ---------------------------------------------
     process_btn.click(
         fn=process_upload,
-        inputs=[input_image],
-        outputs=[output_image, output_report]
+        inputs=input_image,
+        outputs=output_report
     )
-    
-    # Add examples section
-    gr.Markdown("### Example X-rays")
-    gr.Markdown("*Upload your own chest X-ray image above to get started*")
 
+    gr.Markdown("### Need an example?  \nUse any frontal CXR PNG file and click **Generate Report**.")
+
+# ------------------------------------------------------------------
 if __name__ == "__main__":
     demo.launch()