Update src/streamlit_app.py

src/streamlit_app.py

@@ -267,8 +267,31 @@ def make_gradcam_heatmap(img_array, model, layer_name, pred_index=None):
     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
+
 def create_real_gradcam_heatmap(img, model, predictions):
-    """Create a real Grad-CAM heatmap with fallback strategies."""
+    """Create a real Grad-CAM heatmap with enhanced visualization."""
     try:
         # Preprocess image
         img_resized = img.resize((224, 224))
@@ -285,7 +308,7 @@ def create_real_gradcam_heatmap(img, model, predictions):
         layer_name, layer_status = find_best_layer_for_gradcam(model)
         
         if layer_name is None:
-            return None, f"❌ {layer_status}"
+            return None, f"❌ {layer_status}", None
         
         # Generate Grad-CAM heatmap
         heatmap, error = make_gradcam_heatmap(
@@ -296,7 +319,7 @@ def create_real_gradcam_heatmap(img, model, predictions):
         )
         
         if error:
-            return None, f"❌ {error} (Layer: {layer_name})"
+            return None, f"❌ {error} (Layer: {layer_name})", None
         
         if heatmap is not None:
             # Resize heatmap to match input image size
@@ -308,12 +331,23 @@ def create_real_gradcam_heatmap(img, model, predictions):
             else:
                 heatmap_resized = heatmap
             
-            return heatmap_resized, f"✅ Grad-CAM successful using layer: {layer_name}"
+            # 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))
+            }
+            
+            return heatmap_enhanced, f"✅ Grad-CAM successful using layer: {layer_name}", stats
         else:
-            return None, "❌ Failed to generate heatmap"
+            return None, "❌ Failed to generate heatmap", None
             
     except Exception as e:
-        return None, f"❌ Grad-CAM error: {str(e)}"
+        return None, f"❌ Grad-CAM error: {str(e)}", None
 
 def predict_stroke(img, model):
     """Predict stroke type from image."""
@@ -356,12 +390,22 @@ def create_simulated_heatmap(img, predictions):
             mask = (x - center_x)**2 + (y - center_y)**2
             heatmap = np.exp(-mask / (2 * (50**2))) * confidence
         
-        return heatmap, "⚠️ Using simulated heatmap (Grad-CAM failed)"
+        # Enhance simulated heatmap too
+        heatmap_enhanced = enhance_heatmap_contrast(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))
+        }
+        
+        return heatmap_enhanced, "⚠️ Using simulated heatmap (Grad-CAM failed)", stats
     except Exception as e:
-        return None, f"❌ Simulated heatmap error: {str(e)}"
+        return None, f"❌ Simulated heatmap error: {str(e)}", None
 
-def create_overlay_visualization(img, predictions, model, force_gradcam=True):
-    """Create overlay visualization with debugging."""
+def create_enhanced_overlay_visualization(img, predictions, model, force_gradcam=True, colormap='hot'):
+    """Create enhanced overlay visualization with better colors."""
     if not MPL_AVAILABLE:
         return None, "❌ Matplotlib not available"
     
@@ -372,52 +416,64 @@ def create_overlay_visualization(img, predictions, model, force_gradcam=True):
         
         heatmap = None
         status_message = ""
+        stats = None
         
         # 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
+            result = create_real_gradcam_heatmap(img, model, predictions)
+            if result and len(result) == 3:
+                heatmap, gradcam_status, stats = result
+                status_message = gradcam_status
         
         # Fallback to simulated if Grad-CAM failed
         if heatmap is None:
-            heatmap, sim_status = create_simulated_heatmap(img, predictions)
-            if status_message:
-                status_message += f" | {sim_status}"
-            else:
-                status_message = sim_status
+            result = create_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"
         
-        # Create visualization
-        fig, ax = plt.subplots(figsize=(10, 8))
+        # Create enhanced visualization with multiple views
+        fig, axes = plt.subplots(1, 3, figsize=(15, 5))
         
-        # Show original image
-        ax.imshow(img_array)
+        # 1. Original image
+        axes[0].imshow(img_array)
+        axes[0].set_title("Original Image", fontsize=12, fontweight='bold')
+        axes[0].axis('off')
         
-        # Overlay heatmap
-        im = ax.imshow(heatmap, cmap='jet', alpha=0.4, interpolation='bilinear')
+        # 2. Heatmap only with enhanced colors
+        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
+        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 (Grad-CAM)"
+            title = "🎯 Real AI Attention Overlay"
             title_color = 'green'
         else:
-            title = "🎨 Simulated Attention (Grad-CAM Failed)"
+            title = "🎨 Simulated Attention Overlay"
             title_color = 'orange'
         
-        ax.set_title(title, fontsize=14, fontweight='bold', pad=20, color=title_color)
-        ax.axis('off')
-        
-        # Add colorbar
-        cbar = plt.colorbar(im, ax=ax, fraction=0.046, pad=0.04)
-        cbar.set_label('Attention Intensity', rotation=270, labelpad=20)
+        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
+        
+        return fig, status_message, stats
     
     except Exception as e:
-        return None, f"❌ Visualization error: {str(e)}"
+        return None, f"❌ Visualization error: {str(e)}", None
 
 # Main App
 def main():
@@ -479,13 +535,6 @@ def main():
                 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:
@@ -494,12 +543,6 @@ def main():
                 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 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"):
@@ -524,22 +567,27 @@ def main():
             help="Try Grad-CAM even with non-CNN models (experimental)"
         )
         
+        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)
         
         st.markdown("---")
-        st.header("ℹ️ About")
+        st.header("🎨 Color Scheme Guide")
         st.info("""
-        **Model Architecture:** Detected automatically
-        
-        **Classes:**
-        - Hemorrhagic Stroke
-        - Ischemic Stroke  
-        - No Stroke
-        
-        **Input:** 224×224 RGB images
-        
-        **Attention Method:** Grad-CAM (when possible)
+        **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
         """)
 
     if uploaded_file is not None:
@@ -547,7 +595,7 @@ def main():
         image = Image.open(uploaded_file)
         
         # Main content area
-        col1, col2 = st.columns([1, 1])
+        col1, col2 = st.columns([1, 2])
         
         with col1:
             st.subheader("📋 Classification Results")
@@ -582,20 +630,23 @@ def main():
                 st.error("❌ Model not loaded. Check the debug information above to see available files.")
 
         with col2:
-            st.subheader("🎯 AI Attention Visualization")
+            st.subheader("🎯 Enhanced AI Attention Visualization")
             
             if st.session_state.model is not None and 'predictions' in locals() and predictions is not None:
-                # Create overlay visualization
-                with st.spinner("🎨 Generating attention visualization..."):
-                    result = create_overlay_visualization(
+                # Create enhanced overlay visualization
+                with st.spinner("🎨 Generating enhanced attention visualization..."):
+                    result = create_enhanced_overlay_visualization(
                         image, 
                         predictions, 
                         st.session_state.model, 
-                        force_gradcam
+                        force_gradcam,
+                        colormap
                     )
                 
-                if result and len(result) == 2:
-                    overlay_fig, status_message = result
+                if result and len(result) >= 2:
+                    overlay_fig, status_message = result[0], result[1]
+                    stats = result[2] if len(result) > 2 else None
+                    
                     if overlay_fig is not None:
                         st.pyplot(overlay_fig)
                         plt.close()
@@ -608,6 +659,17 @@ def main():
                                 st.warning(f"⚠️ {status_message}")
                             else:
                                 st.info(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}")
                     else:
                         st.error(f"Could not generate visualization: {status_message}")
                 else:
@@ -618,21 +680,20 @@ def main():
     else:
         # Welcome message
         st.markdown("""
-        ## 👋 Welcome to the Stroke Classification System
+        ## 👋 Welcome to the Enhanced Stroke Classification System
         
-        This AI system analyzes brain scan images and attempts to show you where the AI focuses its attention.
+        This system now provides **better color visualization** and **enhanced contrast** for attention heatmaps.
         
-        ### 🚀 Features:
-        - **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
+        ### 🎨 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
         
-        ### 📋 How to Use:
-        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** - real or simulated based on your model
+        ### 🚀 Why Colors Matter:
+        - **Red/Yellow (hot)**: High attention areas
+        - **Blue/Purple**: Low attention areas  
+        - **Enhanced contrast**: Makes subtle patterns visible
         
         **Get started by uploading an image! 👈**
         """)