Update app.py

app.py

@@ -1,7 +1,21 @@
-# Initialize Model Loader globally so it stays in VRAM
+import gradio as gr
+import dotenv
+import os
+from PIL import Image
+
+dotenv.load_dotenv()
+
+from loader import ModelLoader
+from inference import run_inference
+from tutor import generate_socratic_assessment
+
+# Initialize Model Loader
 loader = ModelLoader()
-# Load the model once during startup
-clinical_model = loader.load_model() 
+# Load model once at startup to save VRAM and time
+clinical_model = loader.load_model()
+
+# Global state to hold results between steps
+current_ai_results = {}
 
 def process_clinical_assessment(image):
     if image is None:
@@ -10,11 +24,94 @@ def process_clinical_assessment(image):
             
     yield image, "Running CNN Inference (utilizing pre-loaded model)..."
     
-    # Use the globally loaded model instead of re-loading
+    # Use the globally loaded model
     results = run_inference(clinical_model, image)
-    # ... (rest of your logic)
+    
+    global current_ai_results
+    current_ai_results = results
+    
+    diagnosis_data = {
+        "modality": "X-Ray",
+        "probabilities": results["all_probabilities"]
+    }
+    
+    yield image, "Generating Medical Tutor Socratic Feedback..."
+    
+    stream_generator = generate_socratic_assessment(diagnosis_data)
+    
+    for partial_text in stream_generator:
+        yield image, partial_text
+
+def reveal_ai_analysis():
+    global current_ai_results
+    if not current_ai_results:
+        return None, "No results yet. Please run the Clinical Assessment first.", {}
+            
+    heatmap = current_ai_results.get("heatmap_image", None)
+    confidence = current_ai_results.get("confidence", 0.0)
+    probs = current_ai_results.get("all_probabilities", {})
+    
+    top_diagnosis = current_ai_results.get("top_diagnosis", "Unknown")
+    confidence_text = f"## Top Diagnosis: **{top_diagnosis}** ({confidence*100:.1f}%)\n\nReview the heatmap to audit for spatial mismatch."
+    
+    return heatmap, confidence_text, probs
+
+def log_audit(audit_status, user_notes):
+    if not audit_status:
+        return "Please select an Audit Status."
+    return f"Audit Logged successfully!\n\nStatus: {audit_status}\nNotes: {user_notes}"
+
+with gr.Blocks(theme=gr.themes.Soft(), title="AI-VECINNA") as demo:
+    gr.Markdown("# AI-VECINNA: Dual-Model Medical Auditing (Powered by Local MedGemma)")
+    gr.Markdown("**Status: Running on NVIDIA T4 GPU (Persistent)**")
+    
+    with gr.Accordion("Step 1: Clinical Assessment & Tutor", open=True):
+        with gr.Row():
+            with gr.Column():
+                image_input = gr.Image(type="pil", label="Upload Chest X-Ray")
+                analyze_btn = gr.Button("Analyze & Generate Tutor Scenario", variant="primary")
+            with gr.Column():
+                gr.Markdown("### MedGemma Socratic Tutor Feedback")
+                scenario_output = gr.Markdown("Tutor instructions will appear here.")
+                human_hypothesis = gr.Textbox(label="Record your initial differential diagnosis here...", lines=3)
+
+    with gr.Accordion("Step 2: AI Reveal & Audit", open=False):
+        reveal_btn = gr.Button("Reveal AI Analysis", variant="secondary")
+        with gr.Row():
+            with gr.Column():
+                heatmap_output = gr.Image(type="pil", label="GradCAM AI Heatmap")
+            with gr.Column():
+                confidence_output = gr.Markdown("Confidence results here.")
+                probs_output = gr.Label(label="Full Probability Distribution")
+
+        audit_radio = gr.Radio(
+            label="AI Safety Audit",
+            choices=[
+                "AI Verified: Heatmap & Confidence clinically align.",
+                "AI Flagged: Spatial mismatch (Hallucination).",
+                "AI Flagged: Low confidence/over-reliance risk."
+            ]
+        )
+        audit_notes = gr.Textbox(label="Audit Notes", lines=2)
+        audit_btn = gr.Button("Submit Safety Audit", variant="primary")
+        audit_result = gr.Textbox(label="Status", interactive=False)
+
+    analyze_btn.click(
+        fn=process_clinical_assessment,
+        inputs=[image_input],
+        outputs=[image_input, scenario_output]
+    )
+    
+    reveal_btn.click(
+        fn=reveal_ai_analysis,
+        outputs=[heatmap_output, confidence_output, probs_output]
+    )
+    
+    audit_btn.click(
+        fn=log_audit,
+        inputs=[audit_radio, audit_notes],
+        outputs=audit_result
+    )
 
-# ... inside your demo block ...
 if __name__ == "__main__":
-    # Disable SSR for maximum stability on Hugging Face
     demo.launch(ssr_mode=False)