Update app.py

app.py

@@ -1,178 +1,70 @@
-# Remove theme and css from here
-with gr.Blocks() as demo: 
-    gr.HTML("<div id='header'><h1 style='color:white; font-family:Consolas;'>🧬 MediVance LLC | NerdMedica Council</h1></div>")
-    # ... rest of your UI code ...
-    import gradio as gr
-from loader import ModelLoader
-from inference import generate_maira2_report, generate_biomedclip_heatmap, overlay_medgemma_bboxes
-from tutor import generate_vqa_response
-import pandas as pd
+import gradio as gr
 import os
-import datetime
-from dotenv import load_dotenv
-
-# Security: Ensure dotenv keys are safely loaded early
-load_dotenv()
-
-# --- Initialization ---
-print("Initializing Triple-Model NerdMedica Platform...")
-loader = ModelLoader()
-# Pre-load MedGemma 1.5 4B and BiomedCLIP into VRAM
-loader.init_startup_models()
+from loader import loader
+from tutor import run_council_deliberation
+from inference import generate_universal_heatmap
 
-AUDIT_LOG_FILE = "audit_log.csv"
-if not os.path.exists(AUDIT_LOG_FILE):
-    pd.DataFrame(columns=["Timestamp", "Modality", "Audit_Action", "Notes"]).to_csv(AUDIT_LOG_FILE, index=False)
-
-def log_audit_action(modality, action, notes):
-    if not action:
-        return "Please select an Audit Decision."
-    try:
-        df = pd.read_csv(AUDIT_LOG_FILE)
-        new_entry = {
-            "Timestamp": datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S"),
-            "Modality": modality,
-            "Audit_Action": action,
-            "Notes": notes
-        }
-        df = pd.concat([df, pd.DataFrame([new_entry])], ignore_index=True)
-        df.to_csv(AUDIT_LOG_FILE, index=False)
-        return gr.update(value=f"Successfully logged from Doctor-in-the-Loop UI:\nDecision: {action}\nNotes: {notes}")
-    except Exception as e:
-        return gr.update(value=f"Log Error: {e}")
-
-def process_modality_routing(image, modality):
-    """Routes the image to MAIRA-2 if X-Ray, else assigns text fallback."""
-    if image is None:
-        return "No image uploaded. Operating in text-only Clinical Generalist mode. Ask your question in the Socratic Tutor."
-        
-    if modality == "Chest X-Ray" or modality == "X-Ray":
-        loader.load_maira2_lazy()
-        maira2_model, maira2_processor = loader.get_maira2()
-        report = generate_maira2_report(maira2_model, maira2_processor, image)
-        
-        # Offload after generation to save memory for subsequent MedGemma/BiomedCLIP calls
-        loader.clear_vram() 
-        return report
-    else:
-        return f"MAIRA-2 is specialized for Chest X-Rays.\n\nFor {modality}, please consult the MedGemma Socratic Tutor directly for an interactive VQA assessment."
+# NerdMedica Dark Tech CSS
+CSS = """
+body, .gradio-container { background-color: #0A0E1A !important; color: #D1D5DB !important; }
+.gr-button-primary { background: linear-gradient(90deg, #E5534B, #A855F7) !important; border: none !important; border-radius: 8px !important; color: white !important; font-weight: bold !important; }
+#header { border-bottom: 2px solid #E5534B; padding-bottom: 15px; margin-bottom: 25px; }
+.message.user { background-color: #1F2937 !important; border-left: 4px solid #E5534B !important; }
+.message.bot { background-color: #111827 !important; border-left: 4px solid #3B82F6 !important; }
+"""
 
-def generate_vqa_wrapper(message, history, modality, image):
-    """Wrapper that passes MedGemma models to the tutor generator."""
-    medgemma_model, medgemma_tokenizer = loader.get_medgemma()
+def medical_audit_pipeline(image, modality, question, history):
+    loader.clear_vram()
     
-    for chunk in generate_vqa_response(medgemma_model, medgemma_tokenizer, message, history, modality, image):
-        yield chunk
-
-def generate_visual_audit(image, query, current_chat_history):
-    """Generates the BiomedCLIP heatmap, and optionally overlays MedGemma boxes from chat history."""
-    if image is None:
-        return None
-        
-    biomed_model, biomed_preprocess, biomed_tokenizer = loader.get_biomedclip()
-    heatmap_pil = generate_biomedclip_heatmap(biomed_model, biomed_preprocess, biomed_tokenizer, image, query)
+    # 1. Visual Evidence (If image exists)
+    heatmap = None
+    if image is not None:
+        clip_model, preprocess = loader.load_biomed_clip()
+        heatmap = generate_universal_heatmap(image, question, clip_model, preprocess)
     
-    # Attempt to draw MedGemma bboxes if present in the last assistant response
-    if current_chat_history and len(current_chat_history) > 0:
-        last_item = current_chat_history[-1]
-        last_response = last_item.get("content", "") if isinstance(last_item, dict) else ""
-        if last_response and "<loc" in last_response:
-            heatmap_pil = overlay_medgemma_bboxes(heatmap_pil, last_response)
-            
-    return heatmap_pil
-
-custom_css = """
-body, .gradio-container { background-color: #0A0E1A !important; color: #D1D5DB !important; font-family: 'Inter', sans-serif; }
-.header-title { border-bottom: 2px solid #E5534B; padding-bottom: 10px; color: #D1D5DB !important; }
-h1, h2, h3, p, span, label, .markdown-text { color: #D1D5DB !important; }
-.primary-btn { background: linear-gradient(90deg, #E5534B, #A855F7) !important; border: none !important; color: white !important; font-weight: bold !important; transition: all 0.3s ease; }
-.primary-btn:hover { opacity: 0.9 !important; transform: translateY(-1px); }
-.shadow-card { background-color: #111827 !important; border: 1px solid #1F2937 !important; border-radius: 8px !important; box-shadow: 0 4px 6px -1px rgba(0, 0, 0, 0.6) !important; padding: 15px; margin-bottom: 10px; }
-.textbox textarea { background-color: #1F2937 !important; color: #D1D5DB !important; border: 1px solid #374151 !important; border-radius: 6px !important; }
-"""
+    # 2. Council Deliberation
+    response = run_council_deliberation(image, modality, question)
+    
+    # Append to Chat History
+    history.append({"role": "user", "content": question})
+    history.append({"role": "assistant", "content": response})
+    
+    return history, heatmap, response
 
-with gr.Blocks(theme=gr.themes.Base(), css=custom_css) as demo:
-    gr.Markdown("# <span class='header-title'>NerdMedica Diagnostics Hub</span>")
-    gr.Markdown("Orchestrating **MAIRA-2** (Radiology), **MedGemma 1.5 4B** (Council of Experts), and **BiomedCLIP** (Zero-Shot Visual Audit).")
+# Gradio 6.0: We build the blocks without passing theme/css here
+with gr.Blocks() as demo:
+    gr.HTML("<div id='header'><h1 style='color:white; font-family:Consolas;'>🧬 MediVance LLC | NerdMedica Council</h1></div>")
     
     with gr.Row():
-        with gr.Column(scale=1, elem_classes="shadow-card"):
-            modality_dropdown = gr.Dropdown(choices=["Chest X-Ray", "CT Scan", "MRI"], value="Chest X-Ray", label="Select Modality")
-            image_input = gr.Image(type="pil", label="Upload Medical Scan")
-            analyze_btn = gr.Button("1. Generate MAIRA-2 Radiology Draft", elem_classes="primary-btn")
-            
-            gr.Markdown("### Audit Visualizer (BiomedCLIP & MedGemma)")
-            visual_output = gr.Image(type="pil", label="Visual Evidence (Heatmap + BBoxes)")
-            visual_query = gr.Textbox(label="BiomedCLIP Visual Query", placeholder="e.g., 'Consolidation' or 'Pleural Effusion'")
-            visual_btn = gr.Button("2. Generate Visual Confirmation", elem_classes="primary-btn")
-            
-            with gr.Accordion("Doctor-In-The-Loop Audit Logging", open=False):
-                audit_decision = gr.Radio(["Approved", "Corrected", "Flagged (Hallucination)"], label="Human Audit Decision")
-                audit_notes = gr.Textbox(label="Corrections / Notes")
-                log_btn = gr.Button("Save to Audit Log", elem_classes="primary-btn")
-                log_status = gr.Markdown("")
-                
-        with gr.Column(scale=1, elem_classes="shadow-card"):
-            gr.Markdown("### Clinical Evaluation Draft")
-            draft_output = gr.Textbox(lines=25, interactive=True, label="Radiology / Clinical Draft (Editable)")
+        with gr.Column(scale=1):
+            img_in = gr.Image(type="pil", label="Medical Scan (Optional)")
+            modality_drop = gr.Dropdown(
+                ["General Inquiry", "Chest X-Ray", "CT Scan", "MRI", "Pathology"], 
+                value="General Inquiry", 
+                label="Modality Focus"
+            )
+            heatmap_out = gr.Image(label="Visual Evidence Audit")
             
-        with gr.Column(scale=1, elem_classes="shadow-card"):
-            gr.Markdown("### Council of Experts (Socratic Tutor)")
-            chatbot = gr.Chatbot(type="messages", height=450, label="Senior Clinical Generalist")
-            msg = gr.Textbox(label="Ask the Council (e.g., 'What justifies this diagnosis?')")
-            consult_btn = gr.Button("Consult the Council", elem_classes="primary-btn")
-            clear = gr.ClearButton([msg, chatbot])
-
-    # Event Wiring
-    analyze_btn.click(
-        process_modality_routing,
-        inputs=[image_input, modality_dropdown],
-        outputs=[draft_output]
-    )
-    
-    def user(user_message, history):
-        if history is None:
-            history = []
-        if user_message.strip() != "":
-            history.append({"role": "user", "content": user_message})
-        return "", history
+        with gr.Column(scale=2):
+            chat = gr.Chatbot(label="Council Deliberation", type="messages", height=500)
+            query = gr.Textbox(placeholder="Ask about 'Isolated Hypertension' or specific scan findings...", label="Doctor's Inquiry")
+            btn = gr.Button("Consult the Council", variant="primary")
 
-    def bot(history, modality, image):
-        if history is None or len(history) == 0:
-            yield history
-            return
-        user_message = history[-1]["content"]
-        bot_generator = generate_vqa_wrapper(user_message, history[:-1], modality, image)
-        history.append({"role": "assistant", "content": ""})
-        for chunk in bot_generator:
-            history[-1]["content"] = chunk
-            yield history
+    with gr.Accordion("🩺 Official Clinical Audit Report", open=False):
+        audit_box = gr.Textbox(label="Final Finding (Editable)", interactive=True, lines=4)
+        status = gr.Radio(["Approved", "Corrected", "Flagged"], label="Clinical Signature")
+        gr.Button("Archive to MediVance API", variant="secondary")
 
-    # Link both text box submission and "Consult the Council" button to the same pipeline
-    msg.submit(user, [msg, chatbot], [msg, chatbot], queue=False).then(
-        bot, [chatbot, modality_dropdown, image_input], chatbot
-    )
-    consult_btn.click(user, [msg, chatbot], [msg, chatbot], queue=False).then(
-        bot, [chatbot, modality_dropdown, image_input], chatbot
-    )
-    
-    visual_btn.click(
-        generate_visual_audit,
-        inputs=[image_input, visual_query, chatbot],
-        outputs=[visual_output]
-    )
-    
-    log_btn.click(
-        log_audit_action,
-        inputs=[modality_dropdown, audit_decision, audit_notes],
-        outputs=[log_status]
+    # Wire up the button
+    btn.click(
+        fn=medical_audit_pipeline, 
+        inputs=[img_in, modality_drop, query, chat], 
+        outputs=[chat, heatmap_out, audit_box]
     )
 
-if __name__ == "__main__":
-    demo.launch(ssr_mode=False)
-# Pass them into launch() instead
+# Gradio 6.0: Theme and CSS MUST be passed in the launch method
 demo.launch(
-    ssr_mode=False, 
-    theme=gr.themes.Base(), 
+    ssr_mode=False,
+    theme=gr.themes.Base(),
     css=CSS
 )