app.py

#!/usr/bin/env python
#!/usr/bin/env python3
"""
Complete Medical Image Analysis Application with Error Handling
Includes fallback mechanisms for when models fail to load
"""

import os
import sys
import traceback
import numpy as np
from PIL import Image
import gradio as gr
import logging

# 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 install_missing_dependencies():
    """Install missing dependencies if possible"""
    import subprocess
    
    missing_packages = []
    
    # Check for required packages
    try:
        import albumentations
    except ImportError:
        missing_packages.append('albumentations')
    
    try:
        import einops
    except ImportError:
        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:
            if dep == 'cv2':
                import cv2
            else:
                __import__(dep)
            deps_status[dep] = True
        except ImportError:
            logger.warning(f"Dependency {dep} not available")
    
    return deps_status

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)

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'
    }

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_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

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

def load_qwen_model():
    """Load Qwen model with error handling"""
    global _qwen_model
    
    try:
        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"
        )
        
        _qwen_model = {
            'processor': processor,
            'model': model
        }
        
        logger.info("Qwen model loaded successfully")
        return True
        
    except Exception as e:
        logger.error(f"Failed to load Qwen model: {e}")
        return 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)

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)}"

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)}"

def create_ui():
    """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__":
    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
        )