Update app.py

app.py

@@ -4,10 +4,30 @@ 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):
-    """Apply FFT to each channel of the image and return shifted FFT channels."""
     fft_channels = []
     for channel in cv2.split(image):
         fft = np.fft.fft2(channel)
@@ -16,7 +36,6 @@ def apply_fft(image):
     return fft_channels
 
 def filter_fft_percentage(fft_channels, percentage):
-    """Filter FFT channels to keep top percentage of magnitudes."""
     filtered_fft = []
     for fft_data in fft_channels:
         magnitude = np.abs(fft_data)
@@ -28,7 +47,6 @@ def filter_fft_percentage(fft_channels, percentage):
     return filtered_fft
 
 def inverse_fft(filtered_fft):
-    """Reconstruct image from filtered FFT channels."""
     reconstructed_channels = []
     for fft_data in filtered_fft:
         fft_ishift = np.fft.ifftshift(fft_data)
@@ -37,8 +55,22 @@ def inverse_fft(filtered_fft):
         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):
-    """Create interactive 3D surface plots using Plotly."""
     fig = make_subplots(
         rows=3, cols=2,
         specs=[[{'type': 'scene'}, {'type': 'scene'}],
@@ -51,33 +83,26 @@ def create_3d_plot(fft_channels, downsample_factor=1):
         )
     )
 
-    channel_names = ['Blue', 'Green', 'Red']
-    
     for i, fft_data in enumerate(fft_channels):
-        # Downsample data for better performance
         fft_down = fft_data[::downsample_factor, ::downsample_factor]
         magnitude = np.abs(fft_down)
         phase = np.angle(fft_down)
         
-        # Create grid coordinates
         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)
 
-        # Magnitude plot
         fig.add_trace(
             go.Surface(x=X, y=Y, z=magnitude, colorscale='Viridis', showscale=False),
             row=i+1, col=1
         )
         
-        # Phase plot
         fig.add_trace(
             go.Surface(x=X, y=Y, z=phase, colorscale='Inferno', showscale=False),
             row=i+1, col=2
         )
 
-    # Update layout for better visualization
     fig.update_layout(
         height=1500,
         width=1200,
@@ -93,80 +118,104 @@ def create_3d_plot(fft_channels, downsample_factor=1):
 
 # Streamlit UI
 st.set_page_config(layout="wide")
-st.title("Interactive Frequency Domain Analysis")
-
-# Introduction Text
-st.subheader("Introduction to FFT and Image Filtering")
-st.write(
-    """Fast Fourier Transform (FFT) is a technique to transform an image from the spatial domain to the frequency domain.
-    In this domain, each frequency represents a different aspect of the image's texture and details.
-    By filtering out certain frequencies, you can modify the image's appearance, enhancing or suppressing certain features."""
-)
-
+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:
-    # Read and display original image
     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)
 
-    # Process FFT and store in session state
-    if 'fft_channels' not in st.session_state:
+    # Apply FFT and store in session state
+    if st.session_state.fft_channels is None:
         st.session_state.fft_channels = apply_fft(image)
 
-    # Create a form to submit frequency percentage selection
-    with st.form(key='fft_form'):
-        percentage = st.slider(
-            "Percentage of frequencies to retain:",
-            min_value=0.1, max_value=100.0, value=10.0, step=0.1,
-            help="Adjust the slider to select what portion of frequency components to keep. Lower values blur the image."
-        )
-        submit_button = st.form_submit_button(label="Apply Filter")
+    # 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."
+    )
 
-    if submit_button:
-        # Apply filtering and reconstruct image
-        filtered_fft = filter_fft_percentage(st.session_state.fft_channels, percentage)
-        reconstructed = inverse_fft(filtered_fft)
-        reconstructed_rgb = cv2.cvtColor(reconstructed, cv2.COLOR_BGR2RGB)
+    # 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)
 
-        # Display FFT Data in Table Format
+        # FFT Data Tables
         st.subheader("Frequency Data of Each Channel")
-        fft_data_dict = {}
         for i, channel_name in enumerate(['Blue', 'Green', 'Red']):
-            magnitude = np.abs(st.session_state.fft_channels[i])
-            phase = np.angle(st.session_state.fft_channels[i])
-            fft_data_dict[channel_name] = {'Magnitude': magnitude, 'Phase': phase}
-
-        # Create DataFrames for each channel's FFT data
-        for channel_name, data in fft_data_dict.items():
             st.write(f"### {channel_name} Channel FFT Data")
-            magnitude_df = pd.DataFrame(data['Magnitude'])
-            phase_df = pd.DataFrame(data['Phase'])
+            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))  # Display first 10 rows for brevity
+            st.dataframe(magnitude_df.head(10))
             st.write("#### Phase Data:")
-            st.dataframe(phase_df.head(10))  # Display first 10 rows for brevity
-
-        # Download button for reconstructed image
-        _, encoded_img = cv2.imencode('.png', reconstructed)
-        st.download_button(
-            "Download Reconstructed Image",
-            encoded_img.tobytes(),
-            "reconstructed.png",
-            "image/png"
-        )
+            st.dataframe(phase_df.head(10))
 
-        # 3D visualization controls
+        # 3D Visualization
         st.subheader("3D Frequency Components Visualization")
         downsample = st.slider(
             "Downsampling factor for 3D plots:",
-            min_value=1, max_value=20, value=5,
-            help="Controls the resolution of the 3D surface plots. Higher values improve performance but reduce the plot's detail."
+            1, 20, 5,
+            help="Controls the resolution of the 3D surface plots."
         )
-        
-        # Generate and display 3D plots
-        fig = create_3d_plot(filtered_fft, downsample)
+        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.")