Update src/streamlit_app.py

src/streamlit_app.py

@@ -119,83 +119,156 @@ 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."""
+def analyze_model_architecture(model):
+    """Comprehensive analysis of model architecture."""
     if model is None:
-        return "No model loaded"
+        return {"error": "No model loaded"}
     
-    layer_info = []
-    conv_layers = []
+    layer_analysis = {
+        'total_layers': len(model.layers),
+        'conv_layers': [],
+        'dense_layers': [],
+        'other_layers': [],
+        'potential_gradcam_layers': [],
+        'model_type': 'Unknown'
+    }
     
     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})")
+        layer_info = {
+            'index': i,
+            'name': layer.name,
+            'type': layer_type,
+            'output_shape': getattr(layer, 'output_shape', 'Unknown')
+        }
+        
+        # Categorize layers
+        if any(conv_type in layer_type for conv_type in [
+            'Conv1D', 'Conv2D', 'Conv3D', 'SeparableConv2D', 'DepthwiseConv2D',
+            'Convolution1D', 'Convolution2D', 'Convolution3D'
+        ]):
+            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'
+        ]):
+            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)
     
-    return {
-        'all_layers': layer_info,
-        'conv_layers': conv_layers,
-        'total_layers': len(model.layers)
-    }
+    # 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)'
+    else:
+        layer_analysis['model_type'] = 'Custom Architecture'
+    
+    return layer_analysis
 
-def find_best_conv_layer(model):
-    """Find the best convolutional layer for Grad-CAM."""
+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"
     
-    conv_layers = []
+    analysis = analyze_model_architecture(model)
     
-    # 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)
+    # 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"
     
-    if not conv_layers:
-        return None, "No convolutional layers found"
+    # 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
     
-    # Return the last convolutional layer
-    return conv_layers[-1], f"Found {len(conv_layers)} conv layers, using: {conv_layers[-1]}"
+    # 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
+    
+    return None, "❌ No suitable layers found for Grad-CAM"
 
-def make_gradcam_heatmap(img_array, model, last_conv_layer_name, pred_index=None):
-    """Generate Grad-CAM heatmap."""
+def make_gradcam_heatmap(img_array, model, layer_name, pred_index=None):
+    """Generate Grad-CAM heatmap with better error handling."""
     try:
-        # Create a model that maps the input image to the activations of the last conv layer
+        # Get the target layer
+        target_layer = model.get_layer(layer_name)
+        
+        # Create a model that maps the input image to the activations of the target layer
         grad_model = tf.keras.Model(
             inputs=[model.inputs],
-            outputs=[model.get_layer(last_conv_layer_name).output, model.output]
+            outputs=[target_layer.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)
+            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)
+        # Get gradients
+        grads = tape.gradient(class_channel, layer_output)
         
-        # Mean intensity of the gradient over a specific feature map channel
-        pooled_grads = tf.reduce_mean(grads, axis=(0, 1, 2))
+        if grads is None:
+            return None, "Gradients are None - layer may not be differentiable"
         
-        # 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)
+        # Handle different layer output shapes
+        if len(layer_output.shape) == 4:  # Standard conv layer (batch, height, width, channels)
+            # Standard Grad-CAM for conv layers
+            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)
         
-        # Normalize the heatmap between 0 & 1
-        heatmap = tf.maximum(heatmap, 0) / tf.math.reduce_max(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
+        
+        else:
+            return None, f"Unsupported layer output shape: {layer_output.shape}"
+        
+        # Normalize the heatmap
+        heatmap = tf.maximum(heatmap, 0)
+        if tf.reduce_max(heatmap) > 0:
+            heatmap = heatmap / tf.reduce_max(heatmap)
         
         return heatmap.numpy(), None
         
     except Exception as e:
-        return None, f"Grad-CAM error: {str(e)}"
+        return None, f"Grad-CAM computation error: {str(e)}"
 
 def create_real_gradcam_heatmap(img, model, predictions):
-    """Create a real Grad-CAM heatmap."""
+    """Create a real Grad-CAM heatmap with fallback strategies."""
     try:
         # Preprocess image
         img_resized = img.resize((224, 224))
@@ -208,31 +281,34 @@ 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 convolutional layer
-        conv_layer_name, layer_status = find_best_conv_layer(model)
+        # Find the best layer for Grad-CAM
+        layer_name, layer_status = find_best_layer_for_gradcam(model)
         
-        if conv_layer_name is None:
+        if layer_name is None:
             return None, f"❌ {layer_status}"
         
         # Generate Grad-CAM heatmap
         heatmap, error = make_gradcam_heatmap(
             img_array, 
             model, 
-            conv_layer_name,
+            layer_name,
             pred_index=np.argmax(predictions)
         )
         
         if error:
-            return None, f"❌ {error}"
+            return None, f"❌ {error} (Layer: {layer_name})"
         
         if heatmap is not None:
             # Resize heatmap to match input image size
-            heatmap_resized = tf.image.resize(
-                heatmap[..., tf.newaxis], 
-                (224, 224)
-            ).numpy()[:, :, 0]
+            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
             
-            return heatmap_resized, f"✅ Grad-CAM successful using layer: {conv_layer_name}"
+            return heatmap_resized, f"✅ Grad-CAM successful using layer: {layer_name}"
         else:
             return None, "❌ Failed to generate heatmap"
             
@@ -380,29 +456,50 @@ 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)
 
-    # Debug model architecture
+    # Enhanced model architecture analysis
     if st.session_state.model is not None:
-        with st.expander("🔍 Model Architecture Debug"):
-            debug_info = debug_model_layers(st.session_state.model)
+        with st.expander("🔍 Enhanced Model Architecture Analysis"):
+            analysis = analyze_model_architecture(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'])}")
+            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'])}")
             
-            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)
+            # 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']}")
+            
+            if analysis['other_layers']:
+                st.write("**⚙️ Other Layers:**")
+                for layer in analysis['other_layers'][:5]:  # Show first 5
+                    st.write(f"  • {layer['name']} ({layer['type']}) - Shape: {layer['output_shape']}")
+                if len(analysis['other_layers']) > 5:
+                    st.write(f"  ... and {len(analysis['other_layers']) - 5} more")
+            
+            # 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('<div class="status-box error">❌ No convolutional layers found - Grad-CAM will not work</div>', unsafe_allow_html=True)
+                st.markdown(f'<div class="status-box error"><strong>Grad-CAM Layer:</strong> {layer_status}</div>', unsafe_allow_html=True)
             
-            with st.expander("All Layers (Advanced)"):
-                for layer_info in debug_info['all_layers']:
-                    st.code(layer_info)
+            # Show model architecture recommendation
+            if analysis['model_type'] == 'MLP (Multi-Layer Perceptron)':
+                st.warning("⚠️ **Note:** Your model appears to be a Multi-Layer Perceptron (MLP) without convolutional layers. Grad-CAM works best with CNNs. The visualization may be limited or use experimental methods.")
+            elif analysis['model_type'] == 'Custom Architecture':
+                st.info("ℹ️ **Note:** Your model has a custom architecture. Grad-CAM compatibility depends on the specific layers used.")
 
     # Manual reload button
     if st.button("🔄 Reload Model", help="Try to reload the model"):
@@ -422,9 +519,9 @@ def main():
         st.header("🎨 Visualization Options")
         
         force_gradcam = st.checkbox(
-            "Force Grad-CAM Attempt", 
+            "Attempt Grad-CAM", 
             value=True,
-            help="Always try Grad-CAM first (recommended)"
+            help="Try Grad-CAM even with non-CNN models (experimental)"
         )
         
         show_probabilities = st.checkbox("Show All Probabilities", value=True)
@@ -433,7 +530,7 @@ def main():
         st.markdown("---")
         st.header("ℹ️ About")
         st.info("""
-        **Model Architecture:** Deep Learning CNN
+        **Model Architecture:** Detected automatically
         
         **Classes:**
         - Hemorrhagic Stroke
@@ -442,7 +539,7 @@ def main():
         
         **Input:** 224×224 RGB images
         
-        **Attention Method:** Grad-CAM
+        **Attention Method:** Grad-CAM (when possible)
         """)
 
     if uploaded_file is not None:
@@ -507,8 +604,10 @@ def main():
                         if show_debug:
                             if "✅ Grad-CAM successful" in status_message:
                                 st.success(f"✅ {status_message}")
-                            else:
+                            elif "⚠️" in status_message:
                                 st.warning(f"⚠️ {status_message}")
+                            else:
+                                st.info(f"ℹ️ {status_message}")
                     else:
                         st.error(f"Could not generate visualization: {status_message}")
                 else:
@@ -521,19 +620,19 @@ def main():
         st.markdown("""
         ## 👋 Welcome to the Stroke Classification System
         
-        This AI system analyzes brain scan images and shows you **exactly where the AI is looking**.
+        This AI system analyzes brain scan images and attempts to show you where the AI focuses its attention.
         
         ### 🚀 Features:
-        - **Deep Learning Classification**: Advanced CNN architecture
-        - **Real AI Attention Maps**: See actual model reasoning with Grad-CAM
-        - **Debug Information**: Understand why Grad-CAM works or fails
-        - **Transparent AI**: Full visibility into the decision process
+        - **Automatic Architecture Detection**: Identifies your model type
+        - **Smart Layer Selection**: Finds the best layer for attention visualization
+        - **Fallback Strategies**: Works with different model architectures
+        - **Transparent Process**: Shows exactly what's happening
         
         ### 📋 How to Use:
-        1. **Check system status** and **model debug info** above
+        1. **Check the Enhanced Model Architecture Analysis** above
         2. **Upload a brain scan image** using the sidebar
         3. **View classification results** with confidence scores
-        4. **Explore attention visualization** - it will tell you if it's real or simulated
+        4. **Explore attention visualization** - real or simulated based on your model
         
         **Get started by uploading an image! 👈**
         """)