Update app.py

app.py

@@ -1,4 +1,4 @@
-##!/usr/bin/env python
+#!/usr/bin/env python
 # -*- coding: utf-8 -*-
 """
 Combined Medical-VLM, SAM-2 Automatic Masking, and CheXagent Demo
@@ -173,21 +173,28 @@ 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:
+except Exception as e:
+    print(f"CheXagent loading error: {e}")
     CHEX_AVAILABLE=False
 
 @torch.no_grad()
 def report_generation(im1, im2):
-    if not CHEX_AVAILABLE: yield "CheXagent unavailable"; return
-    streamer = TextIteratorStreamer(chex_tok, skip_prompt=True)
-    yield "Report streaming not fully implemented"
+    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."
 
 @torch.no_grad()
 def phrase_grounding(image, prompt):
-    if not CHEX_AVAILABLE: return "CheXagent unavailable", None
-    w,h=image.size; draw=ImageDraw.Draw(image)
+    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)
     draw.rectangle([(w*0.25,h*0.25),(w*0.75,h*0.75)], outline='red', width=3)
-    return prompt, image
+    return f"Mock grounding for: {prompt}", image
 
 # =============================================================================
 # Gradio UI
@@ -196,23 +203,98 @@ def phrase_grounding(image, prompt):
 def create_ui():
     try:
         m, p, d = load_qwen_model_and_processor()
-        med = MedicalVLMAgent(m,p,d); QW=True
-    except:
-        QW=False; med=None
-    with gr.Blocks() as demo:
+        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 = 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:
         gr.Markdown("# Medical AI Assistant")
-        gr.Markdown(f"- Qwen: {'✅' if QW else '❌'}  - SAM-2: {'✅' if _mask_generator else '❌'}  - CheX: {'✅' if CHEX_AVAILABLE else '❌'}")
+        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 '❌'}")
+        
         with gr.Tab("Medical Q&A"):
-            txt=gr.Textbox(); img=gr.Image(type='pil'); out=gr.Textbox(); gr.Button("Ask").click(med.run,[txt,img],out)
-        with gr.Tab("Segmentation"):
-            seg=gr.Image(type='pil'); so=gr.Image(); ss=gr.Textbox(); fn=segmentation_interface if _mask_generator else fallback_segmentation; gr.Button("Segment").click(fn,seg,[so,ss])
+            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")
+            
+            fn = segmentation_interface if _mask_generator else fallback_segmentation
+            seg_btn.click(fn, seg, [so, ss])
+        
         with gr.Tab("CheXagent Report"):
-            c1=gr.Image(type='pil');c2=gr.Image(type='pil'); rout=gr.Markdown(); gr.Interface(report_generation,[c1,c2],rout,live=True).render()
+            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
+            )
+        
         with gr.Tab("CheXagent Grounding"):
-            gi=gr.Image(type='pil'); gp=gr.Textbox(); gout=gr.Textbox(); goimg=gr.Image(); gr.Interface(phrase_grounding,[gi,gp],[gout,goimg]).render()
+            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])
+    
     return demo
 
 if __name__ == "__main__":
-    ui=create_ui(); ui.launch(server_name='0.0.0.0',server_port=7860,share=True)
-
-
+    print("Starting Medical AI Assistant...")
+    ui = create_ui()
+    ui.launch(
+        server_name='0.0.0.0',
+        server_port=7860,
+        share=True,
+        show_error=True
+    )===============================================================
+# Gradio UI