Update app.py

app.py

@@ -60,6 +60,10 @@ class TempImageFile:
         try:
             if self.path and os.path.exists(self.path):
                 os.remove(self.path)
+                # Also remove parent directory if empty
+                parent_dir = os.path.dirname(self.path)
+                if os.path.exists(parent_dir) and not os.listdir(parent_dir):
+                    os.rmdir(parent_dir)
         except Exception as e:
             print(f"Error cleaning up temporary file: {e}")
         finally:
@@ -160,130 +164,15 @@ MEASUREMENT_RANGES = {
     "astigmatism": (0.0, 6.0)     # D
 }
 
-def get_case_specific_settings(measurements):
-    """Determine appropriate settings based on measurements"""
-    settings = DEFAULT_SETTINGS.copy()
-    
-    # Extract measurements (implement measurement extraction logic)
-    axial_length = measurements.get('axial_length', 23.5)
-    astigmatism = measurements.get('astigmatism', 0.0)
-    
-    # Adjust settings based on measurements
-    if axial_length < 22.0:
-        settings.update(LENS_MODELS['short_eye'])
-    elif astigmatism > 2.0:
-        settings.update(LENS_MODELS['toric'])
-    
-    return settings
-
-DESCRIPTION = """
-# IOL Report Analyzer
-Upload an IOL screenshot to get a detailed analysis of measurements and recommendations.
-
-## Instructions:
-1. Upload your IOL report image
-2. Click Submit to generate analysis
-3. Review the detailed output below
-"""
-
-SYSTEM_PROMPT = """Please analyze the uploaded IOL report image carefully and provide a detailed analysis using these settings:
-Manufacturer: {manufacturer}
-Lens Model: {lens_model}
-A-Constant: {a_constant}
-Target Refraction: {target_refraction} D
-Surgeon Factor: {surgeon_factor}
-Personalization Factor: {personalization_factor}
-AC Constant: {anterior_chamber_constant}
-Vertex Distance: {vertex_distance} mm
-
-Please provide analysis in the following format:
-
-1. BIOMETRY MEASUREMENTS:
-   - Axial Length (AL): [value] mm
-   - Keratometry: K1: [value] D @ [axis]°, K2: [value] D @ [axis]°
-   - Mean K: [value] D
-   - Astigmatism: [value] D
-   - Anterior Chamber Depth (ACD): [value] mm
-   - Lens Thickness (LT): [value] mm
-   - White-to-White Distance (WTW): [value] mm
-   - Central Corneal Thickness (CCT): [value] μm
-
-2. IOL POWER CALCULATIONS:
-   A. Barrett Universal II Formula:
-      - Recommended IOL Power: [value] D
-      - Predicted Refraction: [value] D
-      - Target Refraction: [value] D
-   
-   B. Other Formula Calculations:
-      - Cooke K6: [value] D
-      - EVO Formula: [value] D
-      - Hill-RBF Method: [value] D
-      - Hoffer QST: [value] D
-      - Kane Formula: [value] D
-      - PEARL-DGS: [value] D
-
-3. ANALYSIS:
-   - Optimal IOL Power Range: [range] D
-   - Recommended A-Constant: [value]
-   - Risk Factors: [list any concerns]
-   
-4. RECOMMENDATIONS:
-   - Primary IOL Choice: [power] D ([manufacturer/model])
-   - Alternative Options: [list alternatives]
-   - Special Considerations: [note any relevant factors]
-
-Please ensure all measurements are accurately extracted and calculations are precise. If any values are missing or unclear, please indicate this in your response."""
-
-css = """
-#component-0 {
-    max-width: 1200px;
-    margin: auto;
-    padding: 20px;
-}
-
-.container {
-    border: 1px solid #e0e0e0;
-    border-radius: 10px;
-    padding: 20px;
-    margin-bottom: 20px;
-    background: white;
-}
-
-.output-box {
-    height: 600px;
-    overflow: auto;
-    font-family: 'Source Code Pro', monospace;
-    line-height: 1.5;
-    padding: 15px;
-    background: #f8f9fa;
-    border-radius: 5px;
+ERROR_MESSAGES = {
+    "ValueError": "⚠️ Invalid Input",
+    "RuntimeError": "🔧 Processing Error",
+    "torch.cuda.OutOfMemoryError": "💾 Memory Error",
+    "Exception": "❌ Unexpected Error"
 }
 
-.title {
-    font-size: 24px;
-    font-weight: bold;
-    margin-bottom: 20px;
-    color: #2c3e50;
-}
-
-.submit-btn {
-    background: #2ecc71 !important;
-    border: none !important;
-}
-
-.submit-btn:hover {
-    background: #27ae60 !important;
-}
-"""
-
 def format_error(error_type, message):
-    error_messages = {
-        "ValueError": "⚠️ Invalid Input",
-        "RuntimeError": "🔧 Processing Error",
-        "torch.cuda.OutOfMemoryError": "💾 Memory Error",
-        "Exception": "❌ Unexpected Error"
-    }
-    return f"{error_messages.get(error_type.__name__, '❌ Error')}: {message}"
+    return f"{ERROR_MESSAGES.get(error_type.__name__, '❌ Error')}: {message}"
 
 def validate_image(image_array):
     if image_array is None:
@@ -337,16 +226,6 @@ def normalize_image_size(image, max_size=1024):
         image = image.resize(new_size, Image.LANCZOS)
     return image
 
-def validate_examples():
-    example_files = ["example1.png", "example2.png"]
-    valid_examples = []
-    
-    for file in example_files:
-        if os.path.exists(file):
-            valid_examples.append([file])
-    
-    return valid_examples if valid_examples else None
-
 def extract_measurements(image_array):
     """
     Extract measurements from the image using OCR or vision model
@@ -466,280 +345,250 @@ EXTRACTED MEASUREMENTS:
 """
 
 def generate_analysis_prompt(measurements, settings, warnings=None, risk_factors=None, quality_score=None):
-    """Enhanced prompt generation with quality indicators"""
-    quality_warning = ""
-    if quality_score is not None:
-        if quality_score < 0.5:
-            quality_warning = "\n⚠️ WARNING: Low image quality may affect measurement accuracy"
-            
-    risk_text = ""
-    if risk_factors:
-        risk_text = "\nRISK FACTORS:\n" + "\n".join(f"• {r}" for r in risk_factors)
-
-    settings_text = f"""
-ANALYSIS SETTINGS:
-╔═════════════════��══════╦═══════════╗
-║ Setting               ║ Value     ║
-╠════════════════════════╬═══════════╣
-║ Manufacturer          ║ {settings['manufacturer']} ║
-║ Lens Model            ║ {settings['lens_model']} ║
-║ A-Constant            ║ {settings['a_constant']} ║
-║ Target Refraction     ║ {settings['target_refraction']} D ║
-║ Surgeon Factor        ║ {settings['surgeon_factor']} ║
-║ Personalization       ║ {settings['personalization_factor']} ║
-║ AC Constant           ║ {settings['anterior_chamber_constant']} ║
-║ Vertex Distance       ║ {settings['vertex_distance']} mm ║
-╚════════════════════════╩═══════════╝
-"""
-
-    measurements_text = format_measurements_text(measurements) if measurements else "\nNo measurements could be extracted."
-
-    prompt = f"""As an expert ophthalmologist, please analyze this IOL report carefully.
-
-{settings_text}
-{measurements_text}
-{risk_text}
-{quality_warning}
-
-Please provide a detailed analysis following this structure:
-
-1. MEASUREMENTS VALIDATION
-   - Verify all measurements are within normal ranges
-   - Flag any concerning values
-
-2. IOL POWER CALCULATIONS
-   - Primary Formula Recommendation
-   - Cross-validation with other formulas
-   - Optimal power range
-
-3. RISK ASSESSMENT
-   - Identify potential complications
-   - Special considerations
-
-4. RECOMMENDATIONS
-   - Specific IOL choice and power
-   - Alternative options
-   - Surgical approach suggestions
-
-Please be thorough and highlight any concerns or special considerations."""
-
+    """Enhanced prompt generation for IOL Master and similar biometry reports"""
+    
+    # Detect case type
+    case_types = detect_case_type(measurements)
+    
+    # Build comprehensive analysis request
+    prompt = f"""Analyze this IOL biometry report image.
+
+TASK:
+1. Verify all visible measurements
+2. Confirm measurement accuracy and consistency
+3. Provide IOL power calculation using:
+   - SRK/T formula
+   - Barrett Universal II
+   - Hill-RBF if applicable
+   - Other formulas as appropriate for this case
+
+CONSIDERATIONS:
+1. Image Quality Assessment
+2. Measurement Validation
+3. Formula Selection Rationale
+4. Risk Factor Analysis
+
+{format_measurements_text(measurements)}
+
+SELECTED IOL PARAMETERS:
+• Manufacturer: {settings.get('manufacturer')}
+• Model: {settings.get('lens_model')}
+• A-Constant: {settings.get('a_constant')}
+• Target: {settings.get('target_refraction', 0)} D
+
+Please provide:
+1. Measurement verification
+2. IOL power calculations with multiple formulas
+3. Specific recommendations for:
+   - Optimal IOL power
+   - Formula selection rationale
+   - Surgical approach
+4. Risk assessment and special considerations
+5. Alternative options if needed"""
+
+    # Add case-specific guidance if detected
+    if case_types:
+        prompt += "\n\nSpecial Considerations:"
+        for _, description in case_types:
+            prompt += f"\n• {description}"
+    
     return prompt
 
 def validate_settings(settings):
+    """Enhanced settings validation with better error handling"""
     if not isinstance(settings, dict):
-        raise ValueError("Settings must be a dictionary")
+        return DEFAULT_SETTINGS  # Fallback to defaults instead of raising error
     
-    required_fields = {
-        'manufacturer': str,
-        'lens_model': str,
-        'a_constant': (int, float),
-        'target_refraction': (int, float),
-        'surgeon_factor': (int, float),
-        'personalization_factor': (int, float),
-        'anterior_chamber_constant': (int, float),
-        'vertex_distance': (int, float)
-    }
+    validated_settings = DEFAULT_SETTINGS.copy()
+    for key, value in settings.items():
+        if key in validated_settings:
+            try:
+                # Type conversion and validation
+                if isinstance(DEFAULT_SETTINGS[key], (int, float)):
+                    value = float(value)
+                    # Validate ranges
+                    if key == 'a_constant' and not 115 <= value <= 122:
+                        continue  # Skip invalid value, keep default
+                    if key == 'target_refraction' and not -10 <= value <= 10:
+                        continue  # Skip invalid value, keep default
+                validated_settings[key] = value
+            except (ValueError, TypeError):
+                continue  # Keep default value if conversion fails
     
-    for field, expected_type in required_fields.items():
-        if field not in settings:
-            raise ValueError(f"Missing required setting: {field}")
-        if not isinstance(settings[field], expected_type):
-            raise ValueError(f"Invalid type for {field}: expected {expected_type}, got {type(settings[field])}")
-            
-        # Validate numeric ranges
-        if isinstance(settings[field], (int, float)):
-            if field == 'a_constant' and not 115 <= settings[field] <= 122:
-                raise ValueError("A-constant must be between 115 and 122")
-            elif field == 'target_refraction' and not -10 <= settings[field] <= 10:
-                raise ValueError("Target refraction must be between -10 and 10")
-
-@spaces.GPU(duration=120)
+    return validated_settings
+
+@spaces.GPU
 def run_example(image, model_id="Qwen/Qwen2-VL-7B-Instruct", settings=None, progress=gr.Progress()):
+    """Enhanced error handling and resource management"""
+    cleanup_temp_files()
+    model = None
+    
     try:
-        if settings:
-            validate_settings(settings)
-            
-        # Enhanced validation flow with quality feedback
-        validate_image(image)
-        quality_score, quality_feedback = validate_image_quality(image)
-        
-        if quality_score < 0.3:
-            return f"⚠️ Image quality too low for reliable analysis: {quality_feedback}"
-            
-        progress(0.2, "Extracting measurements...")
-        measurements = extract_measurements(image)
-        
-        # Add error handling for measurement extraction
-        if not measurements:
-            return "❌ Failed to extract measurements from image. Please ensure the image contains clear IOL measurements."
+        progress(0, desc="Validating input...")
+        if image is None:
+            return "Error: No image provided"
             
-        warnings, risk_factors = validate_measurements(measurements)
+        image = validate_image(image)
+        settings = validate_settings(settings or {})
         
-        if not settings:
-            settings = get_case_specific_settings(measurements)
-            
-        analysis_prompt = generate_analysis_prompt(
-            measurements, 
-            settings, 
-            warnings=warnings,
-            risk_factors=risk_factors,
-            quality_score=quality_score
-        )
-        
-        progress(0.4, "Loading model...")
-        model = model_manager.get_model(model_id)
-        processor = model_manager.get_processor(model_id)
+        progress(0.2, desc="Loading model...")
+        try:
+            model = model_manager.get_model(model_id)
+            processor = model_manager.get_processor(model_id)
+        except Exception as e:
+            return f"Error loading model: {str(e)}"
         
-        progress(0.6, "Processing image...")
-        # Save temporary image for model processing
-        with TempImageFile(image) as temp_path:
-            # Process image and generate response
-            inputs = processor(text=analysis_prompt, images=temp_path, return_tensors="pt")
-            inputs = {k: v.to(model.device) for k, v in inputs.items()}
+        progress(0.4, desc="Processing image...")
+        with TempImageFile(image) as image_path:
+            if not image_path:
+                return "Error: Failed to save temporary image"
+                
+            # Convert and normalize image
+            try:
+                image = Image.fromarray(image).convert("RGB")
+                image = normalize_image_size(image)
+            except Exception as e:
+                return f"Error processing image: {str(e)}"
             
-            progress(0.8, "Generating analysis...")
-            with torch.no_grad():
-                output = model.generate(
-                    **inputs,
-                    max_new_tokens=1000,
-                    do_sample=True,
-                    temperature=0.7,
-                    top_p=0.9,
+            # Generate analysis prompt
+            try:
+                measurements = extract_measurements(image)
+                warnings, risk_factors = validate_measurements(measurements)
+                quality_score, quality_feedback = validate_image_quality(image)
+                
+                prompt = generate_analysis_prompt(
+                    measurements=measurements,
+                    settings=settings,
+                    warnings=warnings,
+                    risk_factors=risk_factors,
+                    quality_score=quality_score
                 )
+            except Exception as e:
+                return f"Error analyzing image: {str(e)}"
             
-            response = processor.decode(output[0], skip_special_tokens=True)
-            
-            # Cleanup
-            model_manager.cleanup_unused_models(keep_model_id=model_id)
-            cleanup_temp_files()
+            # Prepare model inputs
+            try:
+                messages = [{
+                    "role": "user",
+                    "content": [
+                        {"type": "image", "image": image_path},
+                        {"type": "text", "text": prompt}
+                    ]
+                }]
+                
+                text = processor.apply_chat_template(
+                    messages, 
+                    tokenize=False, 
+                    add_generation_prompt=True
+                )
+                
+                image_inputs, video_inputs = process_vision_info(messages)
+                inputs = processor(
+                    text=[text],
+                    images=image_inputs,
+                    videos=video_inputs,
+                    padding=True,
+                    return_tensors="pt"
+                ).to(model.device)
+            except Exception as e:
+                return f"Error preparing model inputs: {str(e)}"
             
-            return response
-
-    except torch.cuda.OutOfMemoryError:
-        return "💾 GPU memory exceeded. Please try with a smaller image or wait a moment."
+            # Model inference with proper resource management
+            try:
+                progress(0.6, desc="Generating analysis...")
+                with torch.inference_mode():
+                    generated_ids = model.generate(
+                        **inputs,
+                        max_new_tokens=1024,
+                        do_sample=True,
+                        temperature=0.7,
+                        top_p=0.9
+                    )
+                
+                generated_ids_trimmed = [
+                    out_ids[len(in_ids):] 
+                    for in_ids, out_ids in zip(inputs.input_ids, generated_ids)
+                ]
+                
+                output_text = processor.batch_decode(
+                    generated_ids_trimmed,
+                    skip_special_tokens=True,
+                    clean_up_tokenization_spaces=False
+                )[0]
+                
+                progress(1.0, desc="Complete!")
+                return output_text
+                
+            except torch.cuda.OutOfMemoryError:
+                return "Error: GPU memory exhausted. Please try with a smaller image."
+            except Exception as e:
+                return f"Error during analysis: {str(e)}"
+                
     except Exception as e:
-        return format_error(type(e), str(e))
+        return f"Unexpected error: {str(e)}"
+        
+    finally:
+        # Ensure proper cleanup
+        try:
+            if model is not None:
+                model_manager.cleanup_unused_models(keep_model_id=model_id)
+            torch.cuda.empty_cache()
+            gc.collect()
+        except Exception as e:
+            print(f"Cleanup error: {e}")
 
 with gr.Blocks(css=css, theme=gr.themes.Soft()) as demo:
     gr.Markdown(DESCRIPTION)
     
     with gr.Row():
-        with gr.Column(scale=1):
-            with gr.Group(elem_classes="container"):
-                gr.Markdown("### Upload Image")
-                input_img = gr.Image(
-                    label="IOL Report Image",
-                    type="numpy",
-                    elem_classes="input-image"
+        with gr.Column():
+            # Image upload
+            input_img = gr.Image(
+                label="IOL Report Image",
+                type="numpy"
+            )
+            
+            # Essential settings only
+            with gr.Accordion("Settings", open=False):
+                manufacturer = gr.Dropdown(
+                    choices=["Alcon", "Johnson & Johnson", "Zeiss", "Bausch & Lomb"],
+                    value=DEFAULT_SETTINGS["manufacturer"],
+                    label="Manufacturer"
                 )
-                
-                # Add collapsible advanced settings
-                with gr.Accordion("Advanced Settings", open=False):
-                    manufacturer = gr.Dropdown(
-                        choices=["Alcon", "Johnson & Johnson", "Zeiss", "Bausch & Lomb"],
-                        value=DEFAULT_SETTINGS["manufacturer"],
-                        label="Manufacturer"
-                    )
-                    lens_model = gr.Textbox(
-                        value=DEFAULT_SETTINGS["lens_model"],
-                        label="Lens Model"
-                    )
-                    a_constant = gr.Number(
-                        value=DEFAULT_SETTINGS["a_constant"],
-                        label="A-Constant"
-                    )
-                    target_refraction = gr.Number(
-                        value=DEFAULT_SETTINGS["target_refraction"],
-                        label="Target Refraction (D)"
-                    )
-                    with gr.Row():
-                        surgeon_factor = gr.Number(
-                            value=DEFAULT_SETTINGS["surgeon_factor"],
-                            label="Surgeon Factor"
-                        )
-                        personalization_factor = gr.Number(
-                            value=DEFAULT_SETTINGS["personalization_factor"],
-                            label="Personalization Factor"
-                        )
-                    with gr.Row():
-                        ac_constant = gr.Number(
-                            value=DEFAULT_SETTINGS["anterior_chamber_constant"],
-                            label="AC Constant"
-                        )
-                        vertex_distance = gr.Number(
-                            value=DEFAULT_SETTINGS["vertex_distance"],
-                            label="Vertex Distance (mm)"
-                        )
-
-                model_selector = gr.Dropdown(
-                    choices=model_manager.available_models,
-                    value="Qwen/Qwen2-VL-7B-Instruct",
-                    label="Model Selection",
-                    interactive=True
+                lens_model = gr.Textbox(
+                    value=DEFAULT_SETTINGS["lens_model"],
+                    label="Lens Model"
                 )
-                submit_btn = gr.Button(
-                    "📊 Analyze Report",
-                    variant="primary",
-                    elem_classes="submit-btn"
+                a_constant = gr.Number(
+                    value=DEFAULT_SETTINGS["a_constant"],
+                    label="A-Constant"
                 )
-
-        with gr.Column(scale=1):
-            with gr.Group(elem_classes="container"):
-                gr.Markdown("### Analysis Results")
-                output_text = gr.Textbox(
-                    label="IOL Analysis",
-                    lines=25,
-                    show_copy_button=True,
-                    elem_classes="output-box"
+                target_refraction = gr.Number(
+                    value=DEFAULT_SETTINGS["target_refraction"],
+                    label="Target Refraction (D)"
                 )
-                
-    # Add example images if available
-    examples = validate_examples()
-    if examples:
-        gr.Examples(
-            examples=examples,
-            inputs=input_img,
-            outputs=output_text,
-            fn=run_example,
-            cache_examples=True,
-        )
-
-    # Event handler
+
+            submit_btn = gr.Button("Analyze Report", variant="primary")
+
+        # Output column
+        with gr.Column():
+            output_text = gr.Textbox(
+                label="Analysis Results",
+                lines=20,
+                show_copy_button=True
+            )
+
+    # Simplified event handler
     submit_btn.click(
         fn=run_example,
         inputs=[
             input_img,
-            model_selector,
-            gr.State(lambda: {
-                "manufacturer": manufacturer.value,
-                "lens_model": lens_model.value,
-                "a_constant": a_constant.value,
-                "target_refraction": target_refraction.value,
-                "surgeon_factor": surgeon_factor.value,
-                "personalization_factor": personalization_factor.value,
-                "anterior_chamber_constant": ac_constant.value,
-                "vertex_distance": vertex_distance.value
-            })
+            gr.State("Qwen/Qwen2-VL-7B-Instruct"),
+            gr.State(DEFAULT_SETTINGS)
         ],
-        outputs=output_text,
-        api_name="analyze"
+        outputs=output_text
     )
 
-# Modify the queue configuration
-demo.queue(
-    max_size=1,
-    api_open=False,
-    concurrency_count=1,  # Limit concurrent executions
-    status_update_rate=0.1,  # More frequent updates
-    timeout=600,  # 10 minute timeout
-    default_concurrency_limit=1
-)
-
-# Launch with queue
-demo.launch(
-    debug=True,
-    show_error=True,
-    share=False,
-    server_name="0.0.0.0",
-    server_port=7860,
-)
+# Simplified launch configuration
+demo.queue(concurrency_count=1)
+demo.launch(server_name="0.0.0.0", server_port=7860)