Update src/streamlit_app.py

src/streamlit_app.py

@@ -129,40 +129,44 @@ def analyze_model_architecture(model):
         'conv_layers': [],
         'dense_layers': [],
         'other_layers': [],
-        'potential_gradcam_layers': [],
+        'all_layers_detailed': [],
         'model_type': 'Unknown'
     }
     
     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')
+            'output_shape': getattr(layer, 'output_shape', 'Unknown'),
+            'trainable': getattr(layer, 'trainable', 'Unknown'),
+            'activation': getattr(layer, 'activation', None)
         }
         
-        # Categorize layers
+        # 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)
-            layer_analysis['potential_gradcam_layers'].append(layer_info)
         
         elif 'Dense' in layer_type or 'Linear' in layer_type:
             layer_analysis['dense_layers'].append(layer_info)
         
-        # Check for other layer types that might work with Grad-CAM
-        elif any(layer_name in layer_type for layer_name in [
-            'Activation', 'BatchNormalization', 'Dropout', 'MaxPooling', 'AveragePooling',
-            'GlobalMaxPooling', 'GlobalAveragePooling', 'Flatten', 'Reshape'
-        ]):
+        else:
             layer_analysis['other_layers'].append(layer_info)
-            
-            # Some of these might be suitable for Grad-CAM if they have spatial dimensions
-            if any(pool_type in layer_type for pool_type in ['MaxPooling2D', 'AveragePooling2D']):
-                layer_analysis['potential_gradcam_layers'].append(layer_info)
     
     # Determine model type
     if layer_analysis['conv_layers']:
@@ -174,124 +178,112 @@ def analyze_model_architecture(model):
     
     return layer_analysis
 
-def find_best_layer_for_gradcam(model):
-    """Find the best layer for Grad-CAM with expanded search."""
-    if model is None:
-        return None, "No model loaded"
-    
-    analysis = analyze_model_architecture(model)
-    
-    # Priority 1: Convolutional layers (best for Grad-CAM)
-    if analysis['conv_layers']:
-        best_layer = analysis['conv_layers'][-1]  # Last conv layer
-        return best_layer['name'], f"✅ Using convolutional layer: {best_layer['name']} ({best_layer['type']})"
-    
-    # Priority 2: Pooling layers with spatial dimensions
-    pooling_layers = [layer for layer in analysis['other_layers'] 
-                     if any(pool_type in layer['type'] for pool_type in ['MaxPooling2D', 'AveragePooling2D'])]
-    if pooling_layers:
-        best_layer = pooling_layers[-1]
-        return best_layer['name'], f"⚠️ Using pooling layer: {best_layer['name']} ({best_layer['type']}) - may not work well"
-    
-    # Priority 3: Try any layer with 4D output (batch, height, width, channels)
-    for layer in reversed(model.layers):
-        try:
-            output_shape = layer.output_shape
-            if isinstance(output_shape, tuple) and len(output_shape) == 4:
-                return layer.name, f"⚠️ Trying 4D layer: {layer.name} ({type(layer).__name__}) - experimental"
-        except:
-            continue
-    
-    # Priority 4: Try the layer before the last dense layer
-    if analysis['dense_layers'] and len(model.layers) > 2:
-        # Find the layer just before the first dense layer
-        for i, layer in enumerate(model.layers):
-            if 'Dense' in type(layer).__name__:
-                if i > 0:
-                    prev_layer = model.layers[i-1]
-                    return prev_layer.name, f"⚠️ Using pre-dense layer: {prev_layer.name} ({type(prev_layer).__name__}) - may not work"
-                break
+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
+    }
     
-    return None, "❌ No suitable layers found for Grad-CAM"
-
-def make_gradcam_heatmap(img_array, model, layer_name, pred_index=None):
-    """Generate Grad-CAM heatmap with better error handling."""
     try:
-        # Get the target layer
+        debug_info['step'] = 'Getting target layer'
         target_layer = model.get_layer(layer_name)
+        debug_info['target_layer_type'] = type(target_layer).__name__
         
-        # Create a model that maps the input image to the activations of the target layer
+        debug_info['step'] = 'Creating grad model'
         grad_model = tf.keras.Model(
             inputs=[model.inputs],
             outputs=[target_layer.output, model.output]
         )
         
-        # Compute the gradient of the top predicted class
+        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()
         
-        # Get gradients
+        debug_info['step'] = 'Computing gradients'
         grads = tape.gradient(class_channel, layer_output)
         
         if grads is None:
-            return None, "Gradients are None - layer may not be differentiable"
+            debug_info['error'] = "Gradients are None - no backpropagation path"
+            return None, debug_info
+        
+        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))
+        }
         
-        # Handle different layer output shapes
-        if len(layer_output.shape) == 4:  # Standard conv layer (batch, height, width, channels)
-            # Standard Grad-CAM for conv layers
+        debug_info['step'] = 'Processing gradients based on layer type'
+        
+        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 (batch, features)
-            # For dense layers, create a simple attention based on gradients
-            grads_abs = tf.abs(grads[0])
-            heatmap = tf.reduce_mean(grads_abs)
-            # Create a uniform heatmap since dense layers don't have spatial structure
-            heatmap = tf.ones((7, 7)) * heatmap  # Small heatmap to be resized later
-        
+            
+        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:
-            return None, f"Unsupported layer output shape: {layer_output.shape}"
+            debug_info['error'] = f"Unsupported layer shape: {layer_output.shape}"
+            return None, debug_info
+        
+        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))
+        }
         
-        # Normalize the heatmap
+        # Apply ReLU (remove negative values)
         heatmap = tf.maximum(heatmap, 0)
-        if tf.reduce_max(heatmap) > 0:
-            heatmap = heatmap / tf.reduce_max(heatmap)
         
-        return heatmap.numpy(), None
+        # 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
+        
+        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:
-        return None, f"Grad-CAM computation error: {str(e)}"
-
-def enhance_heatmap_contrast(heatmap):
-    """Enhance heatmap contrast and dynamic range."""
-    if heatmap is None:
-        return None
-    
-    # Apply histogram equalization-like enhancement
-    heatmap_flat = heatmap.flatten()
-    
-    # Remove zeros for better contrast
-    non_zero_values = heatmap_flat[heatmap_flat > 0]
-    if len(non_zero_values) == 0:
-        return heatmap
-    
-    # Enhance contrast using percentile stretching
-    p2, p98 = np.percentile(non_zero_values, [2, 98])
-    heatmap_enhanced = np.clip((heatmap - p2) / (p98 - p2), 0, 1)
-    
-    # Apply power law transformation for better visibility
-    gamma = 0.5  # Makes mid-tones brighter
-    heatmap_enhanced = np.power(heatmap_enhanced, gamma)
-    
-    return heatmap_enhanced
+        debug_info['error'] = f"Exception in step '{debug_info['step']}': {str(e)}"
+        return None, debug_info
 
-def create_real_gradcam_heatmap(img, model, predictions):
-    """Create a real Grad-CAM heatmap with enhanced visualization."""
+def create_robust_gradcam_heatmap(img, model, predictions):
+    """Create Grad-CAM with comprehensive debugging."""
     try:
         # Preprocess image
         img_resized = img.resize((224, 224))
@@ -304,50 +296,67 @@ def create_real_gradcam_heatmap(img, model, predictions):
         # Normalize and add batch dimension
         img_array = np.expand_dims(img_array, axis=0) / 255.0
         
-        # Find the best layer for Grad-CAM
-        layer_name, layer_status = find_best_layer_for_gradcam(model)
+        # Get model analysis
+        analysis = analyze_model_architecture(model)
         
-        if layer_name is None:
-            return None, f"❌ {layer_status}", None
+        # Try different layers in order of preference
+        layer_candidates = []
         
-        # Generate Grad-CAM heatmap
-        heatmap, error = make_gradcam_heatmap(
-            img_array, 
-            model, 
-            layer_name,
-            pred_index=np.argmax(predictions)
-        )
+        # Add conv layers first
+        for layer in analysis['conv_layers']:
+            layer_candidates.append((layer['name'], f"Conv layer: {layer['name']}"))
         
-        if error:
-            return None, f"❌ {error} (Layer: {layer_name})", None
+        # 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']})"))
         
-        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
-            
-            # Enhance contrast
-            heatmap_enhanced = enhance_heatmap_contrast(heatmap_resized)
-            
-            # Get statistics for debugging
-            stats = {
-                'min': float(np.min(heatmap_enhanced)),
-                'max': float(np.max(heatmap_enhanced)),
-                'mean': float(np.mean(heatmap_enhanced)),
-                'std': float(np.std(heatmap_enhanced))
-            }
+        # 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)
             
-            return heatmap_enhanced, f"✅ Grad-CAM successful using layer: {layer_name}", stats
-        else:
-            return None, "❌ Failed to generate heatmap", None
+            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
+        return None, f"❌ Grad-CAM error: {str(e)}", None, {'error': str(e)}
 
 def predict_stroke(img, model):
     """Predict stroke type from image."""
@@ -374,38 +383,58 @@ def predict_stroke(img, model):
     except Exception as e:
         return None, f"Prediction error: {str(e)}"
 
-def create_simulated_heatmap(img, predictions):
-    """Create a simulated heatmap (fallback)."""
+def create_enhanced_simulated_heatmap(img, predictions):
+    """Create a more realistic simulated heatmap."""
     try:
         confidence = np.max(predictions)
-        np.random.seed(42)
-        heatmap = np.random.rand(224, 224) * confidence
-        
-        try:
-            from scipy import ndimage
-            heatmap = ndimage.gaussian_filter(heatmap, sigma=20)
-        except ImportError:
+        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
-            y, x = np.ogrid[:224, :224]
-            mask = (x - center_x)**2 + (y - center_y)**2
-            heatmap = np.exp(-mask / (2 * (50**2))) * confidence
         
-        # Enhance simulated heatmap too
-        heatmap_enhanced = enhance_heatmap_contrast(heatmap)
+        # Create base pattern
+        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_enhanced)),
-            'max': float(np.max(heatmap_enhanced)),
-            'mean': float(np.mean(heatmap_enhanced)),
-            'std': float(np.std(heatmap_enhanced))
+            'min': float(np.min(heatmap)),
+            'max': float(np.max(heatmap)),
+            'mean': float(np.mean(heatmap)),
+            'std': float(np.std(heatmap))
         }
         
-        return heatmap_enhanced, "⚠️ Using simulated heatmap (Grad-CAM failed)", stats
+        return heatmap, "⚠️ Using enhanced simulated heatmap", stats
     except Exception as e:
         return None, f"❌ Simulated heatmap error: {str(e)}", None
 
-def create_enhanced_overlay_visualization(img, predictions, model, force_gradcam=True, colormap='hot'):
-    """Create enhanced overlay visualization with better colors."""
+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"
     
@@ -417,17 +446,20 @@ def create_enhanced_overlay_visualization(img, predictions, model, force_gradcam
         heatmap = None
         status_message = ""
         stats = None
+        debug_info = None
         
         # Try Grad-CAM first
         if force_gradcam and model is not None:
-            result = create_real_gradcam_heatmap(img, model, predictions)
-            if result and len(result) == 3:
-                heatmap, gradcam_status, stats = result
+            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 simulated if Grad-CAM failed
+        # Fallback to enhanced simulated if Grad-CAM failed
         if heatmap is None:
-            result = create_simulated_heatmap(img, predictions)
+            result = create_enhanced_simulated_heatmap(img, predictions)
             if result and len(result) == 3:
                 heatmap, sim_status, stats = result
                 if status_message:
@@ -436,9 +468,9 @@ def create_enhanced_overlay_visualization(img, predictions, model, force_gradcam
                     status_message = sim_status
         
         if heatmap is None:
-            return None, "❌ Could not generate any heatmap"
+            return None, "❌ Could not generate any heatmap", None, None
         
-        # Create enhanced visualization with multiple views
+        # Create visualization
         fig, axes = plt.subplots(1, 3, figsize=(15, 5))
         
         # 1. Original image
@@ -446,13 +478,13 @@ def create_enhanced_overlay_visualization(img, predictions, model, force_gradcam
         axes[0].set_title("Original Image", fontsize=12, fontweight='bold')
         axes[0].axis('off')
         
-        # 2. Heatmap only with enhanced colors
+        # 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 with better blending
+        # 3. Overlay
         axes[2].imshow(img_array)
         im2 = axes[2].imshow(heatmap, cmap=colormap, alpha=0.6, vmin=0, vmax=1, interpolation='bilinear')
         
@@ -470,10 +502,10 @@ def create_enhanced_overlay_visualization(img, predictions, model, force_gradcam
         
         plt.tight_layout()
         
-        return fig, status_message, stats
+        return fig, status_message, stats, debug_info
     
     except Exception as e:
-        return None, f"❌ Visualization error: {str(e)}", None
+        return None, f"❌ Visualization error: {str(e)}", None, None
 
 # Main App
 def main():
@@ -514,7 +546,7 @@ def main():
 
     # Enhanced model architecture analysis
     if st.session_state.model is not None:
-        with st.expander("🔍 Enhanced Model Architecture Analysis"):
+        with st.expander("🔍 Detailed Model Architecture Analysis"):
             analysis = analyze_model_architecture(st.session_state.model)
             
             st.write("**📊 Model Summary:**")
@@ -524,25 +556,11 @@ def main():
             st.write(f"- **Dense Layers:** {len(analysis['dense_layers'])}")
             st.write(f"- **Other Layers:** {len(analysis['other_layers'])}")
             
-            # Show layer details
-            if analysis['conv_layers']:
-                st.write("**🎯 Convolutional Layers (Best for Grad-CAM):**")
-                for layer in analysis['conv_layers']:
-                    st.write(f"  ✅ {layer['name']} ({layer['type']}) - Shape: {layer['output_shape']}")
-            
-            if analysis['dense_layers']:
-                st.write("**🧠 Dense Layers:**")
-                for layer in analysis['dense_layers']:
-                    st.write(f"  • {layer['name']} ({layer['type']}) - Shape: {layer['output_shape']}")
-            
-            # Test layer selection
-            layer_name, layer_status = find_best_layer_for_gradcam(st.session_state.model)
-            if "✅" in layer_status:
-                st.markdown(f'<div class="status-box success"><strong>Grad-CAM Layer:</strong> {layer_status}</div>', unsafe_allow_html=True)
-            elif "⚠️" in layer_status:
-                st.markdown(f'<div class="status-box warning"><strong>Grad-CAM Layer:</strong> {layer_status}</div>', unsafe_allow_html=True)
-            else:
-                st.markdown(f'<div class="status-box error"><strong>Grad-CAM Layer:</strong> {layer_status}</div>', unsafe_allow_html=True)
+            # 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"):
@@ -564,7 +582,7 @@ def main():
         force_gradcam = st.checkbox(
             "Attempt Grad-CAM", 
             value=True,
-            help="Try Grad-CAM even with non-CNN models (experimental)"
+            help="Try Grad-CAM with comprehensive debugging"
         )
         
         colormap = st.selectbox(
@@ -577,18 +595,7 @@ def main():
         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)
-        
-        st.markdown("---")
-        st.header("🎨 Color Scheme Guide")
-        st.info("""
-        **hot**: Red-Yellow (classic heat)
-        **jet**: Blue-Green-Yellow-Red
-        **viridis**: Purple-Blue-Green-Yellow
-        **plasma**: Purple-Pink-Yellow
-        **inferno**: Black-Purple-Red-Yellow
-        **magma**: Black-Purple-Pink-White
-        **coolwarm**: Blue-White-Red
-        """)
+        show_detailed_debug = st.checkbox("Show Detailed Debug Info", value=False)
 
     if uploaded_file is not None:
         # Load image
@@ -630,12 +637,12 @@ def main():
                 st.error("❌ Model not loaded. Check the debug information above to see available files.")
 
         with col2:
-            st.subheader("🎯 Enhanced AI Attention Visualization")
+            st.subheader("🎯 Comprehensive AI Attention Visualization")
             
             if st.session_state.model is not None and 'predictions' in locals() and predictions is not None:
-                # Create enhanced overlay visualization
-                with st.spinner("🎨 Generating enhanced attention visualization..."):
-                    result = create_enhanced_overlay_visualization(
+                # Create comprehensive visualization
+                with st.spinner("🎨 Generating comprehensive attention visualization..."):
+                    result = create_comprehensive_visualization(
                         image, 
                         predictions, 
                         st.session_state.model, 
@@ -646,6 +653,7 @@ def main():
                 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
                     
                     if overlay_fig is not None:
                         st.pyplot(overlay_fig)
@@ -658,20 +666,30 @@ def main():
                             elif "⚠️" in status_message:
                                 st.warning(f"⚠️ {status_message}")
                             else:
-                                st.info(f"ℹ️ {status_message}")
+                                st.error(f"❌ {status_message}")
                         
                         # Show heatmap statistics
                         if show_stats and stats:
                             st.write("**📈 Heatmap Statistics:**")
-                            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}")
+                            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:
@@ -680,22 +698,23 @@ def main():
     else:
         # Welcome message
         st.markdown("""
-        ## 👋 Welcome to the Enhanced Stroke Classification System
+        ## 👋 Welcome to the Comprehensive Stroke Classification System
         
-        This system now provides **better color visualization** and **enhanced contrast** for attention heatmaps.
+        This system now includes **step-by-step debugging** to identify why Grad-CAM might be failing.
         
-        ### 🎨 New Visualization Features:
-        - **Multiple Color Schemes**: Choose from 7 different color palettes
-        - **Enhanced Contrast**: Better visibility of attention patterns
-        - **Three-Panel View**: Original, heatmap, and overlay side-by-side
-        - **Statistics Display**: See heatmap value distributions
+        ### 🔧 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
         
-        ### 🚀 Why Colors Matter:
-        - **Red/Yellow (hot)**: High attention areas
-        - **Blue/Purple**: Low attention areas  
-        - **Enhanced contrast**: Makes subtle patterns visible
+        ### 🎯 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
         
-        **Get started by uploading an image! 👈**
+        **Upload an image to see detailed debugging! 👈**
         """)
 
     # Medical disclaimer