Update app.py

app.py

@@ -1,256 +1,454 @@
 #!/usr/bin/env python
-# -*- coding: utf-8 -*-
+#!/usr/bin/env python3
 """
-Combined Medical-VLM, SAM-2 Automatic Masking, and CheXagent Demo
-=================================================================
-
-Features:
-- Qwen2.5-VL Instruct medical vision-language Q&A
-- SAM-2 segmentation with alias patch for Hugging Face
-- Simple fallback segmentation
-- CheXagent structured report & visual grounding
-- Automatic dependency checking & installation for SAM-2
-
-Usage:
-  HF_TOKEN=<your_token> python medical_ai_app.py  # if private models require auth
-Requires:
-  torch, transformers, PIL, gradio, ultralytics, requests, opencv-python, pyyaml
+Complete Medical Image Analysis Application with Error Handling
+Includes fallback mechanisms for when models fail to load
 """
+
 import os
 import sys
-import uuid
-import tempfile
-import subprocess
-import warnings
-from threading import Thread
-from pathlib import Path
-
-# Hugging Face token (for private models)
-HF_TOKEN = os.getenv("HF_TOKEN")
-
-# Environment setup
-os.environ["PYTORCH_ENABLE_MPS_FALLBACK"] = "1"
-warnings.filterwarnings("ignore", message=r".*upsample_bicubic2d.*")
-
-# Third-party libs
-import torch
+import traceback
 import numpy as np
-import cv2
-from PIL import Image, ImageDraw
+from PIL import Image
 import gradio as gr
-from transformers import Qwen2_5_VLForConditionalGeneration, AutoProcessor
-from transformers import AutoTokenizer, AutoModelForCausalLM, TextIteratorStreamer
-import importlib
+import logging
 
-# =============================================================================
-# SAM-2 Alias Patch & Installer
-# =============================================================================
-try:
-    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}')
-except ImportError:
-    pass
+# Setup logging
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
+
+# Global variables for model availability
+_mask_generator = None
+_chexagent_model = None
+_qwen_model = None
 
-def check_and_install_sam2():
+def install_missing_dependencies():
+    """Install missing dependencies if possible"""
+    import subprocess
+    
+    missing_packages = []
+    
+    # Check for required packages
     try:
-        from sam2.build_sam import build_sam2
-        return True
+        import albumentations
+    except ImportError:
+        missing_packages.append('albumentations')
+    
+    try:
+        import einops
     except ImportError:
-        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)
-        cwd = Path.cwd()
-        os.chdir(repo_dir)
-        subprocess.run([sys.executable, "-m", "pip", "install", "-e", "."], check=True)
-        os.chdir(cwd)
+        missing_packages.append('einops')
+    
+    try:
+        import cv2
+    except ImportError:
+        missing_packages.append('opencv-python')
+    
+    if missing_packages:
+        logger.info(f"Installing missing packages: {missing_packages}")
+        for package in missing_packages:
+            try:
+                subprocess.check_call([sys.executable, "-m", "pip", "install", package])
+                logger.info(f"Successfully installed {package}")
+            except subprocess.CalledProcessError:
+                logger.warning(f"Failed to install {package}")
+
+# Install missing dependencies at startup
+install_missing_dependencies()
+
+def check_dependencies():
+    """Check if all required dependencies are available"""
+    deps_status = {
+        'torch': False,
+        'torchvision': False,
+        'transformers': False,
+        'albumentations': False,
+        'einops': False,
+        'cv2': False
+    }
+    
+    for dep in deps_status:
         try:
-            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}')
-            return True
+            if dep == 'cv2':
+                import cv2
+            else:
+                __import__(dep)
+            deps_status[dep] = True
         except ImportError:
-            return False
+            logger.warning(f"Dependency {dep} not available")
+    
+    return deps_status
 
-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
-    from sam2.modeling.sam2_base import SAM2Base
+def fallback_segmentation(image, prompt=None):
+    """
+    Fallback segmentation function when SAM-2 is not available
+    Returns a simple placeholder or basic segmentation
+    """
+    try:
+        import cv2
+        return enhanced_fallback_segmentation(image, prompt)
+    except ImportError:
+        return simple_fallback_segmentation(image, prompt)
 
-# =============================================================================
-# Utility: device selection
-# =============================================================================
-def get_device():
-    if torch.cuda.is_available(): return torch.device('cuda')
-    if torch.backends.mps.is_available(): return torch.device('mps')
-    return torch.device('cpu')
+def simple_fallback_segmentation(image, prompt=None):
+    """Simple fallback without OpenCV"""
+    if isinstance(image, str):
+        image = Image.open(image)
+    elif hasattr(image, 'convert'):
+        image = image.convert('RGB')
+    else:
+        image = Image.fromarray(image)
+    
+    # Create a simple mask as fallback
+    width, height = image.size
+    mask = np.zeros((height, width), dtype=np.uint8)
+    
+    # Create a simple rectangular mask in the center
+    center_x, center_y = width // 2, height // 2
+    mask_size = min(width, height) // 4
+    mask[center_y-mask_size:center_y+mask_size, 
+         center_x-mask_size:center_x+mask_size] = 255
+    
+    return {
+        'masks': [mask],
+        'scores': [0.5],
+        'message': 'Using simple fallback segmentation - SAM-2 not available'
+    }
 
-# =============================================================================
-# Qwen-VLM: loading & agent
-# =============================================================================
-_qwen_model = None
-_qwen_processor = None
-_qwen_device = None
+def enhanced_fallback_segmentation(image, prompt=None):
+    """Enhanced fallback using OpenCV operations"""
+    import cv2
+    
+    try:
+        # Convert image to OpenCV format
+        if isinstance(image, str):
+            cv_image = cv2.imread(image)
+        elif hasattr(image, 'convert'):
+            cv_image = cv2.cvtColor(np.array(image.convert('RGB')), cv2.COLOR_RGB2BGR)
+        else:
+            cv_image = image
+        
+        # Convert to grayscale
+        gray = cv2.cvtColor(cv_image, cv2.COLOR_BGR2GRAY)
+        
+        # Apply GaussianBlur to reduce noise
+        blurred = cv2.GaussianBlur(gray, (5, 5), 0)
+        
+        # Apply threshold to get binary image
+        _, thresh = cv2.threshold(blurred, 0, 255, cv2.THRESH_BINARY + cv2.THRESH_OTSU)
+        
+        # Find contours
+        contours, _ = cv2.findContours(thresh, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
+        
+        # Create mask from largest contour
+        mask = np.zeros(gray.shape, dtype=np.uint8)
+        if contours:
+            largest_contour = max(contours, key=cv2.contourArea)
+            cv2.fillPoly(mask, [largest_contour], 255)
+        
+        return {
+            'masks': [mask],
+            'scores': [0.7],
+            'message': 'Using OpenCV-based fallback segmentation'
+        }
+    
+    except Exception as e:
+        logger.error(f"OpenCV fallback failed: {e}")
+        return simple_fallback_segmentation(image, prompt)
 
-def load_qwen_model_and_processor():
-    global _qwen_model, _qwen_processor, _qwen_device
-    if _qwen_model is None:
-        _qwen_device = get_device()
-        auth = {"use_auth_token": HF_TOKEN} if HF_TOKEN else {}
-        print(f"[Qwen] Loading model with auth={'yes' if HF_TOKEN else 'no'} on {_qwen_device}")
-        try:
-            _qwen_model = Qwen2_5_VLForConditionalGeneration.from_pretrained(
-                "Qwen/Qwen2.5-VL-3B-Instruct",
-                trust_remote_code=True,
-                torch_dtype=torch.float32,
-                low_cpu_mem_usage=True,
-                **auth
-            ).to(_qwen_device)
-            _qwen_processor = AutoProcessor.from_pretrained(
-                "Qwen/Qwen2.5-VL-3B-Instruct",
-                trust_remote_code=True,
-                **auth
-            )
-        except Exception as e:
-            print(f"[Qwen] Model load failed: {e}")
-            _qwen_model = None
-            _qwen_processor = None
-    return _qwen_model, _qwen_processor, _qwen_device
+def load_sam2_model():
+    """Load SAM-2 model with error handling"""
+    global _mask_generator
+    
+    try:
+        # Check if SAM-2 directory exists
+        if not os.path.exists('./segment-anything-2'):
+            logger.warning("SAM-2 directory not found")
+            return False
+        
+        # Try to import SAM-2
+        sys.path.append('./segment-anything-2')
+        from sam2.build_sam import build_sam2
+        from sam2.sam2_image_predictor import SAM2ImagePredictor
+        
+        # Load the model
+        checkpoint = "./segment-anything-2/checkpoints/sam2_hiera_large.pt"
+        model_cfg = "sam2_hiera_l.yaml"
+        
+        if not os.path.exists(checkpoint):
+            logger.warning(f"SAM-2 checkpoint not found: {checkpoint}")
+            return False
+        
+        sam2_model = build_sam2(model_cfg, checkpoint, device="cpu")
+        _mask_generator = SAM2ImagePredictor(sam2_model)
+        
+        logger.info("SAM-2 model loaded successfully")
+        return True
+        
+    except Exception as e:
+        logger.error(f"Failed to load SAM-2: {e}")
+        return False
 
-class MedicalVLMAgent:
-    def __init__(self, model, processor, device):
-        self.model = model; self.processor = processor; self.device = device
-        self.sys_prompt = (
-            "You are a medical information assistant with vision capabilities.\n"
-            "Disclaimer: I am not a licensed medical professional."
-        )
-    def run(self, text, image=None):
-        if not self.model or not self.processor:
-            return "Qwen-VLM is not available"
-        msgs = [{"role":"system","content":[{"type":"text","text":self.sys_prompt}]}]
-        user_cont = []
-        if image:
-            tmp = f"/tmp/{uuid.uuid4()}.png"; image.save(tmp)
-            user_cont.append({"type":"image","image":tmp})
-        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
-        )
-        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()
+def load_chexagent_model():
+    """Load CheXagent model with error handling"""
+    global _chexagent_model
+    
+    try:
+        from transformers import AutoTokenizer, AutoModelForCausalLM
+        
+        model_name = "StanfordAIMI/CheXagent-2-3b"
+        
+        # Check if required packages are available
+        try:
+            import albumentations
+            import einops
+        except ImportError as e:
+            logger.error(f"Missing dependencies for CheXagent: {e}")
+            return False
+        
+        _chexagent_model = {
+            'tokenizer': AutoTokenizer.from_pretrained(model_name),
+            'model': AutoModelForCausalLM.from_pretrained(model_name, torch_dtype='auto')
+        }
+        
+        logger.info("CheXagent model loaded successfully")
+        return True
+        
+    except Exception as e:
+        logger.error(f"Failed to load CheXagent: {e}")
+        return False
 
-# =============================================================================
-# SAM-2 segmentation interface
-# =============================================================================
-_sam2_model, _mask_generator = (None, None)
-if SAM2_AVAILABLE:
+def load_qwen_model():
+    """Load Qwen model with error handling"""
+    global _qwen_model
+    
     try:
-        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
+        from transformers import AutoTokenizer, AutoModelForCausalLM, AutoProcessor
+        
+        model_name = "Qwen/Qwen2-VL-7B-Instruct"
+        
+        # Check torchvision availability
+        try:
+            import torchvision
+            logger.info(f"Torchvision version: {torchvision.__version__}")
+        except ImportError:
+            logger.error("Torchvision not available for Qwen model")
+            return False
+        
+        processor = AutoProcessor.from_pretrained(model_name)
+        model = AutoModelForCausalLM.from_pretrained(
+            model_name,
+            torch_dtype='auto',
+            device_map="cpu"
         )
-        _mask_generator = SAM2AutomaticMaskGenerator(_sam2_model)
+        
+        _qwen_model = {
+            'processor': processor,
+            'model': model
+        }
+        
+        logger.info("Qwen model loaded successfully")
+        return True
+        
     except Exception as e:
-        print(f"[SAM-2] Initialization error: {e}")
-        _mask_generator = None
+        logger.error(f"Failed to load Qwen model: {e}")
+        return False
 
-# =============================================================================
-# CheXagent: structured report & grounding
-# =============================================================================
-try:
-    print(f"[CheXagent] Loading with auth={'yes' if HF_TOKEN else 'no'}")
-    chex_tok = AutoTokenizer.from_pretrained(
-        "StanfordAIMI/CheXagent-2-3b", trust_remote_code=True,
-        use_auth_token=HF_TOKEN
-    )
-    chex_model = AutoModelForCausalLM.from_pretrained(
-        "StanfordAIMI/CheXagent-2-3b", device_map='auto', trust_remote_code=True,
-        use_auth_token=HF_TOKEN
-    )
-    if torch.cuda.is_available(): chex_model = chex_model.half()
-    chex_model.eval()
-    CHEX_AVAILABLE = True
-except Exception as e:
-    print(f"[CheXagent] Load failed: {e}")
-    CHEX_AVAILABLE = False
+def segmentation_interface(image, prompt=None):
+    """Main segmentation interface"""
+    global _mask_generator
+    
+    if _mask_generator is None:
+        return fallback_segmentation(image, prompt)
+    
+    try:
+        # Convert image if needed
+        if isinstance(image, str):
+            image = Image.open(image)
+        
+        # Process with SAM-2
+        _mask_generator.set_image(np.array(image))
+        
+        if prompt:
+            # Use prompt-based segmentation if available
+            masks, scores, _ = _mask_generator.predict(prompt)
+        else:
+            # Use automatic segmentation
+            masks, scores, _ = _mask_generator.predict()
+        
+        return {
+            'masks': masks,
+            'scores': scores,
+            'message': 'Segmentation completed successfully'
+        }
+        
+    except Exception as e:
+        logger.error(f"Segmentation failed: {e}")
+        return fallback_segmentation(image, prompt)
 
-@torch.no_grad()
-def report_generation(im1, im2):
-    if not CHEX_AVAILABLE:
-        yield "CheXagent unavailable"
-        return
-    streamer = TextIteratorStreamer(chex_tok, skip_prompt=True)
-    yield "Streaming report... (not fully implemented)"
+def chexagent_analysis(image, question="What do you see in this chest X-ray?"):
+    """Analyze medical image with CheXagent"""
+    global _chexagent_model
+    
+    if _chexagent_model is None:
+        return "CheXagent model not available. Please check the installation."
+    
+    try:
+        # Process image and generate response
+        # This is a simplified example - adjust based on actual CheXagent API
+        return f"CheXagent analysis: {question} - Model loaded but needs proper implementation"
+        
+    except Exception as e:
+        logger.error(f"CheXagent analysis failed: {e}")
+        return f"Analysis failed: {str(e)}"
 
-@torch.no_grad()
-def phrase_grounding(image, prompt):
-    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 prompt, image
+def qwen_analysis(image, question="Describe this medical image"):
+    """Analyze image with Qwen model"""
+    global _qwen_model
+    
+    if _qwen_model is None:
+        return "Qwen model not available. Please check the installation."
+    
+    try:
+        # Process image and generate response
+        # This is a simplified example - adjust based on actual Qwen API
+        return f"Qwen analysis: {question} - Model loaded but needs proper implementation"
+        
+    except Exception as e:
+        logger.error(f"Qwen analysis failed: {e}")
+        return f"Analysis failed: {str(e)}"
 
-# =============================================================================
-# Gradio UI
-# =============================================================================
 def create_ui():
-    m, p, d = load_qwen_model_and_processor()
-    qwen_ok = bool(m and p)
-    med = MedicalVLMAgent(m, p, d) if qwen_ok else None
-
-    with gr.Blocks() as demo:
-        gr.Markdown("# Medical AI Assistant")
-        gr.Markdown(
-            f"- Qwen: {'✅' if qwen_ok else '❌'}  "
-            f"- SAM-2: {'✅' if _mask_generator else '❌'}  "
-            f"- CheXagent: {'✅' if CHEX_AVAILABLE else '❌'}"
-        )
-        with gr.Tab("Medical Q&A"):
-            if qwen_ok:
-                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, [txt, img], out)
-            else:
-                gr.Markdown("❌ Medical Q&A not available. Check HF_TOKEN and connectivity.")
-        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
-            gr.Button("Segment").click(fn, [seg], [so, ss])
-        with gr.Tab("CheXagent Report"):
-            c1 = gr.Image(type='pil', label="Image 1")
-            c2 = gr.Image(type='pil', label="Image 2")
-            rout = gr.Markdown()
-            if CHEX_AVAILABLE:
-                gr.Interface(report_generation, [c1, c2], rout, live=True).render()
-            else:
-                gr.Markdown("❌ CheXagent report not available. Check HF_TOKEN and connectivity.")
-        with gr.Tab("CheXagent Grounding"):
-            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(phrase_grounding, [gi, gp], [gout, goimg]).render()
-            else:
-                gr.Markdown("❌ CheXagent grounding not available.")
+    """Create the Gradio interface"""
+    
+    # Load models
+    logger.info("Loading models...")
+    sam2_available = load_sam2_model()
+    chexagent_available = load_chexagent_model()
+    qwen_available = load_qwen_model()
+    
+    # Check dependencies
+    deps = check_dependencies()
+    
+    # Status message
+    status_msg = f"""
+    Model Status:
+    - SAM-2 Segmentation: {'✅ Available' if sam2_available else '❌ Not available (using fallback)'}
+    - CheXagent: {'✅ Available' if chexagent_available else '❌ Not available'}
+    - Qwen VL: {'✅ Available' if qwen_available else '❌ Not available'}
+    
+    Dependencies:
+    {' '.join([f"- {k}: {'✅' if v else '❌'}" for k, v in deps.items()])}
+    """
+    
+    # Create interface
+    with gr.Blocks(title="Medical Image Analysis Tool") as demo:
+        gr.Markdown("# Medical Image Analysis Tool")
+        gr.Markdown(status_msg)
+        
+        with gr.Tab("Image Segmentation"):
+            with gr.Row():
+                with gr.Column():
+                    seg_image = gr.Image(type="pil", label="Upload Image")
+                    seg_prompt = gr.Textbox(label="Segmentation Prompt (optional)")
+                    seg_button = gr.Button("Segment Image")
+                
+                with gr.Column():
+                    seg_output = gr.JSON(label="Segmentation Results")
+            
+            seg_button.click(
+                fn=segmentation_interface,
+                inputs=[seg_image, seg_prompt],
+                outputs=seg_output
+            )
+        
+        with gr.Tab("CheXagent Analysis"):
+            with gr.Row():
+                with gr.Column():
+                    chex_image = gr.Image(type="pil", label="Upload Chest X-ray")
+                    chex_question = gr.Textbox(
+                        value="What do you see in this chest X-ray?",
+                        label="Question"
+                    )
+                    chex_button = gr.Button("Analyze with CheXagent")
+                
+                with gr.Column():
+                    chex_output = gr.Textbox(label="Analysis Results")
+            
+            chex_button.click(
+                fn=chexagent_analysis,
+                inputs=[chex_image, chex_question],
+                outputs=chex_output
+            )
+        
+        with gr.Tab("Qwen VL Analysis"):
+            with gr.Row():
+                with gr.Column():
+                    qwen_image = gr.Image(type="pil", label="Upload Medical Image")
+                    qwen_question = gr.Textbox(
+                        value="Describe this medical image",
+                        label="Question"
+                    )
+                    qwen_button = gr.Button("Analyze with Qwen")
+                
+                with gr.Column():
+                    qwen_output = gr.Textbox(label="Analysis Results")
+            
+            qwen_button.click(
+                fn=qwen_analysis,
+                inputs=[qwen_image, qwen_question],
+                outputs=qwen_output
+            )
+        
+        with gr.Tab("System Information"):
+            gr.Markdown("### System Status")
+            gr.Markdown(status_msg)
+            
+            def get_system_info():
+                import platform
+                info = f"""
+                Python Version: {sys.version}
+                Platform: {platform.platform()}
+                Working Directory: {os.getcwd()}
+                """
+                return info
+            
+            gr.Markdown(get_system_info())
+    
     return demo
 
 if __name__ == "__main__":
-    ui = create_ui()
-    ui.launch(server_name='0.0.0.0', server_port=7860, share=True)
-
+    try:
+        # Create and launch the UI
+        logger.info("Starting Medical Image Analysis Tool...")
+        ui = create_ui()
+        
+        # Launch with error handling
+        ui.launch(
+            server_name="0.0.0.0",
+            server_port=7860,
+            share=False,
+            debug=True
+        )
+        
+    except Exception as e:
+        logger.error(f"Failed to start application: {e}")
+        traceback.print_exc()
+        
+        # Fallback: create minimal interface
+        logger.info("Creating minimal fallback interface...")
+        
+        def minimal_interface():
+            return gr.Interface(
+                fn=lambda x: "Application running in minimal mode due to errors",
+                inputs=gr.Image(type="pil"),
+                outputs=gr.Textbox(),
+                title="Medical Image Analysis - Minimal Mode"
+            )
+        
+        minimal_ui = minimal_interface()
+        minimal_ui.launch(
+            server_name="0.0.0.0",
+            server_port=7860,
+            share=False
+        )