Update app.py

app.py

@@ -9,6 +9,7 @@ import torch.nn as nn
 import torch.nn.functional as F
 import matplotlib.pyplot as plt
 import plotly.express as px
+import seaborn as sns
 
 # Dummy CNN Model
 class SimpleCNN(nn.Module):
@@ -219,95 +220,178 @@ if uploaded_file is not None:
             st.subheader("CNN Processing Visualization")
             activations, magnitude_tensor = pass_to_cnn(st.session_state.filtered_fft[0])
             
-            # Display input tensor
-            st.write("### Input Magnitude Tensor:")
-            st.image(magnitude_tensor.squeeze().numpy(), 
-                    caption="Magnitude Tensor", 
-                    use_column_width=True,
-                    clamp=True)
+            # Display input tensor with improved visualization
+            st.write("### Input Magnitude Tensor")
+            fig_input, ax_input = plt.subplots(figsize=(8, 8))
+            input_img = magnitude_tensor.squeeze().numpy()
+            im = ax_input.imshow(input_img, cmap='viridis')
+            plt.colorbar(im, ax=ax_input)
+            st.pyplot(fig_input)
             
-            # Display activations with improved visualization
+            # Display activation maps with proper normalization
             st.write("### First Convolution Layer Activations")
             activation = activations.detach().numpy()
             
             if len(activation.shape) == 4:
-                # Create a grid of activation maps
-                cols = 4  # Number of columns in the grid
-                rows = 4  # 16 channels / 4 columns = 4 rows
+                # Create grid layout for activation maps
+                st.write("#### Activation Maps Visualization")
+                cols = 4
+                rows = 4
                 fig, axs = plt.subplots(rows, cols, figsize=(20, 20))
                 
                 for i in range(activation.shape[1]):
-                    act_img = activation[0, i, :, :]
                     ax = axs[i//cols, i%cols]
-                    ax.imshow(act_img, cmap='viridis')
+                    act_img = activation[0, i, :, :]
+                    vmin, vmax = np.percentile(act_img, [1, 99])  # Robust normalization
+                    im = ax.imshow(act_img, cmap='inferno', vmin=vmin, vmax=vmax)
                     ax.set_title(f'Channel {i+1}')
-                    ax.axis('off')
+                    fig.colorbar(im, ax=ax)
                 
+                plt.tight_layout()
                 st.pyplot(fig)
                 
-                # Display sample activation values
-                st.write("### Activation Values Sample")
-                sample_activation = activation[0, 0, :10, :10]  # First 10x10 values
-                st.dataframe(pd.DataFrame(sample_activation))
-            
-            # Additional Steps After Activation Channels
+                # Display activation statistics
+                st.write("#### Activation Value Distribution")
+                flat_activations = activation.flatten()
+                fig_hist = px.histogram(
+                    x=flat_activations,
+                    nbins=100,
+                    title="Activation Value Distribution",
+                    labels={'x': 'Activation Value'}
+                )
+                st.plotly_chart(fig_hist)
+
+            # Second Convolution Layer Visualization
             st.markdown("---")
-            st.subheader("Next Processing Steps in CNN")
-            
-            # Step 2: Second Convolution Layer Visualization
-            st.write("### Second Convolution Layer Features")
+            st.subheader("Second Convolution Layer Features")
             with torch.no_grad():
                 model = SimpleCNN()
-                output, activations = model(magnitude_tensor)
-                second_conv = model.conv2(activations).detach().numpy()
+                _, first_conv = model(magnitude_tensor)
+                second_conv = model.conv2(first_conv).detach().numpy()
             
             if len(second_conv.shape) == 4:
-                cols = 8  # 32 channels / 8 columns = 4 rows
+                # Display sample feature maps
+                st.write("#### Feature Maps Visualization")
+                cols = 8
                 rows = 4
                 fig2, axs2 = plt.subplots(rows, cols, figsize=(20, 10))
                 
-                for i in range(second_conv.shape[1]):
-                    act_img = second_conv[0, i, :, :]
+                for i in range(32):  # For all 32 channels
                     ax = axs2[i//cols, i%cols]
-                    ax.imshow(act_img, cmap='plasma')
-                    ax.set_title(f'Channel {i+1}')
+                    feature_map = second_conv[0, i, :, :]
+                    vmin, vmax = np.percentile(feature_map, [1, 99])
+                    im = ax.imshow(feature_map, cmap='plasma', vmin=vmin, vmax=vmax)
+                    ax.set_title(f'FM {i+1}')
                     ax.axis('off')
                 
+                plt.tight_layout()
                 st.pyplot(fig2)
-            
-            # Step 3: Pooling Layer Visualization
-            st.write("### Adaptive Average Pooling Output")
+
+            # Pooling Layer Visualization
+            st.markdown("---")
+            st.subheader("Pooling Layer Output")
             with torch.no_grad():
                 pooled = F.adaptive_avg_pool2d(torch.tensor(second_conv), (8, 8)).numpy()
 
-            st.write("Pooled Features Shape:", pooled.shape)
+            st.write("#### Pooled Features Dimensionality Reduction")
+
+            # Create a heatmap using seaborn
+            fig_pool, ax_pool = plt.subplots(figsize=(10, 6))
+            sns.heatmap(
+                pooled[0, 0],  # Use the first channel of the pooled features
+                annot=True,     # Show values in each cell
+                fmt=".2f",      # Format values to 2 decimal places
+                cmap="coolwarm",# Use a color map for better visualization
+                ax=ax_pool      # Plot on the created axis
+            )
+            st.pyplot(fig_pool)
 
-            # Normalize and display pooled features
-            pooled_sample = pooled[0, 0]
-            pooled_normalized = (pooled_sample - pooled_sample.min()) / (pooled_sample.max() - pooled_sample.min())
-            st.image(pooled_normalized, 
-                    caption="Sample Pooled Feature Map", 
-                    use_container_width=True,
-                    clamp=True)
+            # Create a grid of pooled feature maps
+            cols = 4
+            rows = 2
+            fig, axs = plt.subplots(rows, cols, figsize=(20, 10))
+
+            for i in range(rows * cols):
+                ax = axs[i // cols, i % cols]
+                sns.heatmap(
+                    pooled[0, i],
+                    annot=True,
+                    fmt=".2f",
+                    cmap="coolwarm",
+                    ax=ax
+                )
+                ax.set_title(f"Channel {i+1}")
+
+            plt.tight_layout()
+            st.pyplot(fig)
+
+            # Fully Connected Layer Visualization
+            st.markdown("---")
+            st.subheader("Fully Connected Layer Analysis")
+            with torch.no_grad():
+                model = SimpleCNN()
+                flattened = model.conv2(model.conv1(magnitude_tensor))
+                flattened = F.adaptive_avg_pool2d(flattened, (8, 8))
+                flattened = flattened.view(flattened.size(0), -1)
+                fc_output = model.fc1(flattened).detach().numpy()
             
-            # Step 4: Final Classification
-            st.write("### Final Classification Scores")
+            st.write("#### FC Layer Activation Patterns")
+            fig_fc = px.imshow(
+                fc_output.T,
+                labels=dict(x="Neurons", y="Features", color="Activation"),
+                color_continuous_scale="viridis"
+            )
+            st.plotly_chart(fig_fc)
+
+            # Final Classification Visualization
+            st.markdown("---")
+            st.subheader("Final Classification Results")
             with torch.no_grad():
                 model = SimpleCNN()
                 output, _ = model(magnitude_tensor)
-                scores = F.softmax(output, dim=1).numpy()
-
+                probabilities = F.softmax(output, dim=1).numpy()[0]
+            
             classes = [f"Class {i}" for i in range(10)]
-            fig3 = px.bar(x=classes, y=scores[0], title="Classification Probabilities")
-            st.plotly_chart(fig3)
+            df = pd.DataFrame({"Class": classes, "Probability": probabilities})
             
-            # Step 5: Full Process Explanation
+            fig_class = px.bar(
+                df,
+                x="Class",
+                y="Probability",
+                color="Probability",
+                color_continuous_scale="tealrose"
+            )
+            st.plotly_chart(fig_class)
+
+            # Full Pipeline Explanation
             st.markdown("""
-            #### Processing Pipeline:
-            1. Input Magnitude Spectrum → 
-            2. Conv1 Features (16 channels) → 
-            3. Conv2 Features (32 channels) → 
-            4. Pooled Features → 
-            5. Fully Connected Layers → 
-            6. Final Classification
-            """)
+                ### Complete Processing Pipeline
+                <div style="
+                    background-color: #f0f2f6;
+                    padding: 30px;
+                    border-radius: 15px;
+                    box-shadow: 0px 4px 6px rgba(0, 0, 0, 0.1);
+                    font-family: 'Arial', sans-serif;
+                    font-size: 16px;
+                    color: #333;
+                    border: 1px solid #dcdcdc;
+                ">
+                <ul style="list-style-type: none; padding-left: 0;">
+                    <li><strong>1. Input Preparation:</strong> Magnitude spectrum from FFT</li>
+                    <li><strong>2. Feature Extraction:</strong>
+                        <ul>
+                            <li>- Conv1: 16 filters (3x3)</li>
+                            <li>- Conv2: 32 filters (3x3)</li>
+                        </ul>
+                    </li>
+                    <li><strong>3. Dimensionality Reduction:</strong> Adaptive average pooling (8x8)</li>
+                    <li><strong>4. Feature Transformation:</strong>
+                        <ul>
+                            <li>- Flattening: 32×8×8 → 2048 features</li>
+                            <li>- FC1: 2048 → 128 dimensions</li>
+                        </ul>
+                    </li>
+                    <li><strong>5. Classification:</strong> FC2: 128 → 10 classes</li>
+                </ul>
+                </div>
+            """, unsafe_allow_html=True)