Update app.py

app.py

@@ -23,6 +23,7 @@ import tempfile
 import subprocess
 import warnings
 from threading import Thread
+from pathlib import Path
 
 # Environment setup
 os.environ["PYTORCH_ENABLE_MPS_FALLBACK"] = "1"
@@ -40,7 +41,6 @@ from transformers import AutoTokenizer, AutoModelForCausalLM, TextIteratorStream
 # =============================================================================
 # SAM-2 Alias Patch & Installer
 # =============================================================================
-# Alias sam_2 package to sam2 namespace
 try:
     import sam_2, importlib
     sys.modules['sam2'] = sam_2
@@ -50,23 +50,20 @@ except ImportError:
     pass
 
 def check_and_install_sam2():
-    """Ensure SAM-2 is installed and aliased as sam2."""
     try:
         from sam2.build_sam import build_sam2
         return True
     except ImportError:
-        # Clone repo
         repo_dir = Path("segment-anything-2")
         if not repo_dir.exists():
             subprocess.run(["git","clone","https://github.com/facebookresearch/segment-anything-2.git"], check=True)
-        # Install editable
-        cwd = os.getcwd()
+        cwd = Path.cwd()
         os.chdir(repo_dir)
         subprocess.run([sys.executable, "-m", "pip", "install", "-e", "."], check=True)
         os.chdir(cwd)
-        # Re-alias
         try:
-            import sam_2, importlib
+            import sam_2
+            importlib.reload(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}')
@@ -75,7 +72,7 @@ def check_and_install_sam2():
             return False
 
 SAM2_AVAILABLE = check_and_install_sam2()
-
+print(f"SAM-2 Available: {SAM2_AVAILABLE}")
 if SAM2_AVAILABLE:
     from sam2.build_sam import build_sam2
     from sam2.automatic_mask_generator import SAM2AutomaticMaskGenerator
@@ -126,9 +123,7 @@ class MedicalVLMAgent:
         user_cont.append({"type":"text","text": text or ""})
         msgs.append({"role":"user","content":user_cont})
         prompt = self.processor.apply_chat_template(msgs, tokenize=False, add_generation_prompt=True)
-        img_in, vid_in = [], []  # assuming no videos
-        inputs = self.processor(text=[prompt], images=img_in, videos=vid_in,
-                                 padding=True, return_tensors='pt').to(self.device)
+        inputs = self.processor(text=[prompt], images=[], videos=[], padding=True, return_tensors='pt').to(self.device)
         out = self.model.generate(**inputs, max_new_tokens=128)
         resp = out[0][inputs.input_ids.shape[1]:]
         return self.processor.decode(resp, skip_special_tokens=True).strip()
@@ -139,15 +134,14 @@ class MedicalVLMAgent:
 _sam2_model, _mask_generator = (None, None)
 if SAM2_AVAILABLE:
     try:
-        # Initialize model
         CKPT="checkpoints/sam2.1_hiera_large.pt"; CFG="configs/sam2.1/sam2.1_hiera_l.yaml"
         os.chdir("segment-anything-2/sam2/sam2")
         _sam2_model = build_sam2(CFG, CKPT, device=get_device(), apply_postprocessing=False)
         _mask_generator = SAM2AutomaticMaskGenerator(_sam2_model)
-    except Exception:
+    except Exception as e:
+        print(f"SAM-2 init error: {e}")
         _mask_generator = None
 
-
 def segmentation_interface(image):
     if image is None: return None, "Upload an image"
     if not _mask_generator: return None, "SAM-2 unavailable"
@@ -157,7 +151,7 @@ def segmentation_interface(image):
     for ann in sorted(anns, key=lambda x: x['area'], reverse=True):
         m = ann['segmentation']; color=np.random.randint(0,255,3)
         overlay[m] = (overlay[m]*0.5 + color*0.5).astype(np.uint8)
-    return Image.fromarray(overlay), f"{len(anns)} masks"
+    return Image.fromarray(overlay), f"{len(anns)} masks found"
 
 # =============================================================================
 # Fallback segmentation
@@ -176,9 +170,7 @@ def fallback_segmentation(image):
 # =============================================================================
 try:
     chex_tok = AutoTokenizer.from_pretrained("StanfordAIMI/CheXagent-2-3b", trust_remote_code=True)
-    chex_model = AutoModelForCausalLM.from_pretrained(
-        "StanfordAIMI/CheXagent-2-3b", device_map='auto', trust_remote_code=True
-    )
+    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:
@@ -188,13 +180,11 @@ except Exception:
 def report_generation(im1, im2):
     if not CHEX_AVAILABLE: yield "CheXagent unavailable"; return
     streamer = TextIteratorStreamer(chex_tok, skip_prompt=True)
-    # ... similar to above, streaming report generation ...
-    yield "Report not implemented in snippet"
+    yield "Report streaming not fully implemented"
 
 @torch.no_grad()
 def phrase_grounding(image, prompt):
     if not CHEX_AVAILABLE: return "CheXagent unavailable", None
-    # simple box
     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
@@ -202,36 +192,27 @@ def phrase_grounding(image, prompt):
 # =============================================================================
 # Gradio UI
 # =============================================================================
+
 def create_ui():
-    # Load agents
     try:
-        q_model,q_proc,q_dev = load_qwen_model_and_processor()
-        med_agent = MedicalVLMAgent(q_model,q_proc,q_dev); QWEN_OK=True
+        m, p, d = load_qwen_model_and_processor()
+        med = MedicalVLMAgent(m,p,d); QW=True
     except:
-        QWEN_OK=False; med_agent=None
-
+        QW=False; med=None
     with gr.Blocks() as demo:
         gr.Markdown("# Medical AI Assistant")
-        gr.Markdown(f"- Qwen VLM: {'✅' if QWEN_OK else '❌'}  "
-                    f"- SAM-2: {'✅' if _mask_generator else '❌'}  "
-                    f"- CheXagent: {'✅' if CHEX_AVAILABLE else '❌'}")
+        gr.Markdown(f"- Qwen: {'✅' if QW else '❌'}  - SAM-2: {'✅' if _mask_generator else '❌'}  - CheX: {'✅' if CHEX_AVAILABLE else '❌'}")
         with gr.Tab("Medical Q&A"):
-            txt=gr.Textbox(); img=gr.Image(type='pil'); out=gr.Textbox(); btn=gr.Button("Ask")
-            btn.click(med_agent.run, [txt,img], out)
+            txt=gr.Textbox(); img=gr.Image(type='pil'); out=gr.Textbox(); gr.Button("Ask").click(med.run,[txt,img],out)
         with gr.Tab("Segmentation"):
-            segin=gr.Image(type='pil'); segout=gr.Image(); stat=gr.Textbox()
-            if _mask_generator: fn=segmentation_interface
-            else: fn=fallback_segmentation
-            gr.Button("Segment").click(fn, segin, [segout, stat])
+            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])
         with gr.Tab("CheXagent Report"):
-            c1=gr.Image(type='pil'); c2=gr.Image(type='pil'); rout=gr.Markdown()
-            gr.Interface(fn=report_generation, inputs=[c1,c2], outputs=rout, live=True).render()
+            c1=gr.Image(type='pil');c2=gr.Image(type='pil'); rout=gr.Markdown(); gr.Interface(report_generation,[c1,c2],rout,live=True).render()
         with gr.Tab("CheXagent Grounding"):
-            gi=gr.Image(type='pil'); gp=gr.Textbox(); gout=gr.Textbox(); goimg=gr.Image()
-            gr.Interface(fn=phrase_grounding, inputs=[gi,gp], outputs=[gout,goimg]).render()
+            gi=gr.Image(type='pil'); gp=gr.Textbox(); gout=gr.Textbox(); goimg=gr.Image(); gr.Interface(phrase_grounding,[gi,gp],[gout,goimg]).render()
     return demo
 
 if __name__ == "__main__":
-    ui = create_ui()
-    ui.launch(server_name='0.0.0.0', server_port=7860, share=True)
+    ui=create_ui(); ui.launch(server_name='0.0.0.0',server_port=7860,share=True)
+