app.py

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.")