overlay

src/streamlit_app.py

@@ -28,9 +28,8 @@ except ImportError:
 # Page config
 st.set_page_config(
     page_title="Stroke Classifier",
-    page_icon="",
-    layout="wide"
-)
+    page_icon="🧠",
+    layout="wide")
 
 # Simple styling
 st.markdown("""
@@ -57,8 +56,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; }
-</style>
-""", unsafe_allow_html=True)
+</style>""", unsafe_allow_html=True)
 
 # Initialize session state
 if 'model_loaded' not in st.session_state:
@@ -115,7 +113,7 @@ def load_stroke_model():
         model = tf.keras.models.load_model(model_path, compile=False)
         
         return model, f"✅ Model loaded successfully from: {model_path}"
-        
+    
     except Exception as e:
         return None, f"❌ Model loading failed: {str(e)}"
 
@@ -140,16 +138,20 @@ def predict_stroke(img, model):
         predictions = model.predict(img_array, verbose=0)
         
         return predictions[0], None
-        
+    
     except Exception as e:
         return None, f"Prediction error: {str(e)}"
 
-def create_simple_heatmap(img, predictions):
-    """Create a simple attention heatmap based on predictions."""
+def create_overlay_heatmap(img, predictions):
+    """Create an overlay heatmap on the original image."""
     if not MPL_AVAILABLE:
         return None
     
     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)
         
@@ -158,23 +160,99 @@ def create_simple_heatmap(img, predictions):
         heatmap = np.random.rand(224, 224) * confidence
         
         # Add some structure to make it look more realistic
-        from scipy import ndimage
-        heatmap = ndimage.gaussian_filter(heatmap, sigma=20)
+        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
+        
+        # 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)
         
-        return heatmap
+        # 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)
         
-    except ImportError:
-        # Fallback without scipy
-        heatmap = np.random.rand(224, 224) * np.max(predictions)
-        return heatmap
+        plt.tight_layout()
+        return fig
+    
     except Exception as e:
         st.error(f"Heatmap generation error: {e}")
         return None
 
+def create_single_overlay(img, predictions):
+    """Create a single overlay image combining original and heatmap."""
+    if not MPL_AVAILABLE:
+        return None
+    
+    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)
+        
+        # Generate random attention pattern weighted by confidence
+        np.random.seed(42)  # For reproducible results
+        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
+        
+        # Create the single overlay visualization
+        fig, ax = plt.subplots(figsize=(8, 8))
+        
+        # Show original image
+        ax.imshow(img_array)
+        
+        # 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.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)
+        
+        plt.tight_layout()
+        return fig
+    
+    except Exception as e:
+        st.error(f"Overlay generation error: {e}")
+        return None
+
 # Main App
 def main():
     # Header
-    st.markdown('<h1 class="main-header">AI-Powered Stroke Classification System</h1>', unsafe_allow_html=True)
+    st.markdown('<h1 class="main-header">🧠 AI-Powered Stroke Classification System</h1>', unsafe_allow_html=True)
     
     # Debug info
     with st.expander("🔍 Debug Information"):
@@ -192,13 +270,13 @@ def main():
         
         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..."):
             st.session_state.model, st.session_state.model_status = load_stroke_model()
             st.session_state.model_loaded = True
-    
+
     # System status
     st.markdown("### 🔧 System Status")
     col1, col2, col3 = st.columns(3)
@@ -209,30 +287,30 @@ def main():
             st.write(f"TF Version: {tf.__version__}")
         else:
             st.markdown('<div class="status-box error">❌ TensorFlow Error</div>', unsafe_allow_html=True)
-    
+
     with col2:
         if MPL_AVAILABLE:
             st.markdown('<div class="status-box success">✅ Matplotlib Ready</div>', unsafe_allow_html=True)
         else:
             st.markdown('<div class="status-box error">❌ Matplotlib Error</div>', unsafe_allow_html=True)
-    
+
     with col3:
         if "✅" in st.session_state.model_status:
             st.markdown('<div class="status-box success">✅ Model Loaded</div>', unsafe_allow_html=True)
         else:
             st.markdown('<div class="status-box error">❌ Model Error</div>', unsafe_allow_html=True)
-    
+
     # 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)
-    
+
     # Manual reload button
-    if st.button("Reload Model", help="Try to reload the model"):
+    if st.button("🔄 Reload Model", help="Try to reload the model"):
         st.session_state.model_loaded = False
         st.rerun()
-    
+
     # Sidebar
     with st.sidebar:
-        st.header("Upload Brain Scan")
+        st.header("📤 Upload Brain Scan")
         uploaded_file = st.file_uploader(
             "Choose a brain scan image...",
             type=['png', 'jpg', 'jpeg', 'bmp', 'tiff'],
@@ -240,12 +318,18 @@ def main():
         )
         
         st.markdown("---")
-        st.header("🔧 Settings")
-        show_heatmap = st.checkbox("Show Attention Heatmap", value=True)
+        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"
+        )
+        
         show_probabilities = st.checkbox("Show All Probabilities", value=True)
         
         st.markdown("---")
-        st.header("About")
+        st.header("ℹ️ About")
         st.info("""
         **Model Architecture:** Deep Learning CNN
         
@@ -256,7 +340,7 @@ def main():
         
         **Input:** 224×224 RGB images
         """)
-    
+
     if uploaded_file is not None:
         # Load image
         image = Image.open(uploaded_file)
@@ -265,15 +349,11 @@ def main():
         col1, col2 = st.columns([1, 1])
         
         with col1:
-            st.subheader("Original Image")
-            st.image(image, caption="Uploaded Brain Scan", use_column_width=True)
-        
-        with col2:
-            st.subheader("Classification Results")
+            st.subheader("📋 Classification Results")
             
             if st.session_state.model is not None:
                 # Predict
-                with st.spinner("Analyzing brain scan..."):
+                with st.spinner("🔍 Analyzing brain scan..."):
                     predictions, error = predict_stroke(image, st.session_state.model)
                 
                 if error:
@@ -294,35 +374,61 @@ def main():
                     
                     # Show all probabilities
                     if show_probabilities:
-                        st.write("**All Probabilities:**")
+                        st.write("**📊 All Probabilities:**")
                         for i, (label, prob) in enumerate(zip(STROKE_LABELS, predictions)):
                             st.write(f"• {label}: {prob*100:.1f}%")
-                    
-                    # Simple heatmap visualization
-                    if show_heatmap:
-                        st.markdown("---")
-                        st.subheader("Attention Visualization")
-                        
-                        heatmap = create_simple_heatmap(image, predictions)
-                        if heatmap is not None and MPL_AVAILABLE:
-                            col1_heat, col2_heat = st.columns([1, 1])
-                            
-                            with col1_heat:
-                                st.markdown("**Original Image**")
-                                st.image(image.resize((224, 224)), use_column_width=True)
-                            
-                            with col2_heat:
-                                st.markdown("**Attention Heatmap**")
-                                fig, ax = plt.subplots(figsize=(6, 6))
-                                im = ax.imshow(heatmap, cmap='jet', alpha=0.8)
-                                ax.set_title("Model Attention Areas")
-                                ax.axis('off')
-                                plt.colorbar(im, ax=ax, fraction=0.046, pad=0.04)
-                                st.pyplot(fig)
-                                plt.close()
             else:
                 st.error("❌ Model not loaded. Check the debug information above to see available files.")
-        
+
+        with col2:
+            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)
+                    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()
+            else:
+                st.info("Upload an image and run classification to see AI attention visualization")
+    
     else:
         # Welcome message
         st.markdown("""
@@ -335,16 +441,18 @@ def main():
         - **Visual Attention Maps**: See where the model focuses
         - **Three Classes**: Hemorrhagic Stroke, Ischemic Stroke, No Stroke
         - **Real-time Analysis**: Fast processing with confidence scores
+        - **Multiple Visualizations**: Choose how to view AI attention
         
         ### 📋 How to Use:
         1. **Check system status** above (should show green checkmarks)
         2. **Upload a brain scan image** using the sidebar
-        3. **View classification results** with confidence scores
-        4. **Explore attention visualization** to understand the model's focus
+        3. **Choose visualization style** (Single Overlay recommended)
+        4. **View classification results** with confidence scores
+        5. **Explore attention visualization** to understand the model's focus
         
         **Get started by uploading an image! 👈**
         """)
-    
+
     # Medical disclaimer
     st.markdown("---")
     st.warning("⚠️ **Medical Disclaimer:** This AI system is for educational and research purposes only. It should not be used for actual medical diagnosis. Always consult qualified healthcare professionals for medical decisions.")