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 | |
| # 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) | |
| # CNN Visualization Section | |
| if st.button("Pass to CNN"): | |
| st.session_state.show_cnn = 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 | |
| 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.") | |