Update app.py

app.py

@@ -7,6 +7,8 @@ import pandas as pd
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
+import matplotlib.pyplot as plt
+import plotly.express as px
 
 # Dummy CNN Model
 class SimpleCNN(nn.Module):
@@ -182,9 +184,36 @@ if uploaded_file is not None:
         fig = create_3d_plot(st.session_state.filtered_fft, downsample)
         st.plotly_chart(fig, use_container_width=True)
 
+        # Custom CSS to style the button
+        st.markdown("""
+            <style>
+                .centered-button {
+                    display: flex;
+                    justify-content: center;
+                    align-items: center;
+                    margin-top: 20px;
+                }
+                .stButton>button {
+                    padding: 20px 40px;
+                    font-size: 20px;
+                    background-color: #4CAF50;
+                    color: white;
+                    border: none;
+                    border-radius: 10px;
+                    cursor: pointer;
+                }
+                .stButton>button:hover {
+                    background-color: #45a049;
+                }
+            </style>
+        """, unsafe_allow_html=True)
+
         # CNN Visualization Section
-        if st.button("Pass to CNN"):
-            st.session_state.show_cnn = True
+        with st.container():
+            st.markdown('<div class="centered-button">', unsafe_allow_html=True)
+            if st.button("Pass to CNN"):
+                st.session_state.show_cnn = True
+            st.markdown('</div>', unsafe_allow_html=True)
 
         if st.session_state.show_cnn:
             st.subheader("CNN Processing Visualization")
@@ -197,25 +226,88 @@ if uploaded_file is not None:
                     use_column_width=True,
                     clamp=True)
             
-            # Display activations
+            # Display activations with improved visualization
             st.write("### First Convolution Layer Activations")
             activation = activations.detach().numpy()
             
-            # Check the shape of the activation tensor
-            if len(activation.shape) == 4:  # [batch_size, channels, height, width]
-                for i in range(activation.shape[1]):  # Loop through channels
-                    act_img = activation[0, i, :, :]  # Select the first batch and current channel
-                    act_img_normalized = (act_img - act_img.min()) / (act_img.max() - act_img.min())  # Normalize
-                    
-                    # Display activation map
-                    st.write(f"#### Activation Channel {i+1}")
-                    st.image(act_img_normalized, 
-                             caption=f"Activation Channel {i+1}", 
-                             use_column_width=True)
-                    
-                    # Display activation values in a table
-                    st.write("##### Activation Values:")
-                    activation_df = pd.DataFrame(act_img)
-                    st.dataframe(activation_df)
-            else:
-                st.error(f"Unexpected activation shape: {activation.shape}. Expected 4 dimensions.")
+            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
+                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')
+                    ax.set_title(f'Channel {i+1}')
+                    ax.axis('off')
+                
+                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
+            st.markdown("---")
+            st.subheader("Next Processing Steps in CNN")
+            
+            # Step 2: Second Convolution Layer Visualization
+            st.write("### Second Convolution Layer Features")
+            with torch.no_grad():
+                model = SimpleCNN()
+                output, activations = model(magnitude_tensor)
+                second_conv = model.conv2(activations).detach().numpy()
+            
+            if len(second_conv.shape) == 4:
+                cols = 8  # 32 channels / 8 columns = 4 rows
+                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, :, :]
+                    ax = axs2[i//cols, i%cols]
+                    ax.imshow(act_img, cmap='plasma')
+                    ax.set_title(f'Channel {i+1}')
+                    ax.axis('off')
+                
+                st.pyplot(fig2)
+            
+            # Step 3: Pooling Layer Visualization
+            st.write("### Adaptive Average Pooling Output")
+            with torch.no_grad():
+                pooled = F.adaptive_avg_pool2d(torch.tensor(second_conv), (8, 8)).numpy()
+
+            st.write("Pooled Features Shape:", pooled.shape)
+
+            # 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)
+            
+            # Step 4: Final Classification
+            st.write("### Final Classification Scores")
+            with torch.no_grad():
+                model = SimpleCNN()
+                output, _ = model(magnitude_tensor)
+                scores = F.softmax(output, dim=1).numpy()
+
+            classes = [f"Class {i}" for i in range(10)]
+            fig3 = px.bar(x=classes, y=scores[0], title="Classification Probabilities")
+            st.plotly_chart(fig3)
+            
+            # Step 5: Full Process 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
+            """)