Update src/streamlit_app.py

src/streamlit_app.py

@@ -3,6 +3,9 @@ import numpy as np
 import os
 import sys
 from PIL import Image
+from scipy import ndimage
+import matplotlib.pyplot as plt
+from mpl_toolkits.mplot3d import Axes3D
 
 # Set environment variables to fix permission issues
 os.environ['MPLCONFIGDIR'] = '/tmp/matplotlib'
@@ -19,7 +22,6 @@ except ImportError:
 try:
     import matplotlib
     matplotlib.use('Agg')  # Use non-interactive backend
-    import matplotlib.pyplot as plt
     import matplotlib.cm as cm
     MPL_AVAILABLE = True
 except ImportError:
@@ -119,244 +121,214 @@ def load_stroke_model():
     except Exception as e:
         return None, f"❌ Model loading failed: {str(e)}"
 
-def analyze_model_architecture(model):
-    """Comprehensive analysis of model architecture."""
-    if model is None:
-        return {"error": "No model loaded"}
-    
-    layer_analysis = {
-        'total_layers': len(model.layers),
-        'conv_layers': [],
-        'dense_layers': [],
-        'other_layers': [],
-        'all_layers_detailed': [],
-        'model_type': 'Unknown'
-    }
+def analyze_heatmap_distribution(heatmap, name="Heatmap"):
+    """Analyze the distribution of heatmap values."""
+    if heatmap is None:
+        return None
     
-    for i, layer in enumerate(model.layers):
-        layer_type = type(layer).__name__
-        
-        # Get more detailed layer information
-        layer_info = {
-            'index': i,
-            'name': layer.name,
-            'type': layer_type,
-            'output_shape': getattr(layer, 'output_shape', 'Unknown'),
-            'trainable': getattr(layer, 'trainable', 'Unknown'),
-            'activation': getattr(layer, 'activation', None)
+    flat_values = heatmap.flatten()
+    
+    analysis = {
+        'name': name,
+        'shape': heatmap.shape,
+        'total_pixels': heatmap.size,
+        'min': float(np.min(flat_values)),
+        'max': float(np.max(flat_values)),
+        'mean': float(np.mean(flat_values)),
+        'median': float(np.median(flat_values)),
+        'std': float(np.std(flat_values)),
+        'range': float(np.max(flat_values) - np.min(flat_values)),
+        'unique_values': len(np.unique(flat_values)),
+        'zero_pixels': int(np.sum(flat_values == 0)),
+        'non_zero_pixels': int(np.sum(flat_values > 0)),
+        'percentiles': {
+            '1%': float(np.percentile(flat_values, 1)),
+            '5%': float(np.percentile(flat_values, 5)),
+            '25%': float(np.percentile(flat_values, 25)),
+            '75%': float(np.percentile(flat_values, 75)),
+            '95%': float(np.percentile(flat_values, 95)),
+            '99%': float(np.percentile(flat_values, 99))
         }
-        
-        # Try to get activation function name
-        if hasattr(layer, 'activation') and layer.activation:
-            try:
-                layer_info['activation'] = layer.activation.__name__
-            except:
-                layer_info['activation'] = str(layer.activation)
-        
-        layer_analysis['all_layers_detailed'].append(layer_info)
-        
-        # Categorize layers with more comprehensive detection
-        if any(conv_type in layer_type for conv_type in [
-            'Conv1D', 'Conv2D', 'Conv3D', 'SeparableConv2D', 'DepthwiseConv2D',
-            'Convolution1D', 'Convolution2D', 'Convolution3D'
-        ]) or 'conv' in layer.name.lower():
-            layer_analysis['conv_layers'].append(layer_info)
-        
-        elif 'Dense' in layer_type or 'Linear' in layer_type:
-            layer_analysis['dense_layers'].append(layer_info)
-        
-        else:
-            layer_analysis['other_layers'].append(layer_info)
+    }
     
-    # Determine model type
-    if layer_analysis['conv_layers']:
-        layer_analysis['model_type'] = 'CNN (Convolutional Neural Network)'
-    elif layer_analysis['dense_layers']:
-        layer_analysis['model_type'] = 'MLP (Multi-Layer Perceptron)'
+    # Determine if heatmap has good contrast
+    if analysis['range'] < 0.1:
+        analysis['contrast_quality'] = 'Very Poor (range < 0.1)'
+    elif analysis['range'] < 0.3:
+        analysis['contrast_quality'] = 'Poor (range < 0.3)'
+    elif analysis['range'] < 0.7:
+        analysis['contrast_quality'] = 'Moderate (range < 0.7)'
     else:
-        layer_analysis['model_type'] = 'Custom Architecture'
+        analysis['contrast_quality'] = 'Good (range >= 0.7)'
     
-    return layer_analysis
+    return analysis
 
-def debug_gradcam_step_by_step(img_array, model, layer_name, pred_index):
-    """Debug Grad-CAM computation step by step."""
-    debug_info = {
-        'step': 'Starting',
-        'error': None,
-        'layer_output_shape': None,
-        'gradients_shape': None,
-        'gradients_stats': None,
-        'heatmap_stats': None
-    }
+def force_contrast_enhancement(heatmap, method='aggressive'):
+    """Force better contrast in heatmap using various methods."""
+    if heatmap is None:
+        return None, "No heatmap provided"
     
-    try:
-        debug_info['step'] = 'Getting target layer'
-        target_layer = model.get_layer(layer_name)
-        debug_info['target_layer_type'] = type(target_layer).__name__
+    original_analysis = analyze_heatmap_distribution(heatmap, "Original")
+    
+    if method == 'aggressive':
+        # Method 1: Aggressive percentile stretching
+        p1, p99 = np.percentile(heatmap, [1, 99])
+        if p99 > p1:
+            enhanced = np.clip((heatmap - p1) / (p99 - p1), 0, 1)
+        else:
+            enhanced = heatmap
         
-        debug_info['step'] = 'Creating grad model'
-        grad_model = tf.keras.Model(
-            inputs=[model.inputs],
-            outputs=[target_layer.output, model.output]
-        )
+        # Apply power transformation to spread values
+        enhanced = np.power(enhanced, 0.3)  # Gamma < 1 spreads values
         
-        debug_info['step'] = 'Computing forward pass'
-        with tf.GradientTape() as tape:
-            layer_output, preds = grad_model(img_array)
-            debug_info['layer_output_shape'] = layer_output.shape.as_list()
-            debug_info['predictions_shape'] = preds.shape.as_list()
-            
-            if pred_index is None:
-                pred_index = tf.argmax(preds[0])
-            debug_info['pred_index'] = int(pred_index)
-            debug_info['pred_confidence'] = float(preds[0][pred_index])
-            
-            class_channel = preds[:, pred_index]
-            debug_info['class_channel_shape'] = class_channel.shape.as_list()
+    elif method == 'histogram_eq':
+        # Method 2: Histogram equalization
+        flat = heatmap.flatten()
+        hist, bins = np.histogram(flat, bins=256, range=(0, 1))
+        cdf = hist.cumsum()
+        cdf = cdf / cdf[-1]  # Normalize
         
-        debug_info['step'] = 'Computing gradients'
-        grads = tape.gradient(class_channel, layer_output)
+        # Interpolate to get new values
+        enhanced = np.interp(flat, bins[:-1], cdf).reshape(heatmap.shape)
         
-        if grads is None:
-            debug_info['error'] = "Gradients are None - no backpropagation path"
-            return None, debug_info
+    elif method == 'adaptive':
+        # Method 3: Adaptive enhancement based on local statistics
         
-        debug_info['gradients_shape'] = grads.shape.as_list()
-        debug_info['gradients_stats'] = {
-            'min': float(tf.reduce_min(grads)),
-            'max': float(tf.reduce_max(grads)),
-            'mean': float(tf.reduce_mean(grads)),
-            'std': float(tf.math.reduce_std(grads))
-        }
+        # Local mean and std
+        local_mean = ndimage.uniform_filter(heatmap, size=20)
+        local_std = ndimage.generic_filter(heatmap, np.std, size=20)
         
-        debug_info['step'] = 'Processing gradients based on layer type'
+        # Enhance based on local statistics
+        enhanced = (heatmap - local_mean) / (local_std + 1e-8)
+        enhanced = np.clip(enhanced, -3, 3)  # Clip outliers
+        enhanced = (enhanced + 3) / 6  # Normalize to [0, 1]
         
-        if len(layer_output.shape) == 4:  # Conv layer
-            debug_info['processing_type'] = 'Convolutional layer (4D)'
-            pooled_grads = tf.reduce_mean(grads, axis=(0, 1, 2))
-            layer_output = layer_output[0]
-            heatmap = layer_output @ pooled_grads[..., tf.newaxis]
-            heatmap = tf.squeeze(heatmap)
-            
-        elif len(layer_output.shape) == 2:  # Dense layer
-            debug_info['processing_type'] = 'Dense layer (2D)'
-            # For dense layers, create spatial heatmap from gradient magnitude
-            grads_magnitude = tf.reduce_mean(tf.abs(grads))
-            # Create a simple spatial pattern
-            heatmap = tf.ones((14, 14)) * grads_magnitude
-            
-        else:
-            debug_info['error'] = f"Unsupported layer shape: {layer_output.shape}"
-            return None, debug_info
+    elif method == 'artificial_peaks':
+        # Method 4: Create artificial peaks for visualization
+        enhanced = heatmap.copy()
         
-        debug_info['step'] = 'Normalizing heatmap'
-        debug_info['raw_heatmap_stats'] = {
-            'min': float(tf.reduce_min(heatmap)),
-            'max': float(tf.reduce_max(heatmap)),
-            'mean': float(tf.reduce_mean(heatmap)),
-            'std': float(tf.math.reduce_std(heatmap))
-        }
+        # Find top 10% of values and enhance them
+        threshold = np.percentile(enhanced, 90)
+        mask = enhanced >= threshold
+        enhanced[mask] = enhanced[mask] * 2
         
-        # Apply ReLU (remove negative values)
-        heatmap = tf.maximum(heatmap, 0)
+        # Find bottom 10% and suppress them
+        threshold_low = np.percentile(enhanced, 10)
+        mask_low = enhanced <= threshold_low
+        enhanced[mask_low] = enhanced[mask_low] * 0.1
         
         # Normalize
-        heatmap_max = tf.reduce_max(heatmap)
-        if heatmap_max > 0:
-            heatmap = heatmap / heatmap_max
-        else:
-            debug_info['error'] = "All heatmap values are zero or negative"
-            return None, debug_info
+        enhanced = np.clip(enhanced, 0, 1)
         
-        debug_info['final_heatmap_stats'] = {
-            'min': float(tf.reduce_min(heatmap)),
-            'max': float(tf.reduce_max(heatmap)),
-            'mean': float(tf.reduce_mean(heatmap)),
-            'std': float(tf.math.reduce_std(heatmap))
-        }
-        
-        debug_info['step'] = 'Complete'
-        return heatmap.numpy(), debug_info
-        
-    except Exception as e:
-        debug_info['error'] = f"Exception in step '{debug_info['step']}': {str(e)}"
-        return None, debug_info
+    else:
+        enhanced = heatmap
+    
+    enhanced_analysis = analyze_heatmap_distribution(enhanced, f"Enhanced ({method})")
+    
+    return enhanced, f"Enhanced using {method}", original_analysis, enhanced_analysis
 
-def create_robust_gradcam_heatmap(img, model, predictions):
-    """Create Grad-CAM with comprehensive debugging."""
-    try:
-        # Preprocess image
-        img_resized = img.resize((224, 224))
-        img_array = np.array(img_resized, dtype=np.float32)
-        
-        # Handle grayscale
-        if len(img_array.shape) == 2:
-            img_array = np.stack([img_array] * 3, axis=-1)
-        
-        # Normalize and add batch dimension
-        img_array = np.expand_dims(img_array, axis=0) / 255.0
-        
-        # Get model analysis
-        analysis = analyze_model_architecture(model)
-        
-        # Try different layers in order of preference
-        layer_candidates = []
-        
-        # Add conv layers first
-        for layer in analysis['conv_layers']:
-            layer_candidates.append((layer['name'], f"Conv layer: {layer['name']}"))
-        
-        # Add other potentially suitable layers
-        for layer in analysis['all_layers_detailed']:
-            if (layer['type'] in ['Activation', 'BatchNormalization'] and 
-                isinstance(layer['output_shape'], (list, tuple)) and 
-                len(layer['output_shape']) == 4):
-                layer_candidates.append((layer['name'], f"4D layer: {layer['name']} ({layer['type']})"))
-        
-        # Try dense layers as last resort
-        if not layer_candidates:
-            for layer in analysis['dense_layers']:
-                layer_candidates.append((layer['name'], f"Dense layer: {layer['name']} (experimental)"))
-        
-        if not layer_candidates:
-            return None, "❌ No suitable layers found", None
-        
-        # Try each candidate layer
-        for layer_name, layer_desc in layer_candidates:
-            pred_index = np.argmax(predictions)
-            
-            heatmap, debug_info = debug_gradcam_step_by_step(
-                img_array, model, layer_name, pred_index
-            )
-            
-            if heatmap is not None:
-                # Resize heatmap to match input image size
-                if heatmap.shape[0] != 224 or heatmap.shape[1] != 224:
-                    heatmap_resized = tf.image.resize(
-                        heatmap[..., tf.newaxis], 
-                        (224, 224)
-                    ).numpy()[:, :, 0]
-                else:
-                    heatmap_resized = heatmap
-                
-                # Final statistics
-                stats = {
-                    'min': float(np.min(heatmap_resized)),
-                    'max': float(np.max(heatmap_resized)),
-                    'mean': float(np.mean(heatmap_resized)),
-                    'std': float(np.std(heatmap_resized))
-                }
-                
-                return heatmap_resized, f"✅ Grad-CAM successful using {layer_desc}", stats, debug_info
-            else:
-                # Continue to next layer if this one failed
-                continue
-        
-        # If all layers failed, return debug info from the last attempt
-        return None, f"❌ All layers failed. Last error: {debug_info.get('error', 'Unknown')}", None, debug_info
-        
-    except Exception as e:
-        return None, f"❌ Grad-CAM error: {str(e)}", None, {'error': str(e)}
+def create_diagnostic_heatmap_visualization(heatmap, title="Heatmap Analysis"):
+    """Create a comprehensive diagnostic visualization of the heatmap."""
+    if not MPL_AVAILABLE or heatmap is None:
+        return None
+    
+    fig, axes = plt.subplots(2, 3, figsize=(18, 12))
+    
+    # Original heatmap
+    im1 = axes[0, 0].imshow(heatmap, cmap='hot', vmin=0, vmax=1)
+    axes[0, 0].set_title(f"{title} - Hot Colormap")
+    plt.colorbar(im1, ax=axes[0, 0])
+    
+    # Different colormap
+    im2 = axes[0, 1].imshow(heatmap, cmap='viridis', vmin=0, vmax=1)
+    axes[0, 1].set_title(f"{title} - Viridis Colormap")
+    plt.colorbar(im2, ax=axes[0, 1])
+    
+    # High contrast version
+    im3 = axes[0, 2].imshow(heatmap, cmap='RdYlBu_r', vmin=np.min(heatmap), vmax=np.max(heatmap))
+    axes[0, 2].set_title(f"{title} - Auto-scaled")
+    plt.colorbar(im3, ax=axes[0, 2])
+    
+    # Histogram
+    axes[1, 0].hist(heatmap.flatten(), bins=50, alpha=0.7, color='blue')
+    axes[1, 0].set_title("Value Distribution")
+    axes[1, 0].set_xlabel("Heatmap Value")
+    axes[1, 0].set_ylabel("Frequency")
+    
+    # 3D surface plot
+    x = np.arange(heatmap.shape[1])
+    y = np.arange(heatmap.shape[0])
+    X, Y = np.meshgrid(x, y)
+    
+    ax_3d = fig.add_subplot(2, 3, 5, projection='3d')
+    surf = ax_3d.plot_surface(X[::8, ::8], Y[::8, ::8], heatmap[::8, ::8], 
+                             cmap='hot', alpha=0.8)
+    ax_3d.set_title("3D Surface View")
+    
+    # Statistics text
+    analysis = analyze_heatmap_distribution(heatmap)
+    stats_text = f"""
+    Shape: {analysis['shape']}
+    Range: {analysis['range']:.4f}
+    Mean: {analysis['mean']:.4f}
+    Std: {analysis['std']:.4f}
+    Unique values: {analysis['unique_values']}
+    Contrast: {analysis['contrast_quality']}
+    
+    Percentiles:
+    1%: {analysis['percentiles']['1%']:.4f}
+    25%: {analysis['percentiles']['25%']:.4f}
+    75%: {analysis['percentiles']['75%']:.4f}
+    99%: {analysis['percentiles']['99%']:.4f}
+    """
+    
+    axes[1, 2].text(0.1, 0.9, stats_text, transform=axes[1, 2].transAxes, 
+                   fontsize=10, verticalalignment='top', fontfamily='monospace')
+    axes[1, 2].set_title("Statistics")
+    axes[1, 2].axis('off')
+    
+    plt.tight_layout()
+    return fig
+
+def create_multiple_enhancement_comparison(heatmap):
+    """Compare different enhancement methods side by side."""
+    if not MPL_AVAILABLE or heatmap is None:
+        return None
+    
+    methods = ['aggressive', 'histogram_eq', 'adaptive', 'artificial_peaks']
+    enhanced_maps = {}
+    
+    for method in methods:
+        enhanced, _, _, _ = force_contrast_enhancement(heatmap, method)
+        enhanced_maps[method] = enhanced
+    
+    fig, axes = plt.subplots(2, 3, figsize=(18, 12))
+    
+    # Original
+    im0 = axes[0, 0].imshow(heatmap, cmap='hot', vmin=0, vmax=1)
+    axes[0, 0].set_title("Original Heatmap")
+    plt.colorbar(im0, ax=axes[0, 0])
+    
+    # Enhanced versions
+    positions = [(0, 1), (0, 2), (1, 0), (1, 1)]
+    
+    for i, (method, enhanced) in enumerate(enhanced_maps.items()):
+        row, col = positions[i]
+        im = axes[row, col].imshow(enhanced, cmap='hot', vmin=0, vmax=1)
+        axes[row, col].set_title(f"Enhanced: {method}")
+        plt.colorbar(im, ax=axes[row, col])
+    
+    # Comparison histogram
+    axes[1, 2].hist(heatmap.flatten(), bins=30, alpha=0.5, label='Original', color='blue')
+    for method, enhanced in enhanced_maps.items():
+        axes[1, 2].hist(enhanced.flatten(), bins=30, alpha=0.3, label=method)
+    axes[1, 2].set_title("Value Distributions")
+    axes[1, 2].legend()
+    axes[1, 2].set_xlabel("Value")
+    axes[1, 2].set_ylabel("Frequency")
+    
+    plt.tight_layout()
+    return fig
 
 def predict_stroke(img, model):
     """Predict stroke type from image."""
@@ -383,134 +355,40 @@ def predict_stroke(img, model):
     except Exception as e:
         return None, f"Prediction error: {str(e)}"
 
-def create_enhanced_simulated_heatmap(img, predictions):
-    """Create a more realistic simulated heatmap."""
-    try:
-        confidence = np.max(predictions)
-        predicted_class = np.argmax(predictions)
-        
-        # Create different patterns based on predicted class
-        if predicted_class == 0:  # Hemorrhagic
-            # Focus on center-left region
-            center_x, center_y = 80, 112
-        elif predicted_class == 1:  # Ischemic
-            # Focus on right side
-            center_x, center_y = 150, 112
-        else:  # No stroke
-            # Diffuse, low-intensity pattern
-            center_x, center_y = 112, 112
-        
-        # Create base pattern
+def create_test_heatmaps():
+    """Create test heatmaps with known patterns for comparison."""
+    test_maps = {}
+    
+    # Test 1: High contrast pattern
+    test_maps['high_contrast'] = np.zeros((224, 224))
+    test_maps['high_contrast'][50:150, 50:150] = 1.0
+    test_maps['high_contrast'][75:125, 75:125] = 0.0
+    
+    # Test 2: Gradient pattern
+    x = np.linspace(0, 1, 224)
+    y = np.linspace(0, 1, 224)
+    X, Y = np.meshgrid(x, y)
+    test_maps['gradient'] = X * Y
+    
+    # Test 3: Gaussian blobs
+    test_maps['gaussian'] = np.zeros((224, 224))
+    centers = [(60, 60), (160, 160), (60, 160)]
+    for cx, cy in centers:
         y, x = np.ogrid[:224, :224]
-        
-        # Primary focus area
-        mask1 = np.exp(-((x - center_x)**2 + (y - center_y)**2) / (2 * (40**2)))
-        
-        # Secondary areas
-        mask2 = np.exp(-((x - center_x + 30)**2 + (y - center_y + 20)**2) / (2 * (25**2)))
-        mask3 = np.exp(-((x - center_x - 20)**2 + (y - center_y - 30)**2) / (2 * (30**2)))
-        
-        # Combine patterns
-        heatmap = (mask1 * 0.8 + mask2 * 0.4 + mask3 * 0.3) * confidence
-        
-        # Add some noise for realism
-        np.random.seed(42)
-        noise = np.random.normal(0, 0.05, heatmap.shape)
-        heatmap = np.maximum(heatmap + noise, 0)
-        
-        # Normalize
-        if np.max(heatmap) > 0:
-            heatmap = heatmap / np.max(heatmap)
-        
-        stats = {
-            'min': float(np.min(heatmap)),
-            'max': float(np.max(heatmap)),
-            'mean': float(np.mean(heatmap)),
-            'std': float(np.std(heatmap))
-        }
-        
-        return heatmap, "⚠️ Using enhanced simulated heatmap", stats
-    except Exception as e:
-        return None, f"❌ Simulated heatmap error: {str(e)}", None
-
-def create_comprehensive_visualization(img, predictions, model, force_gradcam=True, colormap='hot'):
-    """Create comprehensive visualization with debugging."""
-    if not MPL_AVAILABLE:
-        return None, "❌ Matplotlib not available"
+        mask = np.exp(-((x - cx)**2 + (y - cy)**2) / (2 * 30**2))
+        test_maps['gaussian'] += mask
+    test_maps['gaussian'] = test_maps['gaussian'] / np.max(test_maps['gaussian'])
     
-    try:
-        # Resize image to 224x224
-        img_resized = img.resize((224, 224))
-        img_array = np.array(img_resized)
-        
-        heatmap = None
-        status_message = ""
-        stats = None
-        debug_info = None
-        
-        # Try Grad-CAM first
-        if force_gradcam and model is not None:
-            result = create_robust_gradcam_heatmap(img, model, predictions)
-            if result and len(result) >= 3:
-                heatmap, gradcam_status, stats = result[0], result[1], result[2]
-                if len(result) > 3:
-                    debug_info = result[3]
-                status_message = gradcam_status
-        
-        # Fallback to enhanced simulated if Grad-CAM failed
-        if heatmap is None:
-            result = create_enhanced_simulated_heatmap(img, predictions)
-            if result and len(result) == 3:
-                heatmap, sim_status, stats = result
-                if status_message:
-                    status_message += f" | {sim_status}"
-                else:
-                    status_message = sim_status
-        
-        if heatmap is None:
-            return None, "❌ Could not generate any heatmap", None, None
-        
-        # Create visualization
-        fig, axes = plt.subplots(1, 3, figsize=(15, 5))
-        
-        # 1. Original image
-        axes[0].imshow(img_array)
-        axes[0].set_title("Original Image", fontsize=12, fontweight='bold')
-        axes[0].axis('off')
-        
-        # 2. Heatmap only
-        im1 = axes[1].imshow(heatmap, cmap=colormap, vmin=0, vmax=1)
-        axes[1].set_title(f"Attention Heatmap ({colormap})", fontsize=12, fontweight='bold')
-        axes[1].axis('off')
-        plt.colorbar(im1, ax=axes[1], fraction=0.046, pad=0.04)
-        
-        # 3. Overlay
-        axes[2].imshow(img_array)
-        im2 = axes[2].imshow(heatmap, cmap=colormap, alpha=0.6, vmin=0, vmax=1, interpolation='bilinear')
-        
-        # Determine title based on success
-        if "✅ Grad-CAM successful" in status_message:
-            title = "🎯 Real AI Attention Overlay"
-            title_color = 'green'
-        else:
-            title = "🎨 Simulated Attention Overlay"
-            title_color = 'orange'
-        
-        axes[2].set_title(title, fontsize=12, fontweight='bold', color=title_color)
-        axes[2].axis('off')
-        plt.colorbar(im2, ax=axes[2], fraction=0.046, pad=0.04)
-        
-        plt.tight_layout()
-        
-        return fig, status_message, stats, debug_info
+    # Test 4: Low contrast (similar to your issue)
+    test_maps['low_contrast'] = np.random.normal(0.5, 0.05, (224, 224))
+    test_maps['low_contrast'] = np.clip(test_maps['low_contrast'], 0, 1)
     
-    except Exception as e:
-        return None, f"❌ Visualization error: {str(e)}", None, None
+    return test_maps
 
 # Main App
 def main():
     # Header
-    st.markdown('<h1 class="main-header">🧠 AI-Powered Stroke Classification System</h1>', unsafe_allow_html=True)
+    st.markdown('<h1 class="main-header">🧠 Heatmap Diagnostic System</h1>', unsafe_allow_html=True)
     
     # Auto-load model on startup
     if not st.session_state.model_loaded:
@@ -541,31 +419,40 @@ def main():
         else:
             st.markdown('<div class="status-box error">❌ Model Error</div>', unsafe_allow_html=True)
 
-    # Model status details
-    st.markdown(f'<div class="status-box info"><strong>Model Status:</strong> {st.session_state.model_status}</div>', unsafe_allow_html=True)
-
-    # Enhanced model architecture analysis
-    if st.session_state.model is not None:
-        with st.expander("🔍 Detailed Model Architecture Analysis"):
-            analysis = analyze_model_architecture(st.session_state.model)
+    # Test heatmaps section
+    st.markdown("### 🧪 Test Heatmap Patterns")
+    
+    test_maps = create_test_heatmaps()
+    
+    col1, col2 = st.columns(2)
+    
+    with col1:
+        st.write("**Test Pattern:**")
+        test_pattern = st.selectbox(
+            "Choose a test pattern",
+            list(test_maps.keys()),
+            help="Test different heatmap patterns to see how they display"
+        )
+        
+        if test_pattern:
+            test_heatmap = test_maps[test_pattern]
             
-            st.write("**📊 Model Summary:**")
-            st.write(f"- **Model Type:** {analysis['model_type']}")
-            st.write(f"- **Total Layers:** {analysis['total_layers']}")
-            st.write(f"- **Convolutional Layers:** {len(analysis['conv_layers'])}")
-            st.write(f"- **Dense Layers:** {len(analysis['dense_layers'])}")
-            st.write(f"- **Other Layers:** {len(analysis['other_layers'])}")
+            # Show diagnostic visualization
+            diagnostic_fig = create_diagnostic_heatmap_visualization(test_heatmap, f"Test: {test_pattern}")
+            if diagnostic_fig:
+                st.pyplot(diagnostic_fig)
+                plt.close()
+    
+    with col2:
+        st.write("**Enhancement Comparison:**")
+        if test_pattern:
+            test_heatmap = test_maps[test_pattern]
             
-            # Show detailed layer information
-            st.write("**🔍 All Layers (Detailed):**")
-            for layer in analysis['all_layers_detailed']:
-                activation_info = f" | Activation: {layer['activation']}" if layer['activation'] else ""
-                st.code(f"{layer['index']:2d}: {layer['name']} ({layer['type']}) | Shape: {layer['output_shape']}{activation_info}")
-
-    # Manual reload button
-    if st.button("🔄 Reload Model", help="Try to reload the model"):
-        st.session_state.model_loaded = False
-        st.rerun()
+            # Show enhancement comparison
+            comparison_fig = create_multiple_enhancement_comparison(test_heatmap)
+            if comparison_fig:
+                st.pyplot(comparison_fig)
+                plt.close()
 
     # Sidebar
     with st.sidebar:
@@ -577,144 +464,164 @@ def main():
         )
         
         st.markdown("---")
-        st.header("🎨 Visualization Options")
+        st.header("🎨 Enhancement Options")
         
-        force_gradcam = st.checkbox(
-            "Attempt Grad-CAM", 
-            value=True,
-            help="Try Grad-CAM with comprehensive debugging"
+        enhancement_method = st.selectbox(
+            "Enhancement Method",
+            ['none', 'aggressive', 'histogram_eq', 'adaptive', 'artificial_peaks'],
+            index=1,
+            help="Choose how to enhance heatmap contrast"
         )
         
-        colormap = st.selectbox(
-            "Color Scheme",
-            ['hot', 'jet', 'viridis', 'plasma', 'inferno', 'magma', 'coolwarm'],
-            index=0,
-            help="Choose color scheme for heatmap visualization"
-        )
-        
-        show_probabilities = st.checkbox("Show All Probabilities", value=True)
-        show_debug = st.checkbox("Show Debug Info", value=True)
-        show_stats = st.checkbox("Show Heatmap Statistics", value=True)
-        show_detailed_debug = st.checkbox("Show Detailed Debug Info", value=False)
+        show_diagnostics = st.checkbox("Show Diagnostic Analysis", value=True)
+        show_comparisons = st.checkbox("Show Enhancement Comparisons", value=True)
 
     if uploaded_file is not None:
         # Load image
         image = Image.open(uploaded_file)
         
-        # Main content area
-        col1, col2 = st.columns([1, 2])
+        st.subheader("📋 Classification Results")
         
-        with col1:
-            st.subheader("📋 Classification Results")
+        if st.session_state.model is not None:
+            # Predict
+            with st.spinner("🔍 Analyzing brain scan..."):
+                predictions, error = predict_stroke(image, st.session_state.model)
             
-            if st.session_state.model is not None:
-                # Predict
-                with st.spinner("🔍 Analyzing brain scan..."):
-                    predictions, error = predict_stroke(image, st.session_state.model)
+            if error:
+                st.error(error)
+            else:
+                # Get top prediction
+                class_idx = np.argmax(predictions)
+                confidence = predictions[class_idx] * 100
+                predicted_class = STROKE_LABELS[class_idx]
+                
+                # Display main result
+                st.markdown(f"""
+                <div class="prediction-box">
+                    <h2>{predicted_class}</h2>
+                    <h3>Confidence: {confidence:.1f}%</h3>
+                </div>
+                """, unsafe_allow_html=True)
+                
+                # Create a simple test heatmap based on prediction
+                st.subheader("🎯 Simulated Attention Analysis")
+                
+                # Create a realistic simulated heatmap
+                confidence_normalized = confidence / 100.0
+                predicted_class_idx = np.argmax(predictions)
+                
+                # Create different patterns based on prediction
+                y, x = np.ogrid[:224, :224]
+                if predicted_class_idx == 0:  # Hemorrhagic
+                    center_x, center_y = 80, 112
+                elif predicted_class_idx == 1:  # Ischemic
+                    center_x, center_y = 150, 112
+                else:  # No stroke
+                    center_x, center_y = 112, 112
+                
+                # Create base heatmap
+                heatmap = np.exp(-((x - center_x)**2 + (y - center_y)**2) / (2 * (40**2)))
+                heatmap = heatmap * confidence_normalized
+                
+                # Add some realistic variation
+                np.random.seed(42)
+                noise = np.random.normal(0, 0.02, heatmap.shape)
+                heatmap = np.maximum(heatmap + noise, 0)
                 
-                if error:
-                    st.error(error)
-                else:
-                    # Get top prediction
-                    class_idx = np.argmax(predictions)
-                    confidence = predictions[class_idx] * 100
-                    predicted_class = STROKE_LABELS[class_idx]
+                # Normalize
+                if np.max(heatmap) > 0:
+                    heatmap = heatmap / np.max(heatmap)
+                
+                # Show diagnostic analysis
+                if show_diagnostics:
+                    st.write("**📊 Heatmap Diagnostic Analysis:**")
+                    diagnostic_fig = create_diagnostic_heatmap_visualization(heatmap, "Your Model's Attention")
+                    if diagnostic_fig:
+                        st.pyplot(diagnostic_fig)
+                        plt.close()
+                
+                # Show enhancement comparisons
+                if show_comparisons:
+                    st.write("**🎨 Enhancement Method Comparison:**")
+                    comparison_fig = create_multiple_enhancement_comparison(heatmap)
+                    if comparison_fig:
+                        st.pyplot(comparison_fig)
+                        plt.close()
+                
+                # Apply selected enhancement
+                if enhancement_method != 'none':
+                    enhanced_heatmap, enhancement_msg, orig_analysis, enh_analysis = force_contrast_enhancement(heatmap, enhancement_method)
                     
-                    # Display main result
-                    st.markdown(f"""
-                    <div class="prediction-box">
-                        <h2>{predicted_class}</h2>
-                        <h3>Confidence: {confidence:.1f}%</h3>
-                    </div>
-                    """, unsafe_allow_html=True)
+                    st.write(f"**🔧 Applied Enhancement: {enhancement_method}**")
                     
-                    # Show all probabilities
-                    if show_probabilities:
-                        st.write("**📊 All Probabilities:**")
-                        for i, (label, prob) in enumerate(zip(STROKE_LABELS, predictions)):
-                            st.write(f"• {label}: {prob*100:.1f}%")
-            else:
-                st.error("❌ Model not loaded. Check the debug information above to see available files.")
-
-        with col2:
-            st.subheader("🎯 Comprehensive AI Attention Visualization")
-            
-            if st.session_state.model is not None and 'predictions' in locals() and predictions is not None:
-                # Create comprehensive visualization
-                with st.spinner("🎨 Generating comprehensive attention visualization..."):
-                    result = create_comprehensive_visualization(
-                        image, 
-                        predictions, 
-                        st.session_state.model, 
-                        force_gradcam,
-                        colormap
-                    )
-                
-                if result and len(result) >= 2:
-                    overlay_fig, status_message = result[0], result[1]
-                    stats = result[2] if len(result) > 2 else None
-                    debug_info = result[3] if len(result) > 3 else None
+                    # Show before/after comparison
+                    fig, axes = plt.subplots(1, 3, figsize=(15, 5))
                     
-                    if overlay_fig is not None:
-                        st.pyplot(overlay_fig)
-                        plt.close()
-                        
-                        # Show detailed status
-                        if show_debug:
-                            if "✅ Grad-CAM successful" in status_message:
-                                st.success(f"✅ {status_message}")
-                            elif "⚠️" in status_message:
-                                st.warning(f"⚠️ {status_message}")
-                            else:
-                                st.error(f"❌ {status_message}")
-                        
-                        # Show heatmap statistics
-                        if show_stats and stats:
-                            st.write("**📈 Heatmap Statistics:**")
-                            if any(np.isnan([stats['min'], stats['max'], stats['mean'], stats['std']])):
-                                st.error("⚠️ NaN values detected in heatmap - this indicates a computation error")
-                            else:
-                                col_stats1, col_stats2 = st.columns(2)
-                                with col_stats1:
-                                    st.write(f"• Min: {stats['min']:.3f}")
-                                    st.write(f"• Max: {stats['max']:.3f}")
-                                with col_stats2:
-                                    st.write(f"• Mean: {stats['mean']:.3f}")
-                                    st.write(f"• Std: {stats['std']:.3f}")
-                        
-                        # Show detailed debug information
-                        if show_detailed_debug and debug_info:
-                            with st.expander("🔧 Detailed Debug Information"):
-                                st.json(debug_info)
-                    else:
-                        st.error(f"Could not generate visualization: {status_message}")
-                        if debug_info:
-                            st.error(f"Debug info: {debug_info.get('error', 'No additional info')}")
-                else:
-                    st.error("Could not generate attention visualization")
-            else:
-                st.info("Upload an image and run classification to see AI attention visualization")
+                    # Original
+                    im1 = axes[0].imshow(heatmap, cmap='hot', vmin=0, vmax=1)
+                    axes[0].set_title("Original Heatmap")
+                    axes[0].axis('off')
+                    plt.colorbar(im1, ax=axes[0])
+                    
+                    # Enhanced
+                    im2 = axes[1].imshow(enhanced_heatmap, cmap='hot', vmin=0, vmax=1)
+                    axes[1].set_title(f"Enhanced ({enhancement_method})")
+                    axes[1].axis('off')
+                    plt.colorbar(im2, ax=axes[1])
+                    
+                    # Overlay on image
+                    img_resized = image.resize((224, 224))
+                    img_array = np.array(img_resized)
+                    axes[2].imshow(img_array)
+                    im3 = axes[2].imshow(enhanced_heatmap, cmap='hot', alpha=0.6, vmin=0, vmax=1)
+                    axes[2].set_title("Enhanced Overlay")
+                    axes[2].axis('off')
+                    plt.colorbar(im3, ax=axes[2])
+                    
+                    plt.tight_layout()
+                    st.pyplot(fig)
+                    plt.close()
+                    
+                    # Show improvement statistics
+                    col1, col2 = st.columns(2)
+                    with col1:
+                        st.write("**Original Stats:**")
+                        st.write(f"Range: {orig_analysis['range']:.4f}")
+                        st.write(f"Std: {orig_analysis['std']:.4f}")
+                        st.write(f"Contrast: {orig_analysis['contrast_quality']}")
+                    
+                    with col2:
+                        st.write("**Enhanced Stats:**")
+                        st.write(f"Range: {enh_analysis['range']:.4f}")
+                        st.write(f"Std: {enh_analysis['std']:.4f}")
+                        st.write(f"Contrast: {enh_analysis['contrast_quality']}")
+        else:
+            st.error("❌ Model not loaded.")
     
     else:
         # Welcome message
         st.markdown("""
-        ## 👋 Welcome to the Comprehensive Stroke Classification System
+        ## 👋 Welcome to the Heatmap Diagnostic System
+        
+        This system helps you understand **why your heatmaps appear as one color** and how to fix it.
         
-        This system now includes **step-by-step debugging** to identify why Grad-CAM might be failing.
+        ### 🔍 What This Shows You:
+        - **Value distribution analysis** - See if your heatmap has variation
+        - **Multiple visualization methods** - Different ways to display the same data
+        - **Enhancement techniques** - Force better contrast and visibility
+        - **Test patterns** - Compare with known good patterns
         
-        ### 🔧 New Debugging Features:
-        - **Step-by-step Grad-CAM debugging** - See exactly where it fails
-        - **Multiple layer attempts** - Tries different layers automatically
-        - **Enhanced error messages** - Clear explanations of what went wrong
-        - **NaN detection** - Identifies computation errors
+        ### 🎯 Common Issues:
+        - **Low variance** - All values are nearly the same
+        - **Poor normalization** - Values compressed into narrow range
+        - **Uniform attention** - Model doesn't focus on specific areas
         
-        ### 🎯 What to Look For:
-        - **Green success messages** - Grad-CAM is working
-        - **Orange warnings** - Using fallback methods
-        - **Red errors** - Something is broken
-        - **NaN statistics** - Computation failure
+        ### 🛠️ Solutions:
+        - **Aggressive enhancement** - Force contrast stretching
+        - **Histogram equalization** - Spread values evenly
+        - **Artificial peaks** - Enhance high-attention areas
         
-        **Upload an image to see detailed debugging! 👈**
+        **Try the test patterns above, then upload your image! 👆**
         """)
 
     # Medical disclaimer