Update src/streamlit_app.py

src/streamlit_app.py

@@ -25,6 +25,14 @@ 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",
@@ -56,6 +64,7 @@ st.markdown("""
     .success { background-color: #d4edda; border: 1px solid #c3e6cb; color: #155724; }
     .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; }
 </style>""", unsafe_allow_html=True)
 
 # Initialize session state
@@ -142,16 +151,9 @@ def predict_stroke(img, model):
     except Exception as e:
         return None, f"Prediction error: {str(e)}"
 
-def create_overlay_heatmap(img, predictions):
-    """Create an overlay heatmap on the original image."""
-    if not MPL_AVAILABLE:
-        return None
-    
+def create_simulated_heatmap(img, predictions):
+    """Create a simulated heatmap (fallback when Grad-CAM is not available)."""
     try:
-        # Resize image to 224x224 to match heatmap
-        img_resized = img.resize((224, 224))
-        img_array = np.array(img_resized)
-        
         # Create a simple heatmap based on prediction confidence
         confidence = np.max(predictions)
         
@@ -170,36 +172,13 @@ def create_overlay_heatmap(img, predictions):
             mask = (x - center_x)**2 + (y - center_y)**2
             heatmap = np.exp(-mask / (2 * (50**2))) * confidence
         
-        # Create the overlay visualization
-        fig, (ax1, ax2, ax3) = plt.subplots(1, 3, figsize=(15, 5))
-        
-        # Original image
-        ax1.imshow(img_array)
-        ax1.set_title("Original Image", fontsize=12, fontweight='bold')
-        ax1.axis('off')
-        
-        # Heatmap only
-        im2 = ax2.imshow(heatmap, cmap='jet', alpha=0.8)
-        ax2.set_title("Attention Heatmap", fontsize=12, fontweight='bold')
-        ax2.axis('off')
-        plt.colorbar(im2, ax=ax2, fraction=0.046, pad=0.04)
-        
-        # Overlay - Original image with heatmap overlay
-        ax3.imshow(img_array)
-        im3 = ax3.imshow(heatmap, cmap='jet', alpha=0.4, interpolation='bilinear')
-        ax3.set_title("Overlay Visualization", fontsize=12, fontweight='bold')
-        ax3.axis('off')
-        plt.colorbar(im3, ax=ax3, fraction=0.046, pad=0.04)
-        
-        plt.tight_layout()
-        return fig
-    
+        return heatmap
     except Exception as e:
-        st.error(f"Heatmap generation error: {e}")
+        st.error(f"Simulated heatmap generation error: {e}")
         return None
 
-def create_single_overlay(img, predictions):
-    """Create a single overlay image combining original and heatmap."""
+def create_overlay_visualization(img, predictions, model, use_real_gradcam=True):
+    """Create overlay visualization with real or simulated heatmap."""
     if not MPL_AVAILABLE:
         return None
     
@@ -208,26 +187,25 @@ def create_single_overlay(img, predictions):
         img_resized = img.resize((224, 224))
         img_array = np.array(img_resized)
         
-        # Create a simple heatmap based on prediction confidence
-        confidence = np.max(predictions)
+        # Try to get real attention heatmap first
+        heatmap = None
+        heatmap_type = "Simulated"
         
-        # Generate random attention pattern weighted by confidence
-        np.random.seed(42)  # For reproducible results
-        heatmap = np.random.rand(224, 224) * confidence
+        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)"
         
-        # 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
+        # Fallback to simulated heatmap
+        if heatmap is None:
+            heatmap = create_simulated_heatmap(img, predictions)
+            heatmap_type = "Simulated Attention"
         
-        # Create the single overlay visualization
-        fig, ax = plt.subplots(figsize=(8, 8))
+        if heatmap is None:
+            return None, "Could not generate heatmap"
+        
+        # Create the overlay visualization
+        fig, ax = plt.subplots(figsize=(10, 8))
         
         # Show original image
         ax.imshow(img_array)
@@ -235,7 +213,7 @@ def create_single_overlay(img, predictions):
         # Overlay heatmap with transparency
         im = ax.imshow(heatmap, cmap='jet', alpha=0.4, interpolation='bilinear')
         
-        ax.set_title("Brain Scan with AI Attention Overlay", fontsize=14, fontweight='bold', pad=20)
+        ax.set_title(f"Brain Scan with {heatmap_type}", fontsize=14, fontweight='bold', pad=20)
         ax.axis('off')
         
         # Add colorbar
@@ -243,34 +221,17 @@ def create_single_overlay(img, predictions):
         cbar.set_label('Attention Intensity', rotation=270, labelpad=20)
         
         plt.tight_layout()
-        return fig
+        return fig, heatmap_type
     
     except Exception as e:
         st.error(f"Overlay generation error: {e}")
-        return None
+        return None, f"Error: {e}"
 
 # Main App
 def main():
     # Header
     st.markdown('<h1 class="main-header">🧠 AI-Powered Stroke Classification System</h1>', unsafe_allow_html=True)
     
-    # Debug info
-    with st.expander("🔍 Debug Information"):
-        st.write(f"**Python Version:** {sys.version}")
-        st.write(f"**Current Directory:** {os.getcwd()}")
-        st.write(f"**Available Files:**")
-        
-        all_files = []
-        for root, dirs, files in os.walk('.'):
-            for file in files:
-                all_files.append(os.path.join(root, file))
-        
-        for file in all_files[:20]:  # Show first 20 files
-            st.write(f"  - {file}")
-        
-        if len(all_files) > 20:
-            st.write(f"  ... and {len(all_files) - 20} more files")
-
     # Auto-load model on startup
     if not st.session_state.model_loaded:
         with st.spinner("Loading AI model..."):
@@ -279,12 +240,11 @@ def main():
 
     # System status
     st.markdown("### 🔧 System Status")
-    col1, col2, col3 = st.columns(3)
+    col1, col2, col3, col4 = st.columns(4)
     
     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)
 
@@ -295,6 +255,12 @@ 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:
@@ -303,6 +269,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
+        """)
+
     # Manual reload button
     if st.button("🔄 Reload Model", help="Try to reload the model"):
         st.session_state.model_loaded = False
@@ -319,11 +308,12 @@ def main():
         
         st.markdown("---")
         st.header("🎨 Visualization Options")
-        viz_option = st.radio(
-            "Choose visualization style:",
-            ["Single Overlay", "Side-by-Side Comparison", "Heatmap Only"],
-            index=0,
-            help="Select how you want to view the AI attention"
+        
+        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"
         )
         
         show_probabilities = st.checkbox("Show All Probabilities", value=True)
@@ -339,6 +329,8 @@ def main():
         - No Stroke
         
         **Input:** 224×224 RGB images
+        
+        **Attention Method:** Grad-CAM
         """)
 
     if uploaded_file is not None:
@@ -384,48 +376,32 @@ def main():
             st.subheader("🎯 AI Attention Visualization")
             
             if st.session_state.model is not None and 'predictions' in locals() and predictions is not None:
-                if viz_option == "Single Overlay":
-                    # Create single overlay
-                    overlay_fig = create_single_overlay(image, predictions)
+                # Create overlay visualization
+                with st.spinner("🎨 Generating attention visualization..."):
+                    result = create_overlay_visualization(
+                        image, 
+                        predictions, 
+                        st.session_state.model, 
+                        use_real_gradcam
+                    )
+                
+                if result and len(result) == 2:
+                    overlay_fig, heatmap_type = result
                     if overlay_fig is not None:
                         st.pyplot(overlay_fig)
                         plt.close()
-                    else:
-                        st.error("Could not generate overlay visualization")
-                
-                elif viz_option == "Side-by-Side Comparison":
-                    # Create side-by-side comparison
-                    comparison_fig = create_overlay_heatmap(image, predictions)
-                    if comparison_fig is not None:
-                        st.pyplot(comparison_fig)
-                        plt.close()
-                    else:
-                        st.error("Could not generate comparison visualization")
-                
-                elif viz_option == "Heatmap Only":
-                    # Show just the heatmap
-                    if MPL_AVAILABLE:
-                        # Generate heatmap
-                        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:
-                            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
                         
-                        fig, ax = plt.subplots(figsize=(6, 6))
-                        im = ax.imshow(heatmap, cmap='jet', alpha=0.8)
-                        ax.set_title("AI Attention Heatmap", fontweight='bold')
-                        ax.axis('off')
-                        plt.colorbar(im, ax=ax, fraction=0.046, pad=0.04)
-                        st.pyplot(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.")
+                    else:
+                        st.error("Could not generate visualization")
+                else:
+                    st.error("Could not generate attention visualization")
             else:
                 st.info("Upload an image and run classification to see AI attention visualization")
     
@@ -434,21 +410,20 @@ def main():
         st.markdown("""
         ## 👋 Welcome to the Stroke Classification System
         
-        This AI system analyzes brain scan images to detect stroke indicators.
+        This AI system analyzes brain scan images to detect stroke indicators and shows you **exactly where the AI is looking**.
         
         ### 🚀 Features:
         - **Deep Learning Classification**: Advanced CNN architecture
-        - **Visual Attention Maps**: See where the model focuses
+        - **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
-        - **Multiple Visualizations**: Choose how to view AI attention
+        - **Transparent AI**: Understand how the AI makes decisions
         
         ### 📋 How to Use:
-        1. **Check system status** above (should show green checkmarks)
+        1. **Check system status** above (Grad-CAM should show ✅ for real attention)
         2. **Upload a brain scan image** using the sidebar
-        3. **Choose visualization style** (Single Overlay recommended)
-        4. **View classification results** with confidence scores
-        5. **Explore attention visualization** to understand the model's focus
+        3. **View classification results** with confidence scores
+        4. **Explore REAL attention visualization** to see where the AI actually looks
         
         **Get started by uploading an image! 👈**
         """)