Update src/streamlit_app.py

src/streamlit_app.py

@@ -25,14 +25,6 @@ try:
 except ImportError:
     MPL_AVAILABLE = False
 
-# Import our Grad-CAM utilities
-try:
-    from gradcam_utils import create_real_attention_heatmap
-    GRADCAM_AVAILABLE = True
-except ImportError:
-    GRADCAM_AVAILABLE = False
-    st.warning("Grad-CAM utilities not available - using simulated heatmaps")
-
 # Page config
 st.set_page_config(
     page_title="Stroke Classifier",
@@ -65,6 +57,7 @@ st.markdown("""
     .error { background-color: #f8d7da; border: 1px solid #f5c6cb; color: #721c24; }
     .info { background-color: #d1ecf1; border: 1px solid #bee5eb; color: #0c5460; }
     .warning { background-color: #fff3cd; border: 1px solid #ffeaa7; color: #856404; }
+    .debug { background-color: #f8f9fa; border: 1px solid #dee2e6; color: #495057; font-family: monospace; }
 </style>""", unsafe_allow_html=True)
 
 # Initialize session state
@@ -126,6 +119,126 @@ def load_stroke_model():
     except Exception as e:
         return None, f"❌ Model loading failed: {str(e)}"
 
+def debug_model_layers(model):
+    """Debug function to show all model layers."""
+    if model is None:
+        return "No model loaded"
+    
+    layer_info = []
+    conv_layers = []
+    
+    for i, layer in enumerate(model.layers):
+        layer_type = type(layer).__name__
+        layer_info.append(f"{i}: {layer.name} ({layer_type})")
+        
+        if 'conv' in layer.name.lower() or 'Conv' in layer_type:
+            conv_layers.append(f"✅ {layer.name} ({layer_type})")
+    
+    return {
+        'all_layers': layer_info,
+        'conv_layers': conv_layers,
+        'total_layers': len(model.layers)
+    }
+
+def find_best_conv_layer(model):
+    """Find the best convolutional layer for Grad-CAM."""
+    if model is None:
+        return None, "No model loaded"
+    
+    conv_layers = []
+    
+    # Look for convolutional layers
+    for layer in model.layers:
+        layer_type = type(layer).__name__
+        if any(conv_type in layer_type for conv_type in ['Conv2D', 'Conv', 'SeparableConv2D']):
+            conv_layers.append(layer.name)
+    
+    if not conv_layers:
+        return None, "No convolutional layers found"
+    
+    # Return the last convolutional layer
+    return conv_layers[-1], f"Found {len(conv_layers)} conv layers, using: {conv_layers[-1]}"
+
+def make_gradcam_heatmap(img_array, model, last_conv_layer_name, pred_index=None):
+    """Generate Grad-CAM heatmap."""
+    try:
+        # Create a model that maps the input image to the activations of the last conv layer
+        grad_model = tf.keras.Model(
+            inputs=[model.inputs],
+            outputs=[model.get_layer(last_conv_layer_name).output, model.output]
+        )
+        
+        # Compute the gradient of the top predicted class
+        with tf.GradientTape() as tape:
+            last_conv_layer_output, preds = grad_model(img_array)
+            if pred_index is None:
+                pred_index = tf.argmax(preds[0])
+            class_channel = preds[:, pred_index]
+        
+        # Gradient of the output neuron with regard to the output feature map
+        grads = tape.gradient(class_channel, last_conv_layer_output)
+        
+        # Mean intensity of the gradient over a specific feature map channel
+        pooled_grads = tf.reduce_mean(grads, axis=(0, 1, 2))
+        
+        # Multiply each channel by "how important this channel is"
+        last_conv_layer_output = last_conv_layer_output[0]
+        heatmap = last_conv_layer_output @ pooled_grads[..., tf.newaxis]
+        heatmap = tf.squeeze(heatmap)
+        
+        # Normalize the heatmap between 0 & 1
+        heatmap = tf.maximum(heatmap, 0) / tf.math.reduce_max(heatmap)
+        
+        return heatmap.numpy(), None
+        
+    except Exception as e:
+        return None, f"Grad-CAM error: {str(e)}"
+
+def create_real_gradcam_heatmap(img, model, predictions):
+    """Create a real Grad-CAM heatmap."""
+    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
+        
+        # Find the best convolutional layer
+        conv_layer_name, layer_status = find_best_conv_layer(model)
+        
+        if conv_layer_name is None:
+            return None, f"❌ {layer_status}"
+        
+        # Generate Grad-CAM heatmap
+        heatmap, error = make_gradcam_heatmap(
+            img_array, 
+            model, 
+            conv_layer_name,
+            pred_index=np.argmax(predictions)
+        )
+        
+        if error:
+            return None, f"❌ {error}"
+        
+        if heatmap is not None:
+            # Resize heatmap to match input image size
+            heatmap_resized = tf.image.resize(
+                heatmap[..., tf.newaxis], 
+                (224, 224)
+            ).numpy()[:, :, 0]
+            
+            return heatmap_resized, f"✅ Grad-CAM successful using layer: {conv_layer_name}"
+        else:
+            return None, "❌ Failed to generate heatmap"
+            
+    except Exception as e:
+        return None, f"❌ Grad-CAM error: {str(e)}"
+
 def predict_stroke(img, model):
     """Predict stroke type from image."""
     if model is None:
@@ -152,68 +265,72 @@ def predict_stroke(img, model):
         return None, f"Prediction error: {str(e)}"
 
 def create_simulated_heatmap(img, predictions):
-    """Create a simulated heatmap (fallback when Grad-CAM is not available)."""
+    """Create a simulated heatmap (fallback)."""
     try:
-        # Create a simple heatmap based on prediction confidence
         confidence = np.max(predictions)
-        
-        # Generate random attention pattern weighted by confidence
-        np.random.seed(42)  # For reproducible results
+        np.random.seed(42)
         heatmap = np.random.rand(224, 224) * confidence
         
-        # Add some structure to make it look more realistic
         try:
             from scipy import ndimage
             heatmap = ndimage.gaussian_filter(heatmap, sigma=20)
         except ImportError:
-            # Fallback without scipy - create a simple gradient
             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
         
-        return heatmap
+        return heatmap, "⚠️ Using simulated heatmap (Grad-CAM failed)"
     except Exception as e:
-        st.error(f"Simulated heatmap generation error: {e}")
-        return None
+        return None, f"❌ Simulated heatmap error: {str(e)}"
 
-def create_overlay_visualization(img, predictions, model, use_real_gradcam=True):
-    """Create overlay visualization with real or simulated heatmap."""
+def create_overlay_visualization(img, predictions, model, force_gradcam=True):
+    """Create overlay visualization with debugging."""
     if not MPL_AVAILABLE:
-        return None
+        return None, "❌ Matplotlib not available"
     
     try:
-        # Resize image to 224x224 to match heatmap
+        # Resize image to 224x224
         img_resized = img.resize((224, 224))
         img_array = np.array(img_resized)
         
-        # Try to get real attention heatmap first
         heatmap = None
-        heatmap_type = "Simulated"
+        status_message = ""
         
-        if use_real_gradcam and GRADCAM_AVAILABLE and model is not None:
-            heatmap = create_real_attention_heatmap(img, model, predictions)
-            if heatmap is not None:
-                heatmap_type = "Real AI Attention (Grad-CAM)"
+        # Try Grad-CAM first
+        if force_gradcam and model is not None:
+            heatmap, gradcam_status = create_real_gradcam_heatmap(img, model, predictions)
+            status_message = gradcam_status
         
-        # Fallback to simulated heatmap
+        # Fallback to simulated if Grad-CAM failed
         if heatmap is None:
-            heatmap = create_simulated_heatmap(img, predictions)
-            heatmap_type = "Simulated Attention"
+            heatmap, sim_status = create_simulated_heatmap(img, predictions)
+            if status_message:
+                status_message += f" | {sim_status}"
+            else:
+                status_message = sim_status
         
         if heatmap is None:
-            return None, "Could not generate heatmap"
+            return None, "❌ Could not generate any heatmap"
         
-        # Create the overlay visualization
+        # Create visualization
         fig, ax = plt.subplots(figsize=(10, 8))
         
         # Show original image
         ax.imshow(img_array)
         
-        # Overlay heatmap with transparency
+        # Overlay heatmap
         im = ax.imshow(heatmap, cmap='jet', alpha=0.4, interpolation='bilinear')
         
-        ax.set_title(f"Brain Scan with {heatmap_type}", fontsize=14, fontweight='bold', pad=20)
+        # Determine title based on success
+        if "✅ Grad-CAM successful" in status_message:
+            title = "🎯 Real AI Attention (Grad-CAM)"
+            title_color = 'green'
+        else:
+            title = "🎨 Simulated Attention (Grad-CAM Failed)"
+            title_color = 'orange'
+        
+        ax.set_title(title, fontsize=14, fontweight='bold', pad=20, color=title_color)
         ax.axis('off')
         
         # Add colorbar
@@ -221,11 +338,10 @@ def create_overlay_visualization(img, predictions, model, use_real_gradcam=True)
         cbar.set_label('Attention Intensity', rotation=270, labelpad=20)
         
         plt.tight_layout()
-        return fig, heatmap_type
+        return fig, status_message
     
     except Exception as e:
-        st.error(f"Overlay generation error: {e}")
-        return None, f"Error: {e}"
+        return None, f"❌ Visualization error: {str(e)}"
 
 # Main App
 def main():
@@ -240,11 +356,12 @@ def main():
 
     # System status
     st.markdown("### 🔧 System Status")
-    col1, col2, col3, col4 = st.columns(4)
+    col1, col2, col3 = st.columns(3)
     
     with col1:
         if TF_AVAILABLE:
             st.markdown('<div class="status-box success">✅ TensorFlow Ready</div>', unsafe_allow_html=True)
+            st.write(f"TF Version: {tf.__version__}")
         else:
             st.markdown('<div class="status-box error">❌ TensorFlow Error</div>', unsafe_allow_html=True)
 
@@ -255,12 +372,6 @@ def main():
             st.markdown('<div class="status-box error">❌ Matplotlib Error</div>', unsafe_allow_html=True)
 
     with col3:
-        if GRADCAM_AVAILABLE:
-            st.markdown('<div class="status-box success">✅ Grad-CAM Ready</div>', unsafe_allow_html=True)
-        else:
-            st.markdown('<div class="status-box warning">⚠️ Grad-CAM Unavailable</div>', unsafe_allow_html=True)
-
-    with col4:
         if "✅" in st.session_state.model_status:
             st.markdown('<div class="status-box success">✅ Model Loaded</div>', unsafe_allow_html=True)
         else:
@@ -269,28 +380,29 @@ def main():
     # 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)
 
-    # Explanation of heatmap types
-    with st.expander("🔍 Understanding AI Attention Heatmaps"):
-        st.markdown("""
-        ### What do the heatmaps show?
-        
-        **🎯 Real AI Attention (Grad-CAM):**
-        - Shows **actual** regions the AI model focuses on for its decision
-        - Uses gradient-weighted class activation mapping
-        - Highlights pixels that most influence the prediction
-        - **This is the AI's actual reasoning process**
-        
-        **🎨 Simulated Attention:**
-        - Shows **fake** attention patterns (used when Grad-CAM unavailable)
-        - Based on random patterns scaled by prediction confidence
-        - **Does NOT represent actual AI reasoning**
-        - Used as a visual placeholder only
-        
-        ### How to interpret:
-        - **Red/Yellow areas**: High attention (important for decision)
-        - **Blue/Purple areas**: Low attention (less important)
-        - **Intensity**: Stronger colors = more important regions
-        """)
+    # Debug model architecture
+    if st.session_state.model is not None:
+        with st.expander("🔍 Model Architecture Debug"):
+            debug_info = debug_model_layers(st.session_state.model)
+            
+            st.write("**📊 Model Summary:**")
+            st.write(f"- Total layers: {debug_info['total_layers']}")
+            st.write(f"- Convolutional layers found: {len(debug_info['conv_layers'])}")
+            
+            if debug_info['conv_layers']:
+                st.write("**🎯 Available Convolutional Layers:**")
+                for conv_layer in debug_info['conv_layers']:
+                    st.write(f"  {conv_layer}")
+                
+                # Test which layer we'll use
+                conv_layer_name, layer_status = find_best_conv_layer(st.session_state.model)
+                st.markdown(f'<div class="status-box info"><strong>Selected for Grad-CAM:</strong> {layer_status}</div>', unsafe_allow_html=True)
+            else:
+                st.markdown('<div class="status-box error">❌ No convolutional layers found - Grad-CAM will not work</div>', unsafe_allow_html=True)
+            
+            with st.expander("All Layers (Advanced)"):
+                for layer_info in debug_info['all_layers']:
+                    st.code(layer_info)
 
     # Manual reload button
     if st.button("🔄 Reload Model", help="Try to reload the model"):
@@ -309,14 +421,14 @@ def main():
         st.markdown("---")
         st.header("🎨 Visualization Options")
         
-        use_real_gradcam = st.checkbox(
-            "Use Real AI Attention (Grad-CAM)", 
-            value=GRADCAM_AVAILABLE,
-            disabled=not GRADCAM_AVAILABLE,
-            help="Show actual AI reasoning vs simulated patterns"
+        force_gradcam = st.checkbox(
+            "Force Grad-CAM Attempt", 
+            value=True,
+            help="Always try Grad-CAM first (recommended)"
         )
         
         show_probabilities = st.checkbox("Show All Probabilities", value=True)
+        show_debug = st.checkbox("Show Debug Info", value=True)
         
         st.markdown("---")
         st.header("ℹ️ About")
@@ -382,24 +494,23 @@ def main():
                         image, 
                         predictions, 
                         st.session_state.model, 
-                        use_real_gradcam
+                        force_gradcam
                     )
                 
                 if result and len(result) == 2:
-                    overlay_fig, heatmap_type = result
+                    overlay_fig, status_message = result
                     if overlay_fig is not None:
                         st.pyplot(overlay_fig)
                         plt.close()
                         
-                        # Show what type of heatmap is being displayed
-                        if "Real AI Attention" in heatmap_type:
-                            st.success(f"✅ Showing: {heatmap_type}")
-                            st.info("This heatmap shows the **actual** regions the AI focuses on for its decision.")
-                        else:
-                            st.warning(f"⚠️ Showing: {heatmap_type}")
-                            st.info("This is a **simulated** heatmap and does NOT represent actual AI reasoning.")
+                        # Show detailed status
+                        if show_debug:
+                            if "✅ Grad-CAM successful" in status_message:
+                                st.success(f"✅ {status_message}")
+                            else:
+                                st.warning(f"⚠️ {status_message}")
                     else:
-                        st.error("Could not generate visualization")
+                        st.error(f"Could not generate visualization: {status_message}")
                 else:
                     st.error("Could not generate attention visualization")
             else:
@@ -410,20 +521,19 @@ def main():
         st.markdown("""
         ## 👋 Welcome to the Stroke Classification System
         
-        This AI system analyzes brain scan images to detect stroke indicators and shows you **exactly where the AI is looking**.
+        This AI system analyzes brain scan images and shows you **exactly where the AI is looking**.
         
         ### 🚀 Features:
         - **Deep Learning Classification**: Advanced CNN architecture
         - **Real AI Attention Maps**: See actual model reasoning with Grad-CAM
-        - **Three Classes**: Hemorrhagic Stroke, Ischemic Stroke, No Stroke
-        - **Real-time Analysis**: Fast processing with confidence scores
-        - **Transparent AI**: Understand how the AI makes decisions
+        - **Debug Information**: Understand why Grad-CAM works or fails
+        - **Transparent AI**: Full visibility into the decision process
         
         ### 📋 How to Use:
-        1. **Check system status** above (Grad-CAM should show ✅ for real attention)
+        1. **Check system status** and **model debug info** above
         2. **Upload a brain scan image** using the sidebar
         3. **View classification results** with confidence scores
-        4. **Explore REAL attention visualization** to see where the AI actually looks
+        4. **Explore attention visualization** - it will tell you if it's real or simulated
         
         **Get started by uploading an image! 👈**
         """)
@@ -434,3 +544,5 @@ def main():
 
 if __name__ == "__main__":
     main()
+
+</merged_code>