Spaces:
Sleeping
Sleeping
| import streamlit as st | |
| import numpy as np | |
| import cv2 | |
| import plotly.graph_objects as go | |
| from plotly.subplots import make_subplots | |
| 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): | |
| def __init__(self): | |
| super(SimpleCNN, self).__init__() | |
| self.conv1 = nn.Conv2d(1, 16, kernel_size=3, padding=1) | |
| self.conv2 = nn.Conv2d(16, 32, kernel_size=3, padding=1) | |
| self.fc1 = nn.Linear(32 * 8 * 8, 128) | |
| self.fc2 = nn.Linear(128, 10) | |
| def forward(self, x): | |
| x1 = F.relu(self.conv1(x)) # First conv layer activation | |
| x2 = F.relu(self.conv2(x1)) | |
| x3 = F.adaptive_avg_pool2d(x2, (8, 8)) | |
| x4 = x3.view(x3.size(0), -1) | |
| x5 = F.relu(self.fc1(x4)) | |
| x6 = self.fc2(x5) | |
| return x6, x1 # Return both output and first layer activations | |
| # FFT processing functions | |
| def apply_fft(image): | |
| fft_channels = [] | |
| for channel in cv2.split(image): | |
| fft = np.fft.fft2(channel) | |
| fft_shifted = np.fft.fftshift(fft) | |
| fft_channels.append(fft_shifted) | |
| return fft_channels | |
| def filter_fft_percentage(fft_channels, percentage): | |
| filtered_fft = [] | |
| for fft_data in fft_channels: | |
| magnitude = np.abs(fft_data) | |
| sorted_mag = np.sort(magnitude.flatten())[::-1] | |
| num_keep = int(len(sorted_mag) * percentage / 100) | |
| threshold = sorted_mag[num_keep - 1] if num_keep > 0 else 0 | |
| mask = magnitude >= threshold | |
| filtered_fft.append(fft_data * mask) | |
| return filtered_fft | |
| def inverse_fft(filtered_fft): | |
| reconstructed_channels = [] | |
| for fft_data in filtered_fft: | |
| fft_ishift = np.fft.ifftshift(fft_data) | |
| img_reconstructed = np.fft.ifft2(fft_ishift).real | |
| img_normalized = cv2.normalize(img_reconstructed, None, 0, 255, cv2.NORM_MINMAX) | |
| reconstructed_channels.append(img_normalized.astype(np.uint8)) | |
| return cv2.merge(reconstructed_channels) | |
| # CNN Pass Visualization | |
| def pass_to_cnn(fft_image): | |
| model = SimpleCNN() | |
| magnitude_tensor = torch.tensor(np.abs(fft_image), dtype=torch.float32).unsqueeze(0).unsqueeze(0) | |
| with torch.no_grad(): | |
| output, activations = model(magnitude_tensor) | |
| # Ensure activations have the correct shape [batch_size, channels, height, width] | |
| if len(activations.shape) == 3: | |
| activations = activations.unsqueeze(0) # Add batch dimension if missing | |
| return activations, magnitude_tensor | |
| # 3D plotting function | |
| def create_3d_plot(fft_channels, downsample_factor=1): | |
| fig = make_subplots( | |
| rows=3, cols=2, | |
| specs=[[{'type': 'scene'}, {'type': 'scene'}], | |
| [{'type': 'scene'}, {'type': 'scene'}], | |
| [{'type': 'scene'}, {'type': 'scene'}]], | |
| subplot_titles=( | |
| 'Blue - Magnitude', 'Blue - Phase', | |
| 'Green - Magnitude', 'Green - Phase', | |
| 'Red - Magnitude', 'Red - Phase' | |
| ) | |
| ) | |
| for i, fft_data in enumerate(fft_channels): | |
| fft_down = fft_data[::downsample_factor, ::downsample_factor] | |
| magnitude = np.abs(fft_down) | |
| phase = np.angle(fft_down) | |
| rows, cols = magnitude.shape | |
| x = np.linspace(-cols//2, cols//2, cols) | |
| y = np.linspace(-rows//2, rows//2, rows) | |
| X, Y = np.meshgrid(x, y) | |
| fig.add_trace( | |
| go.Surface(x=X, y=Y, z=magnitude, colorscale='Viridis', showscale=False), | |
| row=i+1, col=1 | |
| ) | |
| fig.add_trace( | |
| go.Surface(x=X, y=Y, z=phase, colorscale='Inferno', showscale=False), | |
| row=i+1, col=2 | |
| ) | |
| fig.update_layout( | |
| height=1500, | |
| width=1200, | |
| margin=dict(l=0, r=0, b=0, t=30), | |
| scene_camera=dict(eye=dict(x=1.5, y=1.5, z=0.5)), | |
| scene=dict( | |
| xaxis=dict(title='Frequency X'), | |
| yaxis=dict(title='Frequency Y'), | |
| zaxis=dict(title='Magnitude/Phase') | |
| ) | |
| ) | |
| return fig | |
| # Streamlit UI | |
| st.set_page_config(layout="wide") | |
| st.title("Interactive Frequency Domain Analysis with CNN") | |
| # Initialize session state | |
| if 'fft_channels' not in st.session_state: | |
| st.session_state.fft_channels = None | |
| if 'filtered_fft' not in st.session_state: | |
| st.session_state.filtered_fft = None | |
| if 'reconstructed' not in st.session_state: | |
| st.session_state.reconstructed = None | |
| if 'show_cnn' not in st.session_state: | |
| st.session_state.show_cnn = False | |
| # Upload image | |
| uploaded_file = st.file_uploader("Upload an image", type=['png', 'jpg', 'jpeg']) | |
| if uploaded_file is not None: | |
| file_bytes = np.frombuffer(uploaded_file.getvalue(), np.uint8) | |
| image = cv2.imdecode(file_bytes, cv2.IMREAD_COLOR) | |
| image_rgb = cv2.cvtColor(image, cv2.COLOR_BGR2RGB) | |
| st.image(image_rgb, caption="Original Image", use_column_width=True) | |
| # Apply FFT and store in session state | |
| if st.session_state.fft_channels is None: | |
| st.session_state.fft_channels = apply_fft(image) | |
| # Frequency percentage slider | |
| percentage = st.slider( | |
| "Percentage of frequencies to retain:", | |
| 0.1, 100.0, 10.0, 0.1, | |
| help="Adjust the slider to select what portion of frequency components to keep." | |
| ) | |
| # Apply FFT filter | |
| if st.button("Apply Filter"): | |
| st.session_state.filtered_fft = filter_fft_percentage(st.session_state.fft_channels, percentage) | |
| st.session_state.reconstructed = inverse_fft(st.session_state.filtered_fft) | |
| st.session_state.show_cnn = False # Reset CNN visualization | |
| # Display reconstructed image and FFT data | |
| if st.session_state.reconstructed is not None: | |
| reconstructed_rgb = cv2.cvtColor(st.session_state.reconstructed, cv2.COLOR_BGR2RGB) | |
| st.image(reconstructed_rgb, caption="Reconstructed Image", use_column_width=True) | |
| # FFT Data Tables | |
| st.subheader("Frequency Data of Each Channel") | |
| for i, channel_name in enumerate(['Blue', 'Green', 'Red']): | |
| st.write(f"### {channel_name} Channel FFT Data") | |
| magnitude_df = pd.DataFrame(np.abs(st.session_state.filtered_fft[i])) | |
| phase_df = pd.DataFrame(np.angle(st.session_state.filtered_fft[i])) | |
| st.write("#### Magnitude Data:") | |
| st.dataframe(magnitude_df.head(10)) | |
| st.write("#### Phase Data:") | |
| st.dataframe(phase_df.head(10)) | |
| # 3D Visualization | |
| st.subheader("3D Frequency Components Visualization") | |
| downsample = st.slider( | |
| "Downsampling factor for 3D plots:", | |
| 1, 20, 5, | |
| help="Controls the resolution of the 3D surface plots." | |
| ) | |
| 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 | |
| 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") | |
| 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 activations with improved visualization | |
| 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 | |
| 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 | |
| """) | |