Update app.py

app.py

@@ -1,6 +1,8 @@
 import streamlit as st
 from transformers import pipeline
 import torchaudio
+import numpy as np
+import matplotlib.pyplot as plt
 from config import MODEL_ID
 
 # Load the model and pipeline using the model_id variable
@@ -35,17 +37,17 @@ if theme == "Light Green":
         """
         <style>
         body, .stApp {
-            background-color: #e8f5e9; /* Light green background */
+            background-color: #e8f5e9;
         }
         .stApp {
-            color: #004d40; /* Dark green text */
+            color: #004d40;
         }
         .stButton > button, .stFileUpload > div {
-            background-color: #004d40; /* Dark green button and file uploader background */
-            color: white; /* White text */
+            background-color: #004d40;
+            color: white;
         }
         .stButton > button:hover, .stFileUpload > div:hover {
-            background-color: #00332c; /* Darker green on hover */
+            background-color: #00332c;
         }
         </style>
         """,
@@ -56,24 +58,23 @@ elif theme == "Light Blue":
         """
         <style>
         body, .stApp {
-            background-color: #e0f7fa; /* Light blue background */
+            background-color: #e0f7fa;
         }
         .stApp {
-            color: #006064; /* Dark blue text */
+            color: #006064;
         }
         .stButton > button, .stFileUpload > div {
-            background-color: #006064; /* Dark blue button and file uploader background */
-            color: white; /* White text */
+            background-color: #006064;
+            color: white;
         }
         .stButton > button:hover, .stFileUpload > div:hover {
-            background-color: #004d40; /* Darker blue on hover */
+            background-color: #004d40;
         }
         </style>
         """,
         unsafe_allow_html=True
     )
 
-
 # File uploader for audio files
 uploaded_file = st.file_uploader("Upload an audio file", type=["wav", "mp3"])
 
@@ -87,11 +88,34 @@ if uploaded_file is not None:
     with open("temp_audio_file.wav", "wb") as f:
         f.write(audio_bytes)
     
+    # Load audio for visualization
+    waveform, sample_rate = torchaudio.load("temp_audio_file.wav")
+    
+    # Visualization selection
+    viz_type = st.radio("Select visualization type:", ["Waveform", "Spectrogram"])
+    
+    # Create visualization
+    fig, ax = plt.subplots(figsize=(10, 4))
+    if viz_type == "Waveform":
+        time = np.arange(waveform.shape[1]) / sample_rate
+        ax.plot(time, waveform[0].numpy())
+        ax.set_title("Audio Waveform")
+        ax.set_xlabel("Time (s)")
+        ax.set_ylabel("Amplitude")
+        ax.set_xlim([0, time[-1]])
+    else:
+        ax.specgram(waveform[0].numpy(), Fs=sample_rate, cmap='viridis', NFFT=1024, noverlap=512)
+        ax.set_title("Spectrogram")
+        ax.set_xlabel("Time (s)")
+        ax.set_ylabel("Frequency (Hz)")
+    
+    st.pyplot(fig)
+    
     # Classify the audio file
     st.write("Classifying the audio...")
     results = classify_audio("temp_audio_file.wav")
     
-    # Display the classification results in a dedicated output box
+    # Display the classification results
     st.subheader("Classification Results")
     results_box = st.empty()
     results_str = "\n".join([f"{label}: {score:.2f}" for label, score in results.items()])
@@ -105,7 +129,32 @@ for example in examples:
         st.subheader(f"Sample Audio: {example}")
         audio_bytes = open(example, 'rb').read()
         st.audio(audio_bytes, format='audio/wav')
+        
+        # Load audio for visualization
+        waveform, sample_rate = torchaudio.load(example)
+        
+        # Visualization selection
+        viz_type = st.radio("Select visualization type:", ["Waveform", "Spectrogram"], key=example)
+        
+        # Create visualization
+        fig, ax = plt.subplots(figsize=(10, 4))
+        if viz_type == "Waveform":
+            time = np.arange(waveform.shape[1]) / sample_rate
+            ax.plot(time, waveform[0].numpy())
+            ax.set_title("Audio Waveform")
+            ax.set_xlabel("Time (s)")
+            ax.set_ylabel("Amplitude")
+            ax.set_xlim([0, time[-1]])
+        else:
+            ax.specgram(waveform[0].numpy(), Fs=sample_rate, cmap='viridis', NFFT=1024, noverlap=512)
+            ax.set_title("Spectrogram")
+            ax.set_xlabel("Time (s)")
+            ax.set_ylabel("Frequency (Hz)")
+        
+        st.pyplot(fig)
+        
+        # Classification results
         results = classify_audio(example)
         st.write("Results:")
         results_str = "\n".join([f"{label}: {score:.2f}" for label, score in results.items()])
-        st.text(results_str)
+        st.text(results_str)