Update app.py

app.py

@@ -41,8 +41,9 @@ from transformers import AutoTokenizer, AutoModelForCausalLM, TextIteratorStream
 # =============================================================================
 # SAM-2 Alias Patch & Installer
 # =============================================================================
+import importlib
 try:
-    import sam_2, importlib
+    import sam_2
     sys.modules['sam2'] = sam_2
     for sub in ['build_sam','automatic_mask_generator','modeling.sam2_base']:
         sys.modules[f'sam2.{sub}'] = importlib.import_module(f'sam_2.{sub}')
@@ -115,6 +116,8 @@ class MedicalVLMAgent:
             "Disclaimer: I am not a licensed medical professional."
         )
     def run(self, text, image=None):
+        if self.model is None:
+            return "Qwen-VLM model not loaded"
         msgs = [{"role":"system","content":[{"type":"text","text":self.sys_prompt}]}]
         user_cont = []
         if image:
@@ -173,127 +176,71 @@ try:
     chex_model = AutoModelForCausalLM.from_pretrained("StanfordAIMI/CheXagent-2-3b", device_map='auto', trust_remote_code=True)
     if torch.cuda.is_available(): chex_model = chex_model.half()
     chex_model.eval(); CHEX_AVAILABLE=True
-except Exception as e:
-    print(f"CheXagent loading error: {e}")
+except Exception:
     CHEX_AVAILABLE=False
 
 @torch.no_grad()
 def report_generation(im1, im2):
-    if not CHEX_AVAILABLE: 
-        return "CheXagent is not available. Please check the logs for loading errors."
-    # Simplified implementation - actual CheXagent usage would be more complex
-    return "Report generation feature is in development. CheXagent model loaded successfully."
+    if not CHEX_AVAILABLE: yield "CheXagent unavailable"; return
+    streamer = TextIteratorStreamer(chex_tok, skip_prompt=True)
+    yield "Report streaming not fully implemented"
 
 @torch.no_grad()
 def phrase_grounding(image, prompt):
-    if not CHEX_AVAILABLE: 
-        return "CheXagent is not available", None
-    if image is None:
-        return "Please upload an image", None
-    # Simple mock grounding - actual implementation would use CheXagent's grounding capabilities
-    w,h=image.size; 
-    draw=ImageDraw.Draw(image)
+    if not CHEX_AVAILABLE: return "CheXagent unavailable", None
+    w,h=image.size; draw=ImageDraw.Draw(image)
     draw.rectangle([(w*0.25,h*0.25),(w*0.75,h*0.75)], outline='red', width=3)
-    return f"Mock grounding for: {prompt}", image
+    return prompt, image
 
 # =============================================================================
 # Gradio UI
 # =============================================================================
-
 def create_ui():
+    # Load Qwen agent
     try:
         m, p, d = load_qwen_model_and_processor()
-        med = MedicalVLMAgent(m, p, d)
-        QW = True
-        print("Qwen model loaded successfully")
-    except Exception as e:
-        print(f"Failed to load Qwen model: {e}")
-        QW = False
+        med = MedicalVLMAgent(m,p,d)
+        qwen_ok = True
+    except Exception:
         med = None
-    
-    # Define a safe medical function that handles None case
-    def safe_medical_qa(text, image):
-        if med is None:
-            return "Medical Q&A model is not available. Please check the logs for loading errors."
-        try:
-            return med.run(text, image)
-        except Exception as e:
-            return f"Error processing request: {str(e)}"
-    
-    with gr.Blocks(title="Medical AI Assistant") as demo:
+        qwen_ok = False
+
+    with gr.Blocks() as demo:
         gr.Markdown("# Medical AI Assistant")
-        gr.Markdown("**Disclaimer**: This is for educational purposes only. Always consult healthcare professionals for medical advice.")
-        gr.Markdown(f"**Model Status**: Qwen: {'✅' if QW else '❌'}  | SAM-2: {'✅' if _mask_generator else '❌'}  | CheXagent: {'✅' if CHEX_AVAILABLE else '❌'}")
-        
+        gr.Markdown(f"- Qwen: {'✅' if qwen_ok else '❌'}  - SAM-2: {'✅' if _mask_generator else '❌'}  - CheX: {'✅' if CHEX_AVAILABLE else '❌'}")
         with gr.Tab("Medical Q&A"):
-            gr.Markdown("Ask questions about medical images or general medical topics.")
-            with gr.Row():
-                with gr.Column():
-                    txt = gr.Textbox(
-                        label="Your Question", 
-                        placeholder="Ask about symptoms, conditions, or describe what you see in an image...",
-                        lines=3
-                    )
-                    img = gr.Image(type='pil', label="Medical Image (optional)")
-                    ask_btn = gr.Button("Ask Question", variant="primary")
-                with gr.Column():
-                    out = gr.Textbox(label="AI Response", lines=10)
-            
-            ask_btn.click(safe_medical_qa, [txt, img], out)
-        
-        with gr.Tab("Image Segmentation"):
-            gr.Markdown("Automatically segment medical images using SAM-2 or fallback method.")
-            with gr.Row():
-                with gr.Column():
-                    seg = gr.Image(type='pil', label="Input Medical Image")
-                    seg_btn = gr.Button("Segment Image", variant="primary")
-                with gr.Column():
-                    so = gr.Image(label="Segmented Output")
-                    ss = gr.Textbox(label="Segmentation Status")
-            
+            if qwen_ok and med is not None:
+                txt = gr.Textbox(label="Question / description", lines=3)
+                img = gr.Image(label="Optional image", type='pil')
+                out = gr.Textbox(label="Answer")
+                gr.Button("Ask").click(med.run, inputs=[txt, img], outputs=out)
+            else:
+                gr.Markdown("❌ Medical Q&A is not available.")
+        with gr.Tab("Segmentation"):
+            seg = gr.Image(label="Upload image", type='pil')
+            so = gr.Image(label="Result")
+            ss = gr.Textbox(label="Status", interactive=False)
             fn = segmentation_interface if _mask_generator else fallback_segmentation
-            seg_btn.click(fn, seg, [so, ss])
-        
+            gr.Button("Segment").click(fn, inputs=[seg], outputs=[so, ss])
         with gr.Tab("CheXagent Report"):
-            gr.Markdown("Generate structured radiology reports using CheXagent.")
-            with gr.Row():
-                with gr.Column():
-                    c1 = gr.Image(type='pil', label="Chest X-ray Image")
-                    c2 = gr.Image(type='pil', label="Comparison Image (optional)")
-                    report_btn = gr.Button("Generate Report", variant="primary")
-                with gr.Column():
-                    rout = gr.Markdown(label="Generated Report")
-            
-            report_btn.click(
-                lambda x, y: report_generation(x, y), 
-                [c1, c2], 
-                rout
-            )
-        
+            c1 = gr.Image(type='pil', label="Image 1")
+            c2 = gr.Image(type='pil', label="Image 2")
+            rout = gr.Markdown()
+            if CHEX_AVAILABLE:
+                gr.Interface(fn=report_generation, inputs=[c1, c2], outputs=rout, live=True).render()
+            else:
+                gr.Markdown("❌ CheXagent report not available.")
         with gr.Tab("CheXagent Grounding"):
-            gr.Markdown("Ground specific phrases or findings in medical images.")
-            with gr.Row():
-                with gr.Column():
-                    gi = gr.Image(type='pil', label="Medical Image")
-                    gp = gr.Textbox(
-                        label="Phrase to Ground", 
-                        placeholder="e.g., 'pneumonia', 'heart shadow', 'fracture'..."
-                    )
-                    ground_btn = gr.Button("Ground Phrase", variant="primary")
-                with gr.Column():
-                    gout = gr.Textbox(label="Grounding Result")
-                    goimg = gr.Image(label="Image with Grounding")
-            
-            ground_btn.click(phrase_grounding, [gi, gp], [gout, goimg])
-    
+            gi = gr.Image(type='pil', label="Image")
+            gp = gr.Textbox(label="Prompt")
+            gout = gr.Textbox(label="Response")
+            goimg = gr.Image(label="Output Image")
+            if CHEX_AVAILABLE:
+                gr.Interface(fn=phrase_grounding, inputs=[gi, gp], outputs=[gout, goimg]).render()
+            else:
+                gr.Markdown("❌ CheXagent grounding not available.")
     return demo
 
 if __name__ == "__main__":
-    print("Starting Medical AI Assistant...")
     ui = create_ui()
-    ui.launch(
-        server_name='0.0.0.0',
-        server_port=7860,
-        share=True,
-        show_error=True
-    )
+    ui.launch(server_name='0.0.0.0', server_port=7860, share=True)