Upload app.py

app.py

@@ -1,571 +1,309 @@
-# -*- coding: utf-8 -*-
-# 🏥 Gemma 3N SOAP Note Generator with Unsloth
-# Optimized for offline medical documentation
-
-import torch
-import gradio as gr
-import io
-import base64
-from datetime import datetime
-import os
-import easyocr
-from PIL import Image, ImageDraw, ImageFont
-import cv2
-import numpy as np
-import psutil
-
-# Import Unsloth for optimized Gemma 3n
-try:
-    from unsloth import FastModel
-    print("✅ Unsloth imported successfully")
-    UNSLOTH_AVAILABLE = True
-except ImportError:
-    print("❌ Unsloth not available. Install with: pip install unsloth")
-    UNSLOTH_AVAILABLE = False
-
-# Device setup
-def setup_device():
-    device = "cuda" if torch.cuda.is_available() else "cpu"
-    print(f"🖥️  Using device: {device}")
-    
-    if torch.cuda.is_available():
-        print(f"🚀 GPU: {torch.cuda.get_device_name(0)}")
-        print(f"💾 GPU Memory: {torch.cuda.get_device_properties(0).total_memory / 1e9:.1f} GB")
-    else:
-        print("⚠️  Running on CPU - will be slower but works offline")
-    
-    return device
-
-# Load Unsloth Gemma 3n model
-def load_unsloth_gemma_model(device):
-    """Load optimized Gemma 3n model using Unsloth"""
-    
-    if not UNSLOTH_AVAILABLE:
-        print("❌ Unsloth not available. Using fallback method.")
-        return load_fallback_model()
-    
-    try:
-        print("📡 Loading Unsloth-optimized Gemma 3n model...")
-        
-        # Use the 4-bit quantized model for efficiency
-        model_name = "unsloth/gemma-3n-E4B-it-unsloth-bnb-4bit"
-        
-        print(f"🔧 Loading model: {model_name}")
-        
-        # Load with Unsloth optimizations
-        model, tokenizer = FastModel.from_pretrained(
-            model_name=model_name,
-            dtype=None,  # Auto-detect
-            max_seq_length=1024,  # Good for medical notes
-            load_in_4bit=True,   # 4-bit quantization for efficiency
-            full_finetuning=False,
-        )
-        
-        print("✅ Unsloth Gemma 3n model loaded successfully!")
-        print(f"📊 Model: {model_name}")
-        print(f"💾 Memory optimized with 4-bit quantization")
-        print(f"🎯 Ready for medical SOAP note generation!")
-        
-        return model, tokenizer
-        
-    except Exception as e:
-        print(f"❌ Error loading Unsloth model: {e}")
-        print("💡 Trying fallback model...")
-        return load_fallback_model()
-
-def load_fallback_model():
-    """Fallback model if Unsloth fails"""
-    try:
-        from transformers import AutoTokenizer, AutoModelForCausalLM
-        
-        print("🔄 Loading fallback model...")
-        model_name = "microsoft/DialoGPT-medium"
-        
-        tokenizer = AutoTokenizer.from_pretrained(model_name)
-        if tokenizer.pad_token is None:
-            tokenizer.pad_token = tokenizer.eos_token
-            
-        model = AutoModelForCausalLM.from_pretrained(
-            model_name,
-            torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
-            low_cpu_mem_usage=True
-        )
-        
-        print("✅ Fallback model loaded!")
-        return model, tokenizer
-        
-    except Exception as e:
-        print(f"❌ Fallback model also failed: {e}")
-        return None, None
-
-# Enhanced SOAP Note Generation with Gemma 3n
-def generate_soap_note_gemma(doctor_notes, model=None, tokenizer=None, include_timestamp=True):
-    """Generate SOAP note using Gemma 3n model"""
-    
-    if not doctor_notes.strip():
-        return "❌ Please enter some medical notes to process."
-    
-    if model is None or tokenizer is None:
-        return generate_template_soap(doctor_notes, include_timestamp)
-    
-    # Medical-specific prompt for Gemma 3n
-    prompt = f"""<bos><start_of_turn>user
-You are a medical AI assistant specialized in creating SOAP notes. Convert the following unstructured medical notes into a professional SOAP note format.
-
-Medical Notes:
-{doctor_notes}
-
-Please create a structured SOAP note with these sections:
-- SUBJECTIVE: Patient's reported symptoms, complaints, and relevant history
-- OBJECTIVE: Physical examination findings, vital signs, and observable data  
-- ASSESSMENT: Clinical diagnosis, differential diagnosis, and medical reasoning
-- PLAN: Treatment recommendations, medications, tests, and follow-up care
-
-<end_of_turn>
-<start_of_turn>model
-SOAP NOTE:
-
-SUBJECTIVE:"""
-
-    try:
-        # Tokenize input
-        inputs = tokenizer(
-            prompt, 
-            return_tensors="pt", 
-            truncation=True, 
-            max_length=512,
-            padding=True
-        )
-        
-        # Generate with optimized settings for medical text
-        with torch.no_grad():
-            outputs = model.generate(
-                **inputs,
-                max_new_tokens=400,
-                temperature=0.2,  # Lower temperature for medical precision
-                top_p=0.9,
-                do_sample=True,
-                repetition_penalty=1.1,
-                pad_token_id=tokenizer.eos_token_id,
-                eos_token_id=tokenizer.eos_token_id
-            )
-        
-        # Decode response
-        generated_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
-        
-        # Extract only the SOAP note part
-        if "SOAP NOTE:" in generated_text:
-            soap_response = generated_text.split("SOAP NOTE:")[1].strip()
-        else:
-            soap_response = generated_text[len(prompt):].strip()
-        
-        # Clean up response
-        soap_response = clean_soap_response(soap_response)
-        
-        # Add professional header
-        if include_timestamp:
-            timestamp = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
-            header = f"""📋 SOAP NOTE - Generated by Gemma 3n
-🕐 Timestamp: {timestamp}
-🤖 Model: Unsloth-optimized Gemma 3n (4-bit quantized)
-🔒 Processed locally on device
-🏥 Medical Documentation Assistant
-
-{'='*60}
-"""
-            return header + soap_response
-        
-        return soap_response
-        
-    except Exception as e:
-        print(f"❌ Generation error: {e}")
-        return generate_template_soap(doctor_notes, include_timestamp)
-
-def clean_soap_response(response):
-    """Clean and format SOAP note response"""
-    
-    # Remove any incomplete sentences at the end
-    lines = response.split('\n')
-    cleaned_lines = []
-    
-    for line in lines:
-        line = line.strip()
-        if line:
-            # Ensure proper SOAP section headers
-            if line.upper().startswith(('SUBJECTIVE', 'OBJECTIVE', 'ASSESSMENT', 'PLAN')):
-                if not line.endswith(':'):
-                    line += ':'
-                cleaned_lines.append(f"\n{line}")
-            else:
-                cleaned_lines.append(line)
-    
-    return '\n'.join(cleaned_lines).strip()
-
-# Template-based SOAP generation (enhanced fallback)
-def generate_template_soap(doctor_notes, include_timestamp=True):
-    """Enhanced template-based SOAP note generation"""
-    
-    notes_lower = doctor_notes.lower()
-    lines = doctor_notes.split('\n')
-    
-    # Enhanced keyword extraction
-    subjective_info = extract_section_info(lines, [
-        'complains', 'reports', 'states', 'denies', 'pain', 'symptoms', 
-        'history', 'onset', 'duration', 'patient says', 'chief complaint'
-    ])
-    
-    objective_info = extract_section_info(lines, [
-        'vital signs', 'vs:', 'bp', 'hr', 'temp', 'examination', 'exam',
-        'physical', 'inspection', 'palpation', 'auscultation', 'laboratory'
-    ])
-    
-    assessment_info = extract_section_info(lines, [
-        'diagnosis', 'impression', 'assessment', 'likely', 'possible', 
-        'rule out', 'differential', 'icd', 'condition'
-    ])
-    
-    plan_info = extract_section_info(lines, [
-        'plan', 'treatment', 'medication', 'prescribe', 'follow', 'return',
-        'therapy', 'intervention', 'monitoring', 'referral'
-    ])
-    
-    # Build comprehensive SOAP note
-    soap_note = build_soap_sections(subjective_info, objective_info, assessment_info, plan_info)
-    
-    if include_timestamp:
-        timestamp = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
-        header = f"""📋 SOAP NOTE (Template-Enhanced)
-🕐 Timestamp: {timestamp}
-🔒 Processed locally - HIPAA compliant
-🏥 Scribbled Docs Medical Assistant
-
-{'='*60}
-"""
-        return header + soap_note
-    
-    return soap_note
-
-def extract_section_info(lines, keywords):
-    """Extract relevant lines for each SOAP section"""
-    relevant_lines = []
-    for line in lines:
-        if any(keyword in line.lower() for keyword in keywords):
-            relevant_lines.append(line.strip())
-    return relevant_lines
-
-def build_soap_sections(subjective, objective, assessment, plan):
-    """Build formatted SOAP sections"""
-    
-    soap = "SUBJECTIVE:\n"
-    if subjective:
-        soap += '\n'.join(f"• {line}" for line in subjective[:5])  # Limit to 5 most relevant
-    else:
-        soap += "• Patient complaints and reported symptoms as documented"
-    
-    soap += "\n\nOBJECTIVE:\n"
-    if objective:
-        soap += '\n'.join(f"• {line}" for line in objective[:5])
-    else:
-        soap += "• Physical examination findings and clinical observations as documented"
-    
-    soap += "\n\nASSESSMENT:\n"
-    if assessment:
-        soap += '\n'.join(f"• {line}" for line in assessment[:3])
-    else:
-        soap += "• Clinical assessment based on presenting symptoms and examination findings"
-    
-    soap += "\n\nPLAN:\n"
-    if plan:
-        soap += '\n'.join(f"• {line}" for line in plan[:5])
-    else:
-        soap += "• Treatment plan and follow-up care as clinically indicated"
-    
-    return soap
-
-# OCR Functions (same as before but optimized)
-def initialize_ocr():
-    """Initialize OCR reader for handwritten notes"""
-    try:
-        # Initialize with English and medical text optimization
-        reader = easyocr.Reader(['en'], gpu=torch.cuda.is_available())
-        print("✅ EasyOCR initialized for handwritten medical notes")
-        return reader
-    except Exception as e:
-        print(f"⚠️ EasyOCR initialization failed: {e}")
-        return None
-
-def extract_text_from_image(image, ocr_reader=None):
-    """Enhanced OCR for medical handwriting"""
-    if image is None:
-        return "❌ No image provided"
-    
-    try:
-        # Preprocess specifically for medical handwriting
-        processed_img = preprocess_medical_image(image)
-        
-        extracted_text = ""
-        
-        # Try EasyOCR (better for handwritten text)
-        if ocr_reader is not None:
-            try:
-                results = ocr_reader.readtext(processed_img, detail=0, paragraph=True)
-                if results:
-                    extracted_text = ' '.join(results)
-                    if len(extracted_text.strip()) > 10:
-                        return clean_medical_text(extracted_text)
-            except Exception as e:
-                print(f"EasyOCR failed: {e}")
-        
-        # Fallback to Tesseract with medical optimization
-        try:
-            import pytesseract
-            
-            # Medical-optimized Tesseract config
-            custom_config = r'--oem 3 --psm 6 -c tessedit_char_whitelist=ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789.,;:()[]{}/-+= '
-            
-            tesseract_text = pytesseract.image_to_string(processed_img, config=custom_config)
-            
-            if len(tesseract_text.strip()) > 5:
-                return clean_medical_text(tesseract_text)
-                
-        except Exception as e:
-            print(f"Tesseract failed: {e}")
-        
-        return "❌ Could not extract text from image. Please ensure the image is clear and try again."
-        
-    except Exception as e:
-        return f"❌ Error processing image: {str(e)}"
-
-def preprocess_medical_image(image):
-    """Optimized preprocessing for medical handwriting"""
-    try:
-        img_array = np.array(image)
-        
-        # Convert to grayscale
-        if len(img_array.shape) == 3:
-            gray = cv2.cvtColor(img_array, cv2.COLOR_RGB2GRAY)
-        else:
-            gray = img_array
-        
-        # Resize for optimal OCR (medical notes are often small)
-        height, width = gray.shape
-        if height < 400 or width < 400:
-            scale_factor = max(400/height, 400/width)
-            new_width = int(width * scale_factor)
-            new_height = int(height * scale_factor)
-            gray = cv2.resize(gray, (new_width, new_height), interpolation=cv2.INTER_CUBIC)
-        
-        # Advanced preprocessing for handwritten medical text
-        # 1. Noise reduction
-        denoised = cv2.fastNlMeansDenoising(gray)
-        
-        # 2. Contrast enhancement specifically for handwriting
-        clahe = cv2.createCLAHE(clipLimit=3.0, tileGridSize=(8,8))
-        enhanced = clahe.apply(denoised)
-        
-        # 3. Morphological operations to clean up text
-        kernel = cv2.getStructuringElement(cv2.MORPH_RECT, (1,1))
-        cleaned = cv2.morphologyEx(enhanced, cv2.MORPH_CLOSE, kernel)
-        
-        # 4. Adaptive thresholding (better for varying lighting)
-        thresh = cv2.adaptiveThreshold(
-            cleaned, 255, cv2.ADAPTIVE_THRESH_GAUSSIAN_C, cv2.THRESH_BINARY, 11, 2
-        )
-        
-        return thresh
-        
-    except Exception as e:
-        print(f"❌ Image preprocessing error: {e}")
-        return np.array(image)
-
-def clean_medical_text(text):
-    """Clean extracted text with medical context awareness"""
-    # Remove excessive whitespace and empty lines
-    lines = [line.strip() for line in text.split('\n') if line.strip()]
-    
-    # Medical text cleaning
-    cleaned_lines = []
-    for line in lines:
-        # Remove obvious OCR artifacts
-        line = line.replace('|', 'l').replace('_', ' ').replace('~', '-')
-        
-        # Fix common medical abbreviations that OCR might misread
-        medical_corrections = {
-            'BP': 'BP', 'HR': 'HR', 'RR': 'RR', 'O2': 'O2',
-            'mg': 'mg', 'ml': 'ml', 'cc': 'cc', 'cm': 'cm'
-        }
-        
-        for wrong, correct in medical_corrections.items():
-            line = line.replace(wrong.lower(), correct)
-        
-        if len(line) > 1:  # Filter out single characters
-            cleaned_lines.append(line)
-    
-    return '\n'.join(cleaned_lines)
-
-# Enhanced Gradio Interface
-def gradio_generate_soap(medical_notes, uploaded_image, model_data):
-    """Main Gradio interface function"""
-    model, tokenizer = model_data if model_data else (None, None)
-    ocr_reader = getattr(gradio_generate_soap, 'ocr_reader', None)
-    
-    text_to_process = medical_notes.strip() if medical_notes else ""
-    
-    # Process uploaded image with enhanced OCR
-    if uploaded_image is not None:
-        try:
-            print("🔍 Extracting text from medical image...")
-            extracted_text = extract_text_from_image(uploaded_image, ocr_reader)
-            
-            if not extracted_text.startswith("❌"):
-                if not text_to_process:
-                    text_to_process = f"--- Extracted from uploaded image ---\n{extracted_text}"
-                else:
-                    text_to_process = f"{text_to_process}\n\n--- Additional text from image ---\n{extracted_text}"
-            else:
-                return extracted_text
-                
-        except Exception as e:
-            return f"❌ Error processing image: {str(e)}"
-    
-    if not text_to_process:
-        return "❌ Please enter medical notes manually or upload an image with medical text"
-    
-    # Generate SOAP note using Gemma 3n
-    try:
-        return generate_soap_note_gemma(text_to_process, model, tokenizer)
-    except Exception as e:
-        return f"❌ Error generating SOAP note: {str(e)}"
-
-# Example medical notes for testing
-medical_examples = {
-    'chest_pain': """Patient: John Smith, 45yo M
-CC: Chest pain x 2 hours
-HPI: Sudden onset sharp substernal chest pain 7/10, radiating to L arm. Associated SOB, diaphoresis. No N/V.
-PMH: HTN, no CAD
-VS: BP 150/90, HR 110, RR 22, O2 96% RA
-PE: Anxious, diaphoretic. RRR, no murmur. CTAB. No edema.
-A: Acute chest pain, r/o MI
-P: EKG, troponins, CXR, ASA 325mg, monitor""",
-
-    'diabetes': """Patient: Maria Garcia, 52yo F
-CC: Increased thirst, urination x 3 weeks  
-HPI: Polyuria, polydipsia, 10lb weight loss. FH DM. No fever, abd pain.
-VS: BP 140/85, HR 88, BMI 28
-PE: Mild dehydration, dry MM. RRR. No diabetic foot changes.
-Labs: Random glucose 280, HbA1c pending
-A: New onset DM Type 2
-P: HbA1c, CMP, diabetic education, metformin, f/u 2 weeks""",
-
-    'pediatric': """Patient: Emma Thompson, 8yo F
-CC: Fever, sore throat x 2 days
-HPI: Fever 102F, sore throat, odynophagia, decreased appetite. No cough/rhinorrhea.
-VS: T 101.8F, HR 110, RR 20, O2 99%
-PE: Alert, mildly ill. Throat erythematous w/ tonsillar exudate. Anterior cervical LAD.
-A: Strep pharyngitis (probable)
-P: Rapid strep, throat culture, amoxicillin if +, supportive care, RTC PRN"""
-}
-
-# Initialize everything
-def initialize_app():
-    """Initialize the complete application"""
-    print("🚀 Initializing Scribbled Docs SOAP Generator...")
-    
-    # Setup device
-    device = setup_device()
-    
-    # Load model
-    model, tokenizer = load_unsloth_gemma_model(device)
-    
-    # Initialize OCR
-    ocr_reader = initialize_ocr()
-    gradio_generate_soap.ocr_reader = ocr_reader
-    
-    return model, tokenizer
-
-# Create the main Gradio interface
-def create_interface(model, tokenizer):
-    """Create the main Gradio interface"""
-    
-    interface = gr.Interface(
-        fn=lambda notes, image: gradio_generate_soap(notes, image, (model, tokenizer)),
-        inputs=[
-            gr.Textbox(
-                lines=8,
-                placeholder="Enter medical notes here...\n\nExample:\nPatient: John Doe, 45yo M\nCC: Chest pain\nVS: BP 140/90, HR 88\n...",
-                label="📝 Medical Notes (Manual Entry)",
-                info="Enter unstructured medical notes or upload an image below"
-            ),
-            gr.Image(
-                type="pil",
-                label="📷 Upload Medical Image (Handwritten/Typed Notes)",
-                sources=["upload", "webcam"],
-                info="Upload PNG/JPG images of medical notes - handwritten or typed"
-            )
-        ],
-        outputs=[
-            gr.Textbox(
-                lines=20,
-                label="📋 Generated SOAP Note",
-                show_copy_button=True,
-                info="Professional SOAP note generated from your input"
-            )
-        ],
-        title="🏥 Scribbled Docs - Medical SOAP Note Generator",
-        description="""
-        **Transform medical notes into professional SOAP documentation using Gemma 3n AI**
-        
-        🔒 **100% Offline & HIPAA Compliant** - All processing happens locally on your device
-        🤖 **Powered by Unsloth-optimized Gemma 3n** - 4-bit quantized for efficiency
-        📝 **Supports handwritten & typed notes** - Advanced OCR for medical handwriting
-        
-        **Instructions:**
-        1. Enter medical notes manually OR upload an image
-        2. Click Submit to generate a structured SOAP note
-        3. Copy the result for use in your medical records
-        
-        **Perfect for:** Emergency medicine, family practice, internal medicine, pediatrics
-        """,
-        examples=[
-            [medical_examples['chest_pain'], None],
-            [medical_examples['diabetes'], None], 
-            [medical_examples['pediatric'], None]
-        ],
-        theme=gr.themes.Soft(
-            primary_hue="blue",
-            secondary_hue="green"
-        ),
-        allow_flagging="never",
-        analytics_enabled=False
-    )
-    
-    return interface
-
-# Main execution
-if __name__ == "__main__":
-    try:
-        # Initialize app
-        model, tokenizer = initialize_app()
-        
-        # Create and launch interface
-        interface = create_interface(model, tokenizer)
-        
-        print("\n🎯 Scribbled Docs SOAP Generator Ready!")
-        print("📱 Features:")
-        print("   ✅ Offline processing (HIPAA compliant)")
-        print("   ✅ Unsloth-optimized Gemma 3n model")
-        print("   ✅ Handwritten note OCR")
-        print("   ✅ Professional SOAP formatting")
-        print("   ✅ Medical terminology aware")
-        
-        # Launch interface
-        interface.launch(
-            share=True,  # Creates public link
-            server_port=7860,
-            show_error=True,
-            quiet=False
-        )
-        
-    except Exception as e:
-        print(f"❌ Error launching application: {e}")
-        print("💡 Make sure you have installed: pip install unsloth gradio easyocr opencv-python")
+import gradio as gr
+import torch
+import numpy as np
+from PIL import Image
+import time
+import io
+import subprocess
+import sys
+
+# Install required packages
+def install_packages():
+    packages = [
+        "transformers",
+        "accelerate", 
+        "timm",
+        "easyocr"
+    ]
+    for package in packages:
+        try:
+            subprocess.check_call([sys.executable, "-m", "pip", "install", package])
+        except:
+            print(f"Warning: Could not install {package}")
+
+# Install packages at startup
+install_packages()
+
+from transformers import AutoProcessor, AutoModelForImageTextToText, AutoConfig
+
+# Global variables for model
+processor = None
+model = None
+config = None
+ocr_reader = None
+
+def load_model():
+    """Load the Gemma 3n model"""
+    global processor, model, config, ocr_reader
+    
+    try:
+        print("🚀 Loading Gemma 3n model...")
+        GEMMA_PATH = "google/gemma-3n-e2b-it"
+        
+        # Load configuration
+        config = AutoConfig.from_pretrained(GEMMA_PATH, trust_remote_code=True)
+        print("✅ Config loaded")
+        
+        # Load processor
+        processor = AutoProcessor.from_pretrained(GEMMA_PATH, trust_remote_code=True)
+        print("✅ Processor loaded")
+        
+        # Load model
+        model = AutoModelForImageTextToText.from_pretrained(
+            GEMMA_PATH,
+            config=config,
+            torch_dtype="auto",
+            device_map="auto",
+            trust_remote_code=True
+        )
+        print("✅ Model loaded successfully!")
+        
+        # Set up compilation fix
+        import torch._dynamo
+        torch._dynamo.config.suppress_errors = True
+        
+        # Initialize OCR
+        try:
+            import easyocr
+            ocr_reader = easyocr.Reader(['en'], gpu=False, verbose=False)
+            print("✅ EasyOCR initialized")
+        except Exception as e:
+            print(f"⚠️ EasyOCR not available: {e}")
+            ocr_reader = None
+            
+        return True
+        
+    except Exception as e:
+        print(f"❌ Model loading failed: {e}")
+        return False
+
+def generate_soap_note(text):
+    """Generate SOAP note using Gemma 3n"""
+    if model is None or processor is None:
+        return "❌ Model not loaded. Please wait for initialization."
+    
+    soap_prompt = f"""You are a medical AI assistant. Convert the following medical notes into a properly formatted SOAP note.
+
+Medical notes:
+{text}
+
+Please format as:
+S - SUBJECTIVE: (chief complaint, history of present illness, past medical history, medications, allergies)
+O - OBJECTIVE: (vital signs, physical examination findings)  
+A - ASSESSMENT: (diagnosis/clinical impression)
+P - PLAN: (treatment plan, follow-up instructions)
+
+Generate a complete, professional SOAP note:"""
+    
+    messages = [{
+        "role": "system",
+        "content": [{"type": "text", "text": "You are an expert medical AI assistant specialized in creating SOAP notes from medical documentation."}]
+    }, {
+        "role": "user", 
+        "content": [{"type": "text", "text": soap_prompt}]
+    }]
+    
+    try:
+        inputs = processor.apply_chat_template(
+            messages,
+            add_generation_prompt=True,
+            tokenize=True,
+            return_dict=True,
+            return_tensors="pt"
+        ).to(model.device)
+        
+        input_len = inputs["input_ids"].shape[-1]
+        
+        with torch.no_grad():
+            outputs = model.generate(
+                **inputs,
+                max_new_tokens=400,
+                do_sample=True,
+                temperature=0.1,
+                top_p=0.95,
+                pad_token_id=processor.tokenizer.eos_token_id,
+                disable_compile=True
+            )
+        
+        response = processor.batch_decode(
+            outputs[:, input_len:],
+            skip_special_tokens=True
+        )[0].strip()
+        
+        return response
+        
+    except Exception as e:
+        return f"❌ SOAP generation failed: {str(e)}"
+
+def extract_text_from_image(image):
+    """Extract text using EasyOCR"""
+    if ocr_reader is None:
+        return "❌ OCR not available"
+    
+    try:
+        if hasattr(image, 'convert'):
+            image = image.convert('RGB')
+        img_array = np.array(image)
+        
+        results = ocr_reader.readtext(img_array, detail=0, paragraph=True)
+        if results:
+            return ' '.join(results).strip()
+        else:
+            return "❌ No text detected in image"
+            
+    except Exception as e:
+        return f"❌ OCR failed: {str(e)}"
+
+def process_medical_input(image, text):
+    """Main processing function for the Gradio interface"""
+    
+    if image is not None and text.strip():
+        return "⚠️ Please provide either an image OR text, not both.", ""
+    
+    if image is not None:
+        # Process image
+        print("🔍 Extracting text from image...")
+        extracted_text = extract_text_from_image(image)
+        
+        if extracted_text.startswith('❌'):
+            return extracted_text, ""
+        
+        print("🤖 Generating SOAP note...")
+        soap_note = generate_soap_note(extracted_text)
+        
+        return extracted_text, soap_note
+        
+    elif text.strip():
+        # Process text directly
+        print("🤖 Generating SOAP note from text...")
+        soap_note = generate_soap_note(text.strip())
+        return text.strip(), soap_note
+        
+    else:
+        return "❌ Please provide either an image or text input.", ""
+
+def create_demo():
+    """Create the Gradio demo interface"""
+    
+    # Sample text for demonstration
+    sample_text = """Patient: John Smith, 45yo male
+CC: Chest pain
+Vitals: BP 140/90, HR 88, RR 16, O2 98%, Temp 98.6F
+HPI: Patient reports crushing chest pain x 2 hours, radiating to left arm
+PMH: HTN, DM Type 2
+Current Meds: Lisinopril 10mg daily, Metformin 500mg BID
+PE: Diaphoretic, anxious appearance
+EKG: ST elevation in leads II, III, aVF"""
+    
+    with gr.Blocks(title="Medical OCR SOAP Generator", theme=gr.themes.Soft()) as demo:
+        
+        gr.Markdown("""
+        # 🏥 Medical OCR SOAP Generator
+        ### Powered by Gemma 3n - Convert handwritten medical notes to professional SOAP format
+        
+        **Instructions:**
+        - **Option 1:** Upload an image of handwritten medical notes
+        - **Option 2:** Enter medical text directly
+        - The system will generate a properly formatted SOAP note
+        
+        ⚠️ **Note:** First generation may take ~60-90 seconds as the model loads
+        """)
+        
+        with gr.Row():
+            with gr.Column():
+                image_input = gr.Image(
+                    type="pil", 
+                    label="📷 Upload Medical Image",
+                    height=300
+                )
+                
+                text_input = gr.Textbox(
+                    label="📝 Or Enter Medical Text",
+                    placeholder=sample_text,
+                    lines=8,
+                    max_lines=15
+                )
+                
+                submit_btn = gr.Button(
+                    "Generate SOAP Note",
+                    variant="primary",
+                    size="lg"
+                )
+                
+            with gr.Column():
+                extracted_output = gr.Textbox(
+                    label="📋 Extracted/Input Text",
+                    lines=6,
+                    max_lines=10
+                )
+                
+                soap_output = gr.Textbox(
+                    label="🏥 Generated SOAP Note",
+                    lines=12,
+                    max_lines=20
+                )
+        
+        # Example section
+        gr.Markdown("### 📋 Quick Test Example")
+        example_btn = gr.Button("Try Sample Medical Text", variant="secondary")
+        
+        def load_example():
+            return sample_text, None
+            
+        example_btn.click(
+            load_example,
+            outputs=[text_input, image_input]
+        )
+        
+        # Process function
+        submit_btn.click(
+            process_medical_input,
+            inputs=[image_input, text_input],
+            outputs=[extracted_output, soap_output]
+        )
+        
+        gr.Markdown("""
+        ---
+        **About:** This application uses Google's Gemma 3n model for medical text understanding and EasyOCR for handwriting recognition. 
+        All processing is done locally for HIPAA compliance.
+        
+        **Competition Entry:** Medical AI Innovation Challenge 2024
+        """)
+    
+    return demo
+
+# Initialize the application
+if __name__ == "__main__":
+    print("🚀 Starting Medical OCR SOAP Generator...")
+    
+    # Load model
+    model_loaded = load_model()
+    
+    if model_loaded:
+        print("✅ All systems ready!")
+        demo = create_demo()
+        demo.launch(
+            share=True,
+            server_name="0.0.0.0",
+            server_port=7860
+        )
+    else:
+        print("❌ Failed to load model. Creating fallback demo...")
+        
+        def fallback_demo():
+            return "❌ Model loading failed. Please check the logs.", "❌ Model not available."
+        
+        demo = gr.Interface(
+            fn=fallback_demo,
+            inputs=[
+                gr.Image(type="pil", label="Upload Medical Image"),
+                gr.Textbox(label="Enter Medical Text", lines=5)
+            ],
+            outputs=[
+                gr.Textbox(label="Status"),
+                gr.Textbox(label="Error Message")
+            ],
+            title="❌ Medical OCR - Model Loading Failed"
+        )
+        
+        demo.launch(share=True)