Update app.py

app.py

@@ -12,26 +12,47 @@ import matplotlib.pyplot as plt
 import plotly.express as px
 import soundfile as sf
 from scipy.signal import stft
+import math
 
-# Dummy CNN Model for Audio
+# -------------------------------
+# CNN Model for Audio Analysis
+# -------------------------------
 class AudioCNN(nn.Module):
     def __init__(self):
         super(AudioCNN, self).__init__()
+        # Convolutional layers
         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 * 32 * 8, 128)  # Adjusted for typical spectrogram size
-        self.fc2 = nn.Linear(128, 10)
+        self.conv3 = nn.Conv2d(32, 64, kernel_size=3, padding=1)
+        # Pooling layer
+        self.pool = nn.MaxPool2d(kernel_size=2, stride=2)
+        # Fully connected layers (with dynamic sizing)
+        self.fc1 = None
+        self.fc2 = nn.Linear(256, 128)
+        self.fc3 = nn.Linear(128, 10)
+        # Dropout for regularization
+        self.dropout = nn.Dropout(0.5)
 
     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, 32))
-        x4 = x3.view(x3.size(0), -1)
-        x5 = F.relu(self.fc1(x4))
-        x6 = self.fc2(x5)
-        return x6, x1
-
-# Audio processing functions
+        x1 = F.relu(self.conv1(x))
+        x2 = self.pool(x1)
+        x3 = F.relu(self.conv2(x2))
+        x4 = self.pool(x3)
+        x5 = F.relu(self.conv3(x4))
+        x6 = self.pool(x5)
+        if self.fc1 is None:
+            fc1_input_size = x6.numel() // x6.size(0)
+            self.fc1 = nn.Linear(fc1_input_size, 256)
+        x7 = x6.view(x6.size(0), -1)
+        x8 = F.relu(self.fc1(x7))
+        x9 = self.dropout(x8)
+        x10 = F.relu(self.fc2(x9))
+        x11 = self.fc3(x10)
+        return x11, [x2, x4, x6], x8
+
+# -------------------------------
+# Audio Processing Functions
+# -------------------------------
 def load_audio(file):
     audio, sr = librosa.load(file, sr=None, mono=True)
     return audio, sr
@@ -53,11 +74,13 @@ def filter_fft(fft, percentage):
 def create_spectrogram(audio, sr):
     n_fft = 2048
     hop_length = 512
-    stft = librosa.stft(audio, n_fft=n_fft, hop_length=hop_length)
-    spectrogram = np.abs(stft)
+    S = librosa.stft(audio, n_fft=n_fft, hop_length=hop_length)
+    spectrogram = np.abs(S)
     return spectrogram, n_fft, hop_length
 
-# Visualization functions
+# -------------------------------
+# Visualization Functions
+# -------------------------------
 def plot_waveform(audio, sr, title):
     fig = go.Figure()
     time = np.arange(len(audio)) / sr
@@ -65,41 +88,111 @@ def plot_waveform(audio, sr, title):
     fig.update_layout(title=title, xaxis_title='Time (s)', yaxis_title='Amplitude')
     return fig
 
+def create_waveform_table(audio, sr, num_samples=100):
+    time = np.arange(len(audio)) / sr
+    indices = np.linspace(0, len(audio)-1, num_samples, dtype=int)
+    df = pd.DataFrame({"Time (s)": time[indices], "Amplitude": audio[indices]})
+    return df
+
 def plot_fft(magnitude, phase, sr):
     fig = make_subplots(rows=2, cols=1, subplot_titles=('Magnitude Spectrum', 'Phase Spectrum'))
     freq = np.fft.fftfreq(len(magnitude), 1/sr)
-    
     fig.add_trace(go.Scatter(x=freq, y=magnitude, mode='lines', name='Magnitude'), row=1, col=1)
     fig.add_trace(go.Scatter(x=freq, y=phase, mode='lines', name='Phase'), row=2, col=1)
-    
     fig.update_xaxes(title_text='Frequency (Hz)', row=1, col=1)
     fig.update_xaxes(title_text='Frequency (Hz)', row=2, col=1)
     fig.update_yaxes(title_text='Magnitude', row=1, col=1)
     fig.update_yaxes(title_text='Phase (radians)', row=2, col=1)
-    
     return fig
 
-def plot_3d_fft(magnitude, phase, sr):
+def plot_fft_bands(magnitude, phase, sr):
     freq = np.fft.fftfreq(len(magnitude), 1/sr)
-    fig = go.Figure(data=[go.Scatter3d(
-        x=freq,
-        y=magnitude,
-        z=phase,
+    pos_mask = freq >= 0
+    freq, magnitude, phase = freq[pos_mask], magnitude[pos_mask], phase[pos_mask]
+    bass_mask = (freq >= 20) & (freq < 250)
+    mid_mask = (freq >= 250) & (freq < 4000)
+    treble_mask = (freq >= 4000) & (freq <= sr/2)
+    fig = make_subplots(rows=2, cols=1, subplot_titles=('Magnitude Spectrum by Bands', 'Phase Spectrum by Bands'))
+    fig.add_trace(go.Scatter(x=freq[bass_mask], y=magnitude[bass_mask], mode='lines', name='Bass'), row=1, col=1)
+    fig.add_trace(go.Scatter(x=freq[mid_mask], y=magnitude[mid_mask], mode='lines', name='Mid'), row=1, col=1)
+    fig.add_trace(go.Scatter(x=freq[treble_mask], y=magnitude[treble_mask], mode='lines', name='Treble'), row=1, col=1)
+    fig.add_trace(go.Scatter(x=freq[bass_mask], y=phase[bass_mask], mode='lines', name='Bass'), row=2, col=1)
+    fig.add_trace(go.Scatter(x=freq[mid_mask], y=phase[mid_mask], mode='lines', name='Mid'), row=2, col=1)
+    fig.add_trace(go.Scatter(x=freq[treble_mask], y=phase[treble_mask], mode='lines', name='Treble'), row=2, col=1)
+    fig.update_xaxes(title_text='Frequency (Hz)', row=1, col=1)
+    fig.update_xaxes(title_text='Frequency (Hz)', row=2, col=1)
+    fig.update_yaxes(title_text='Magnitude', row=1, col=1)
+    fig.update_yaxes(title_text='Phase (radians)', row=2, col=1)
+    return fig
+
+def create_fft_table(magnitude, phase, sr, num_samples=100):
+    freq = np.fft.fftfreq(len(magnitude), 1/sr)
+    pos_mask = freq >= 0
+    freq, magnitude, phase = freq[pos_mask], magnitude[pos_mask], phase[pos_mask]
+    indices = np.linspace(0, len(freq)-1, num_samples, dtype=int)
+    df = pd.DataFrame({
+        "Frequency (Hz)": freq[indices],
+        "Magnitude": magnitude[indices],
+        "Phase (radians)": phase[indices]
+    })
+    return df
+
+def plot_3d_polar_fft(magnitude, phase, sr):
+    # Get positive frequencies
+    freq = np.fft.fftfreq(len(magnitude), 1/sr)
+    pos_mask = freq >= 0
+    freq, mag, ph = freq[pos_mask], magnitude[pos_mask], phase[pos_mask]
+    # Convert polar to Cartesian coordinates
+    x = mag * np.cos(ph)
+    y = mag * np.sin(ph)
+    z = freq  # Use frequency as z-axis
+
+    # Downsample the data to avoid huge message sizes.
+    # Compute a decimation factor so that approximately 500 points are plotted.
+    step = max(1, len(x) // 500)
+    x, y, z, ph = x[::step], y[::step], z[::step], ph[::step]
+
+    # Create a coarser grid for the contour surface.
+    n_rep = 10
+    X_surface = np.tile(x, (n_rep, 1))
+    Y_surface = np.tile(y, (n_rep, 1))
+    Z_surface = np.tile(z, (n_rep, 1))
+    
+    surface = go.Surface(
+        x=X_surface,
+        y=Y_surface,
+        z=Z_surface,
+        colorscale='Viridis',
+        opacity=0.6,
+        showscale=False,
+        contours={
+            "x": {"show": True, "start": float(np.min(x)), "end": float(np.max(x)), "size": float((np.max(x)-np.min(x))/10)},
+            "y": {"show": True, "start": float(np.min(y)), "end": float(np.max(y)), "size": float((np.max(y)-np.min(y))/10)},
+            "z": {"show": True, "start": float(np.min(z)), "end": float(np.max(z)), "size": float((np.max(z)-np.min(z))/10)},
+        },
+    )
+    
+    scatter = go.Scatter3d(
+        x=x,
+        y=y,
+        z=z,
         mode='markers',
         marker=dict(
-            size=5,
-            color=phase,  # Color by phase
-            colorscale='Viridis',  # Choose a colorscale
-            opacity=0.8
+            size=3,
+            color=ph,  # color by phase
+            colorscale='Viridis',
+            opacity=0.8,
+            colorbar=dict(title='Phase (radians)')
         )
-    )])
+    )
     
+    fig = go.Figure(data=[surface, scatter])
     fig.update_layout(scene=dict(
-        xaxis_title='Frequency (Hz)',
-        yaxis_title='Magnitude',
-        zaxis_title='Phase (radians)'
-    ))
-    
+        xaxis_title='Real Component',
+        yaxis_title='Imaginary Component',
+        zaxis_title='Frequency (Hz)',
+        camera=dict(eye=dict(x=1.5, y=1.5, z=0.5))
+    ), margin=dict(l=0, r=0, b=0, t=0))
     return fig
 
 def plot_spectrogram(spectrogram, sr, hop_length):
@@ -110,117 +203,185 @@ def plot_spectrogram(spectrogram, sr, hop_length):
     plt.title('Spectrogram')
     return fig
 
-def create_fft_table(magnitude, phase, sr):
-    freq = np.fft.fftfreq(len(magnitude), 1/sr)
-    df = pd.DataFrame({
-        'Frequency (Hz)': freq,
-        'Magnitude': magnitude,
-        'Phase (radians)': phase
-    })
+def create_spectrogram_table(spectrogram, num_rows=10, num_cols=10):
+    sub_spec = spectrogram[:num_rows, :num_cols]
+    df = pd.DataFrame(sub_spec, 
+                      index=[f'Freq Bin {i}' for i in range(sub_spec.shape[0])],
+                      columns=[f'Time Bin {j}' for j in range(sub_spec.shape[1])])
     return df
 
-# Streamlit UI
+def create_activation_table(activation, num_rows=10, num_cols=10):
+    sub_act = activation[:num_rows, :num_cols]
+    df = pd.DataFrame(sub_act, 
+                      index=[f'Row {i}' for i in range(sub_act.shape[0])],
+                      columns=[f'Col {j}' for j in range(sub_act.shape[1])])
+    return df
+
+# -------------------------------
+# Streamlit UI & Main App
+# -------------------------------
 st.set_page_config(layout="wide")
-st.title("Audio Frequency Analysis with CNN")
+st.title("Audio Frequency Analysis with CNN and FFT")
 
-# Initialize session state
-if 'audio_data' not in st.session_state:
-    st.session_state.audio_data = None
-if 'sr' not in st.session_state:
-    st.session_state.sr = None
-if 'fft' not in st.session_state:
-    st.session_state.fft = None
+st.markdown("""
+    ### Welcome to the Audio Frequency Analysis Tool!
+    This application allows you to:
+    - **Upload an audio file** and visualize its waveform along with a data table.
+    - **Analyze frequency components** using FFT (with both 2D and enhanced 3D polar plots).
+    - **Highlight frequency bands:** Bass (20–250 Hz), Mid (250–4000 Hz), Treble (4000 Hz to Nyquist).
+    - **Filter frequency components** and reconstruct the waveform.
+    - **Generate a spectrogram** for time-frequency analysis with a sample data table.
+    - **Inspect CNN activations** (pooling and dense layers) arranged in grid layouts.
+    - **Final Audio Classification:** Classify the audio for gender (Male/Female) and tone.
+""")
 
 # File uploader
-uploaded_file = st.file_uploader("Upload an audio file", type=['wav', 'mp3', 'ogg'])
+uploaded_file = st.file_uploader("Upload an audio file (WAV, MP3, OGG)", type=['wav', 'mp3', 'ogg'])
 
 if uploaded_file is not None:
-    # Load and process audio
     audio, sr = load_audio(uploaded_file)
-    st.session_state.audio_data = audio
-    st.session_state.sr = sr
     
-    # Display original waveform
-    st.subheader("Original Audio Waveform")
-    st.plotly_chart(plot_waveform(audio, sr, "Original Waveform"), use_container_width=True)
+    # --- Section 1: Raw Audio Waveform ---
+    st.header("1. Raw Audio Waveform")
+    st.markdown("""
+        The waveform represents the amplitude over time.
+        **Graph:** Amplitude vs. Time.
+        **Data Table:** Sampled values.
+    """)
+    waveform_fig = plot_waveform(audio, sr, "Original Waveform")
+    st.plotly_chart(waveform_fig, use_container_width=True)
+    st.dataframe(create_waveform_table(audio, sr))
     
-    # Apply FFT
+    # --- Section 2: Frequency Domain Analysis ---
+    st.header("2. Frequency Domain Analysis")
+    st.markdown("""
+        **FFT Analysis:** Decompose the audio into frequency components.
+        - **Magnitude Spectrum:** Strength of frequencies.
+        - **Phase Spectrum:** Phase angles.
+    """)
     fft, magnitude, phase = apply_fft(audio)
-    st.session_state.fft = fft
-    
-    # Display FFT results
-    st.subheader("Frequency Domain Analysis")
-    st.plotly_chart(plot_fft(magnitude, phase, sr), use_container_width=True)
-    
-    # 3D FFT Plot
-    st.subheader("3D Frequency Domain Analysis")
-    st.plotly_chart(plot_3d_fft(magnitude, phase, sr), use_container_width=True)
+    col1, col2 = st.columns(2)
+    with col1:
+        st.subheader("2D FFT Plot")
+        st.plotly_chart(plot_fft(magnitude, phase, sr), use_container_width=True)
+    with col2:
+        st.subheader("Enhanced 3D Polar FFT Plot with Contours")
+        st.plotly_chart(plot_3d_polar_fft(magnitude, phase, sr), use_container_width=True)
+    st.subheader("FFT Data Table (Sampled)")
+    st.dataframe(create_fft_table(magnitude, phase, sr))
+    st.subheader("Frequency Bands: Bass, Mid, Treble")
+    st.plotly_chart(plot_fft_bands(magnitude, phase, sr), use_container_width=True)
     
-    # FFT Table
-    st.subheader("FFT Values Table")
-    fft_table = create_fft_table(magnitude, phase, sr)
-    st.dataframe(fft_table)
-    
-    # Frequency filtering
+    # --- Section 3: Frequency Filtering ---
+    st.header("3. Frequency Filtering")
+    st.markdown("""
+        Filter the audio signal by retaining a percentage of the strongest frequencies.
+        Adjust the slider for retention percentage.
+        **Graph:** Filtered waveform.
+        **Data Table:** Sampled values.
+    """)
     percentage = st.slider("Percentage of frequencies to retain:", 0.1, 100.0, 10.0, 0.1)
-    
     if st.button("Apply Frequency Filter"):
-        filtered_fft = filter_fft(st.session_state.fft, percentage)
+        filtered_fft = filter_fft(fft, percentage)
         reconstructed = np.fft.ifft(filtered_fft).real
-        
-        # Display reconstructed waveform
-        st.subheader("Reconstructed Audio")
-        st.plotly_chart(plot_waveform(reconstructed, sr, "Filtered Waveform"), use_container_width=True)
-        
-        # Play audio
-        st.audio(reconstructed, sample_rate=sr)
+        col1, col2 = st.columns(2)
+        with col1:
+            st.plotly_chart(plot_waveform(reconstructed, sr, "Filtered Waveform"), use_container_width=True)
+        with col2:
+            st.audio(reconstructed, sample_rate=sr)
+        st.dataframe(create_waveform_table(reconstructed, sr))
     
-    # Spectrogram creation
-    st.subheader("Spectrogram Analysis")
+    # --- Section 4: Spectrogram Analysis ---
+    st.header("4. Spectrogram Analysis")
+    st.markdown("""
+        A spectrogram shows how frequency content evolves over time.
+        **Graph:** Spectrogram (log-frequency scale).
+        **Data Table:** A subsection of the spectrogram matrix.
+    """)
     spectrogram, n_fft, hop_length = create_spectrogram(audio, sr)
     st.pyplot(plot_spectrogram(spectrogram, sr, hop_length))
+    st.dataframe(create_spectrogram_table(spectrogram))
     
-    # CNN Processing
-    if st.button("Process with CNN"):
-        # Convert spectrogram to tensor
+    # --- Section 5: CNN Analysis (Pooling & Dense Activations) ---
+    st.header("5. CNN Analysis: Pooling and Dense Activations")
+    st.markdown("""
+        Instead of classification probabilities, inspect internal activations:
+        - **Pooling Layer Outputs:** Arranged in a grid layout.
+        - **Dense Layer Activation:** Feature vector from the dense layer.
+    """)
+    if st.button("Run CNN Analysis"):
         spec_tensor = torch.tensor(spectrogram[np.newaxis, np.newaxis, ...], dtype=torch.float32)
-        
         model = AudioCNN()
         with torch.no_grad():
-            output, activations = model(spec_tensor)
-        
-        # Visualize activations
-        st.subheader("CNN Layer Activations")
-        
-        # Input spectrogram
-        st.write("### Input Spectrogram")
-        fig_input, ax = plt.subplots()
-        ax.imshow(spectrogram, aspect='auto', origin='lower')
-        st.pyplot(fig_input)
+            output, pooling_outputs, dense_activation = model(spec_tensor)
+        for idx, activation in enumerate(pooling_outputs):
+            st.subheader(f"Pooling Layer {idx+1} Output")
+            act = activation[0].cpu().numpy()
+            num_channels = act.shape[0]
+            ncols = 4
+            nrows = math.ceil(num_channels / ncols)
+            fig, axes = plt.subplots(nrows, ncols, figsize=(3*ncols, 3*nrows))
+            axes = axes.flatten()
+            for i in range(nrows * ncols):
+                if i < num_channels:
+                    axes[i].imshow(act[i], aspect='auto', origin='lower', cmap='viridis')
+                    axes[i].set_title(f'Channel {i+1}', fontsize=8)
+                    axes[i].axis('off')
+                else:
+                    axes[i].axis('off')
+            st.pyplot(fig)
+            st.markdown("**Data Table for Pooling Layer Activation (Channel 1, Sampled)**")
+            df_act = create_activation_table(act[0])
+            st.dataframe(df_act)
+        st.subheader("Dense Layer Activation")
+        dense_act = dense_activation[0].cpu().numpy()
+        df_dense = pd.DataFrame({
+            "Feature Index": np.arange(len(dense_act)),
+            "Activation Value": dense_act
+        })
+        st.plotly_chart(px.bar(df_dense, x="Feature Index", y="Activation Value"), use_container_width=True)
+        st.dataframe(df_dense)
+    
+    # --- Section 6: Final Audio Classification (Gender & Tone) ---
+    st.header("6. Final Audio Classification: Gender and Tone")
+    st.markdown("""
+        In this final step, a pretrained model classifies the audio as Male or Female,
+        and determines its tone (High Tone vs. Low Tone).
         
-        # First conv layer activations
-        st.write("### First Convolution Layer Activations")
-        activation = activations.detach().numpy()[0]
+        **Note:** This example uses a placeholder model. Replace the dummy model and random outputs with your actual pretrained model.
+    """)
+    if st.button("Run Final Classification"):
+        # Extract MFCC features as an example (adjust as needed)
+        mfccs = librosa.feature.mfcc(y=audio, sr=sr, n_mfcc=40)
+        features = np.mean(mfccs, axis=1)  # average over time
+        features_tensor = torch.tensor(features, dtype=torch.float32).unsqueeze(0)
         
-        cols = 4
-        rows = 4
-        fig, axs = plt.subplots(rows, cols, figsize=(20, 20))
-        for i in range(16):
-            ax = axs[i//cols, i%cols]
-            ax.imshow(activation[i], aspect='auto', origin='lower')
-            ax.set_title(f'Channel {i+1}')
-        plt.tight_layout()
-        st.pyplot(fig)
+        # Dummy classifier model for demonstration
+        class GenderToneClassifier(nn.Module):
+            def __init__(self):
+                super(GenderToneClassifier, self).__init__()
+                self.fc = nn.Linear(40, 4)  # 4 outputs: [Male, Female, High Tone, Low Tone]
+            def forward(self, x):
+                return self.fc(x)
         
-        # Classification results
-        st.write("### Classification Output")
-        probabilities = F.softmax(output, dim=1).numpy()[0]
-        classes = [f"Class {i}" for i in range(10)]
-        df = pd.DataFrame({"Class": classes, "Probability": probabilities})
-        fig = px.bar(df, x="Class", y="Probability", color="Probability")
-        st.plotly_chart(fig)
-
-# Add some styling
+        classifier = GenderToneClassifier()
+        # In practice, load your pretrained weights here.
+        with torch.no_grad():
+            output = classifier(features_tensor)
+            probs = F.softmax(output, dim=1).numpy()[0]
+        # Interpret outputs: assume first 2 are gender, next 2 are tone.
+        gender = "Male" if probs[0] > probs[1] else "Female"
+        tone = "High Tone" if probs[2] > probs[3] else "Low Tone"
+        st.markdown(f"**Predicted Gender:** {gender}")
+        st.markdown(f"**Predicted Tone:** {tone}")
+        categories = ["Male", "Female", "High Tone", "Low Tone"]
+        df_class = pd.DataFrame({"Category": categories, "Probability": probs})
+        st.plotly_chart(px.bar(df_class, x="Category", y="Probability"), use_container_width=True)
+        st.dataframe(df_class)
+
+# -------------------------------
+# Style Enhancements
+# -------------------------------
 st.markdown("""
     <style>
         .stButton>button {