app.py

# app.py - Medical AI using LLaVA (Large Language and Vision Assistant)
import gradio as gr
import torch
from transformers import LlavaNextProcessor, LlavaNextForConditionalGeneration
from PIL import Image
import logging
from collections import defaultdict, Counter
import time

# Configure logging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

# Usage tracking
class UsageTracker:
    def __init__(self):
        self.stats = {
            'total_analyses': 0,
            'successful_analyses': 0,
            'failed_analyses': 0,
            'average_processing_time': 0.0,
            'question_types': Counter()
        }
    
    def log_analysis(self, success, duration, question_type=None):
        self.stats['total_analyses'] += 1
        if success:
            self.stats['successful_analyses'] += 1
        else:
            self.stats['failed_analyses'] += 1
        
        total_time = self.stats['average_processing_time'] * (self.stats['total_analyses'] - 1)
        self.stats['average_processing_time'] = (total_time + duration) / self.stats['total_analyses']
        
        if question_type:
            self.stats['question_types'][question_type] += 1

# Rate limiting
class RateLimiter:
    def __init__(self, max_requests_per_hour=30):
        self.max_requests_per_hour = max_requests_per_hour
        self.requests = defaultdict(list)
    
    def is_allowed(self, user_id="default"):
        current_time = time.time()
        hour_ago = current_time - 3600
        self.requests[user_id] = [req_time for req_time in self.requests[user_id] if req_time > hour_ago]
        if len(self.requests[user_id]) < self.max_requests_per_hour:
            self.requests[user_id].append(current_time)
            return True
        return False

# Initialize components
usage_tracker = UsageTracker()
rate_limiter = RateLimiter()

# Model configuration - Using LLaVA-Next (latest version)
MODEL_ID = "llava-hf/llava-v1.6-mistral-7b-hf"

# Global variables
model = None
processor = None

def load_llava():
    """Load LLaVA model for medical analysis"""
    global model, processor
    
    try:
        logger.info(f"Loading LLaVA model: {MODEL_ID}")
        
        # Load processor
        processor = LlavaNextProcessor.from_pretrained(MODEL_ID)
        
        # Load model with appropriate settings
        model = LlavaNextForConditionalGeneration.from_pretrained(
            MODEL_ID, 
            torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
            device_map="auto" if torch.cuda.is_available() else None,
            low_cpu_mem_usage=True
        )
        
        logger.info("✅ LLaVA model loaded successfully!")
        return True
        
    except Exception as e:
        logger.error(f"❌ Error loading LLaVA: {str(e)}")
        return False

# Load model at startup
llava_ready = load_llava()

def analyze_medical_image_llava(image, clinical_question, patient_history=""):
    """Analyze medical image using LLaVA"""
    start_time = time.time()
    
    # Rate limiting
    if not rate_limiter.is_allowed():
        usage_tracker.log_analysis(False, time.time() - start_time)
        return "⚠️ Rate limit exceeded. Please wait before trying again."
    
    if not llava_ready or model is None:
        usage_tracker.log_analysis(False, time.time() - start_time)
        return "❌ LLaVA model not loaded. Please refresh the page and wait for model loading."
    
    if image is None:
        return "⚠️ Please upload a medical image first."
    
    if not clinical_question.strip():
        return "⚠️ Please provide a clinical question."
    
    try:
        logger.info("Starting LLaVA medical analysis...")
        
        # Prepare comprehensive medical prompt
        medical_prompt = f"""You are a highly skilled medical AI assistant with expertise in medical image analysis. You have extensive knowledge in:

- **Radiology**: X-rays, CT scans, MRI, ultrasound interpretation
- **Pathology**: Histological analysis, tissue examination, cellular patterns
- **Dermatology**: Skin lesions, rashes, dermatological conditions
- **Ophthalmology**: Retinal imaging, eye examinations, ocular pathology
- **General Medical Imaging**: Cross-sectional anatomy, normal variants, pathological findings

**Patient Information:**
{f"Patient History: {patient_history}" if patient_history.strip() else "No specific patient history provided"}

**Clinical Question:** {clinical_question}

**Instructions:**
Please provide a comprehensive medical analysis of this image following this structure:

1. **IMAGE QUALITY ASSESSMENT**
   - Technical adequacy of the image
   - Any artifacts or limitations
   - Overall diagnostic quality

2. **SYSTEMATIC OBSERVATION**
   - Describe what you see in detail
   - Identify anatomical structures visible
   - Note any normal findings

3. **ABNORMAL FINDINGS**
   - Identify any pathological changes
   - Describe abnormalities in detail
   - Note their location and characteristics

4. **CLINICAL SIGNIFICANCE**
   - Explain the importance of findings
   - Relate to potential diagnoses
   - Discuss clinical implications

5. **DIFFERENTIAL DIAGNOSIS**
   - List possible conditions
   - Explain reasoning for each
   - Prioritize based on imaging findings

6. **RECOMMENDATIONS**
   - Suggest additional imaging if needed
   - Recommend clinical correlation
   - Advise on follow-up or further evaluation

Please be thorough, educational, and professional in your analysis. Always emphasize that this is for educational purposes and requires professional medical validation."""
        
        # Prepare conversation for LLaVA
        conversation = [
            {
                "role": "user",
                "content": [
                    {"type": "text", "text": medical_prompt},
                    {"type": "image", "image": image}
                ]
            }
        ]
        
        # Apply chat template
        prompt = processor.apply_chat_template(conversation, add_generation_prompt=True)
        
        # Process inputs
        inputs = processor(prompt, image, return_tensors='pt')
        
        # Move to appropriate device
        if torch.cuda.is_available() and hasattr(model, 'device'):
            inputs = {k: v.to(model.device) for k, v in inputs.items()}
        
        # Generate response
        logger.info("Generating comprehensive medical analysis...")
        with torch.inference_mode():
            output = model.generate(
                **inputs,
                max_new_tokens=2000,
                do_sample=True,
                temperature=0.2,  # Lower temperature for more focused medical analysis
                top_p=0.9,
                repetition_penalty=1.1,
                pad_token_id=processor.tokenizer.eos_token_id
            )
        
        # Decode response
        generated_text = processor.decode(output[0][inputs["input_ids"].shape[-1]:], skip_special_tokens=True)
        
        # Clean up response
        response = generated_text.strip()
        
        # Format the response with medical structure
        formatted_response = f"""# 🏥 **LLaVA Medical Image Analysis**

## **Clinical Question:** {clinical_question}
{f"## **Patient History:** {patient_history}" if patient_history.strip() else ""}

---

## 🔍 **Comprehensive Medical Analysis**

{response}

---

## 📋 **Summary and Clinical Correlation**

**Key Points:**
- This analysis provides a systematic approach to medical image interpretation
- All findings should be correlated with clinical presentation and patient history
- The AI assessment serves as an educational tool and decision support aid

**Clinical Workflow:**
1. **Review** the systematic analysis above
2. **Correlate** findings with patient symptoms and history  
3. **Consult** with appropriate medical specialists as needed
4. **Document** findings in the patient's medical record
5. **Follow up** with recommended additional studies if indicated

**Educational Value:**
This analysis demonstrates structured medical image interpretation methodology and clinical reasoning processes used in healthcare settings.
"""
        
        # Add comprehensive medical disclaimer
        disclaimer = """
---
## ⚠️ **IMPORTANT MEDICAL DISCLAIMER**

**FOR EDUCATIONAL AND RESEARCH PURPOSES ONLY**

- **Not a Medical Diagnosis**: This AI analysis does not constitute a medical diagnosis, treatment recommendation, or professional medical advice
- **Professional Review Required**: All findings must be validated by qualified healthcare professionals
- **Emergency Situations**: For urgent medical concerns, contact emergency services immediately (911 in US)
- **Clinical Correlation**: AI findings must be correlated with clinical examination and patient history
- **Liability**: This system is not intended for clinical decision-making and users assume all responsibility
- **Educational Tool**: Designed for medical education, training, and research applications only
- **Data Privacy**: Do not upload images containing patient identifiable information

**Always consult qualified healthcare professionals for medical diagnosis and treatment decisions.**

---
**Powered by**: LLaVA (Large Language and Vision Assistant) | **Model**: {MODEL_ID}
        """
        
        # Log successful analysis
        duration = time.time() - start_time
        question_type = classify_question(clinical_question)
        usage_tracker.log_analysis(True, duration, question_type)
        
        logger.info("✅ LLaVA medical analysis completed successfully")
        return formatted_response + disclaimer
        
    except Exception as e:
        duration = time.time() - start_time
        usage_tracker.log_analysis(False, duration)
        logger.error(f"❌ LLaVA analysis error: {str(e)}")
        
        if "memory" in str(e).lower() or "cuda" in str(e).lower():
            return "❌ **Memory Error**: The model requires more memory. Try using a smaller image or upgrading to GPU hardware."
        else:
            return f"❌ **Analysis Failed**: {str(e)}\n\nPlease try again with a different image or contact support if the issue persists."

def classify_question(question):
    """Classify clinical question type"""
    question_lower = question.lower()
    if any(word in question_lower for word in ['describe', 'findings', 'observe', 'see']):
        return 'descriptive'
    elif any(word in question_lower for word in ['diagnosis', 'differential', 'condition', 'disease']):
        return 'diagnostic'
    elif any(word in question_lower for word in ['abnormal', 'pathology', 'lesion', 'mass']):
        return 'pathological'
    elif any(word in question_lower for word in ['analyze', 'assess', 'evaluate', 'review']):
        return 'analytical'
    else:
        return 'general'

def get_usage_stats():
    """Get comprehensive usage statistics"""
    stats = usage_tracker.stats
    if stats['total_analyses'] == 0:
        return "📊 **Usage Statistics**\n\nNo analyses performed yet."
    
    success_rate = (stats['successful_analyses'] / stats['total_analyses']) * 100
    
    return f"""📊 **LLaVA Medical AI Usage Statistics**

**Performance Metrics:**
- **Total Analyses**: {stats['total_analyses']}
- **Success Rate**: {success_rate:.1f}%
- **Average Processing Time**: {stats['average_processing_time']:.2f} seconds
- **Failed Analyses**: {stats['failed_analyses']}

**Question Type Distribution:**
{chr(10).join([f"- **{qtype.title()}**: {count} ({count/stats['total_analyses']*100:.1f}%)" for qtype, count in stats['question_types'].most_common()])}

**System Information:**
- **Model**: LLaVA-v1.6-Mistral-7B
- **Capabilities**: Medical image analysis and clinical reasoning
- **Device**: {'GPU' if torch.cuda.is_available() else 'CPU'}
- **Status**: {'🟢 Operational' if llava_ready else '🔴 Offline'}
"""

# Create comprehensive Gradio interface
def create_interface():
    with gr.Blocks(
        title="LLaVA Medical Image Analysis", 
        theme=gr.themes.Soft(),
        css="""
        .gradio-container { max-width: 1400px !important; }
        .disclaimer { background-color: #fef2f2; border: 1px solid #fecaca; border-radius: 8px; padding: 16px; margin: 16px 0; }
        .success { background-color: #f0f9ff; border: 1px solid #bae6fd; border-radius: 8px; padding: 16px; margin: 16px 0; }
        .warning { background-color: #fffbeb; border: 1px solid #fed7aa; border-radius: 8px; padding: 16px; margin: 16px 0; }
        """
    ) as demo:
        
        # Header
        gr.Markdown("""
        # 🏥 LLaVA Medical Image Analysis
        
        **Advanced Medical AI powered by LLaVA (Large Language and Vision Assistant)**
        
        **Specialized Medical Capabilities:**
        🫁 **Radiology** • 🔬 **Pathology** • 🩺 **Dermatology** • 👁️ **Ophthalmology** • 🧠 **Clinical Reasoning**
        """)
        
        # Status display
        if llava_ready:
            gr.Markdown("""
            <div class="success">
            ✅ <strong>LLAVA MEDICAL AI READY</strong><br>
            LLaVA vision-language model loaded successfully. Ready for comprehensive medical image analysis with clinical reasoning.
            </div>
            """)
        else:
            gr.Markdown("""
            <div class="warning">
            ⚠️ <strong>MODEL LOADING IN PROGRESS</strong><br>
            LLaVA model is loading. This may take a few minutes. Please wait and refresh the page.
            </div>
            """)
        
        # Medical disclaimer
        gr.Markdown("""
        <div class="disclaimer">
        ⚠️ <strong>CRITICAL MEDICAL DISCLAIMER</strong><br>
        This AI tool provides <strong>educational medical analysis only</strong>. It is NOT a substitute for professional medical diagnosis.
        <br><br>
        <strong>Do NOT upload real patient data or PHI.</strong> Always consult qualified healthcare professionals for medical decisions.
        </div>
        """)
        
        with gr.Row():
            # Left column - Main interface
            with gr.Column(scale=2):
                with gr.Row():
                    with gr.Column():
                        gr.Markdown("## 📤 Medical Image Upload")
                        image_input = gr.Image(
                            label="Upload Medical Image", 
                            type="pil",
                            height=350,
                            sources=["upload", "clipboard"]
                        )
                        
                    with gr.Column():
                        gr.Markdown("## 💬 Clinical Information")
                        clinical_question = gr.Textbox(
                            label="Clinical Question *",
                            placeholder="Examples:\n• Analyze this chest X-ray comprehensively\n• What pathological findings are visible in this image?\n• Provide differential diagnosis based on these imaging findings\n• Describe abnormalities and their clinical significance\n• Evaluate this medical image systematically",
                            lines=5,
                            max_lines=8
                        )
                        
                        patient_history = gr.Textbox(
                            label="Patient History & Clinical Context (Optional)",
                            placeholder="e.g., 58-year-old female presenting with chest pain and shortness of breath. History of hypertension and smoking. Recent onset of symptoms.",
                            lines=3,
                            max_lines=5
                        )
                
                with gr.Row():
                    clear_btn = gr.Button("🗑️ Clear All", variant="secondary")
                    analyze_btn = gr.Button("🔍 Analyze with LLaVA", variant="primary", size="lg")
                
                gr.Markdown("## 📋 LLaVA Medical Analysis Results")
                output = gr.Textbox(
                    label="Comprehensive Medical Analysis",
                    lines=30,
                    max_lines=50,
                    show_copy_button=True,
                    placeholder="Upload a medical image and provide a clinical question to receive comprehensive AI-powered medical analysis..." if llava_ready else "LLaVA model is loading. Please wait and refresh the page."
                )
            
            # Right column - System info and controls
            with gr.Column(scale=1):
                gr.Markdown("## ℹ️ System Status")
                
                model_status = "✅ Ready" if llava_ready else "🔄 Loading"
                device_info = "GPU" if torch.cuda.is_available() else "CPU"
                
                gr.Markdown(f"""
                **Model Status:** {model_status}  
                **AI Model:** LLaVA-v1.6-Mistral-7B  
                **Device:** {device_info}  
                **Capabilities:** Medical image analysis + clinical reasoning  
                **Context Length:** 32K tokens  
                **Rate Limit:** 30 requests/hour
                """)
                
                gr.Markdown("## 📊 Usage Analytics")
                stats_display = gr.Markdown("")
                refresh_stats_btn = gr.Button("🔄 Refresh Statistics", size="sm")
                
                if llava_ready:
                    gr.Markdown("## 🎯 Quick Clinical Examples")
                    
                    radiology_btn = gr.Button("🫁 Chest X-ray Analysis", size="sm")
                    pathology_btn = gr.Button("🔬 Pathology Review", size="sm")
                    dermatology_btn = gr.Button("🩺 Skin Lesion Analysis", size="sm")
                    differential_btn = gr.Button("🧠 Differential Diagnosis", size="sm")
                    
                    gr.Markdown("## 🏥 Medical Specialties")
                    gr.Markdown("""
                    **LLaVA excels in:**
                    - Radiology interpretation
                    - Pathological analysis  
                    - Dermatological assessment
                    - Ophthalmological evaluation
                    - Clinical reasoning & education
                    """)
        
        # Comprehensive example cases
        if llava_ready:
            with gr.Accordion("📚 Sample Medical Cases & Examples", open=False):
                examples = gr.Examples(
                    examples=[
                        [
                            "https://upload.wikimedia.org/wikipedia/commons/c/c8/Chest_Xray_PA_3-8-2010.png",
                            "Please perform a comprehensive systematic analysis of this chest X-ray. Evaluate image quality, assess cardiac silhouette, examine lung fields bilaterally, review mediastinal structures, and identify any pathological findings. Provide differential diagnosis if abnormalities are present.",
                            "Adult patient presenting with acute onset chest pain and shortness of breath. No significant past medical history."
                        ],
                        [
                            None,
                            "Analyze this medical image systematically. Describe normal anatomical structures, identify any abnormal findings, assess clinical significance, and provide appropriate differential diagnoses based on imaging characteristics.",
                            "Patient with acute presentation requiring medical imaging evaluation"
                        ],
                        [
                            None,
                            "What pathological changes are visible in this medical image? Please provide detailed morphological analysis, clinical correlation, and discuss potential diagnoses with supporting evidence from the imaging findings.",
                            ""
                        ]
                    ],
                    inputs=[image_input, clinical_question, patient_history],
                    label="Click any example to load it into the interface"
                )
        
        # Event handlers
        analyze_btn.click(
            fn=analyze_medical_image_llava,
            inputs=[image_input, clinical_question, patient_history],
            outputs=output,
            show_progress=True
        )
        
        def clear_all_fields():
            return None, "", "", ""
        
        clear_btn.click(
            fn=clear_all_fields,
            outputs=[image_input, clinical_question, patient_history, output]
        )
        
        refresh_stats_btn.click(
            fn=get_usage_stats,
            outputs=stats_display
        )
        
        # Quick example button handlers
        if llava_ready:
            radiology_btn.click(
                fn=lambda: ("Perform systematic radiological analysis of this medical image. Assess technical quality, identify normal anatomical structures, detect any pathological findings, and provide clinical interpretation with differential diagnosis.", "Adult patient with respiratory symptoms"),
                outputs=[clinical_question, patient_history]
            )
            
            pathology_btn.click(
                fn=lambda: ("Analyze this pathological specimen or medical image. Describe morphological features, identify cellular patterns, assess for pathological changes, and provide histopathological interpretation with clinical significance.", "Tissue sample for pathological evaluation"),
                outputs=[clinical_question, patient_history]
            )
            
            dermatology_btn.click(
                fn=lambda: ("Examine this dermatological image systematically. Describe the lesion characteristics, assess morphological features, evaluate for concerning signs, and provide differential diagnosis with clinical recommendations.", "Patient presenting with skin lesion requiring evaluation"),
                outputs=[clinical_question, patient_history]
            )
            
            differential_btn.click(
                fn=lambda: ("Based on the imaging findings in this medical image, provide a comprehensive differential diagnosis. List possible conditions in order of likelihood, explain supporting evidence for each diagnosis, and recommend additional studies if needed.", "Patient requiring diagnostic workup based on imaging findings"),
                outputs=[clinical_question, patient_history]
            )
        
        # Comprehensive footer with detailed information
        gr.Markdown("""
        ---
        ## 🤖 About LLaVA Medical AI
        
        **LLaVA (Large Language and Vision Assistant)** is a state-of-the-art multimodal AI model that combines advanced computer vision with natural language processing for comprehensive medical image analysis.
        
        ### 🔬 Key Capabilities
        
        **Medical Image Analysis:**
        - **Radiology**: X-rays, CT scans, MRI, ultrasound interpretation
        - **Pathology**: Histological analysis, tissue examination, cellular morphology
        - **Dermatology**: Skin lesion analysis, dermatological condition assessment
        - **Ophthalmology**: Retinal imaging, ocular pathology evaluation
        
        **Clinical Reasoning:**
        - Systematic medical image interpretation
        - Differential diagnosis generation
        - Clinical correlation and significance assessment
        - Educational medical content and explanations
        
        ### 🏥 Medical Education Applications
        
        - **Medical Student Training**: Interactive case-based learning
        - **Resident Education**: Systematic approach to image interpretation  
        - **Continuing Medical Education**: Advanced diagnostic reasoning
        - **Research Applications**: Medical imaging analysis and documentation
        
        ### 🔒 Privacy & Compliance
        
        - **No Data Storage**: All images processed in real-time, not stored
        - **Educational Purpose**: Designed specifically for medical education and training
        - **Privacy Protection**: No patient identifiable information should be uploaded
        - **Professional Standards**: Adheres to medical AI ethics and best practices
        
        ### ⚡ Technical Specifications
        
        - **Model**: LLaVA-v1.6-Mistral-7B (Latest version)
        - **Context Window**: 32,000 tokens for comprehensive analysis
        - **Processing**: Real-time inference with detailed medical reasoning
        - **Accuracy**: Research-grade performance on medical imaging tasks
        
        ### 📞 Support & Resources
        
        For technical support, feature requests, or educational partnerships, please contact our support team.
        
        ---
        **Powered by**: LLaVA (Large Language and Vision Assistant) | **License**: Apache 2.0 | **Purpose**: Medical Education & Research
        """)
    
    return demo

# Launch the application
if __name__ == "__main__":
    demo = create_interface()
    demo.launch(
        server_name="0.0.0.0",
        server_port=7860,
        show_error=True
    )