Update app.py

app.py

@@ -1,42 +1,154 @@
 import streamlit as st
-from transformers import TFAutoModelForAudioClassification, AutoFeatureExtractor, pipeline
+import time
+from transformers import pipeline
+import librosa
+import numpy as np
+import plotly.graph_objects as go
+import tempfile
 import os
+import soundfile as sf
 
-# Title of the Streamlit app
-st.title('Music Genre Classification')
+# Set page config
+st.set_page_config(page_title="🎵 Music Genre Classifier", layout="wide")
 
-# Load the model and feature extractor
-model = TFAutoModelForAudioClassification.from_pretrained("sandychoii/distilhubert-finetuned-gtzan-audio-classificationyour-tensorflow-compatible-model")
-feature_extractor = AutoFeatureExtractor.from_pretrained("sandychoii/distilhubert-finetuned-gtzan-audio-classification")
+# Custom CSS for UI
+st.markdown("""
+<style>
+    .main-title {
+        font-size: 3rem;
+        color: #1DB954;
+        text-align: center;
+        padding: 2rem 0;
+        text-shadow: 2px 2px 4px rgba(0,0,0,0.1);
+    }
+    .sub-title {
+        font-size: 1.5rem;
+        color: #191414;
+        text-align: center;
+        margin-bottom: 2rem;
+    }
+    .stAudio {
+        margin: 2rem auto;
+        display: block;
+    }
+    .genre-result {
+        font-size: 2rem;
+        font-weight: bold;
+        text-align: center;
+        color: #1DB954;
+        margin: 1rem 0;
+    }
+    .prediction-time {
+        font-size: 1.2rem;
+        color: #191414;
+        text-align: center;
+    }
+</style>
+""", unsafe_allow_html=True)
 
-# Create the audio classification pipeline
-pipe = pipeline(task="audio-classification", model=model, feature_extractor=feature_extractor)
+@st.cache_resource
+def load_model():
+    return pipeline("audio-classification", model="juangtzi/wav2vec2-base-finetuned-gtzan")
 
-# File uploader to upload an audio file
-audio_file = st.file_uploader("Upload an audio file", type=["wav", "mp3", "flac"])
+pipe = load_model()
 
-if audio_file is not None:
-    # Save the uploaded file to a temporary directory
-    if not os.path.exists("temp_audio"):
-        os.makedirs("temp_audio")
-    
-    temp_path = os.path.join("temp_audio", audio_file.name)
-    with open(temp_path, "wb") as f:
-        f.write(audio_file.getbuffer())
-    
-    # Display a loading message
-    st.text("Classifying the audio file... Please wait.")
-    
-    # Perform classification
-    result = pipe(temp_path)
+def convert_to_wav(audio_file):
+    """Converts uploaded audio file to WAV format."""
+    with tempfile.NamedTemporaryFile(delete=False, suffix=".wav") as tmp_wav:
+        # Use soundfile to load and save the audio file as WAV
+        audio_data, samplerate = sf.read(audio_file)
+        sf.write(tmp_wav.name, audio_data, samplerate)
+        return tmp_wav.name
+
+def classify_audio(audio_file):
+    """Classifies the audio file using the loaded model."""
+    start_time = time.time()
+
+    # Convert to WAV format before passing to the model
+    wav_file = convert_to_wav(audio_file)
     
-    # Display the results
-    st.subheader("Classification Results:")
-    for label in result:
-        st.write(f"**Genre**: {label['label']} | **Confidence**: {label['score']:.4f}")
+    try:
+        # Use the wav file with the model
+        preds = pipe(wav_file)
+        outputs = {p["label"]: p["score"] for p in preds}
+        end_time = time.time()
+        prediction_time = end_time - start_time
+        return outputs, prediction_time
+    finally:
+        os.unlink(wav_file)  # Remove the temp file
+
+# Page title and subtitle
+st.markdown("<h1 class='main-title'>🎵 Music Genre Classifier</h1>", unsafe_allow_html=True)
+st.markdown("<p class='sub-title'>Upload a music file and let AI detect its genre!</p>", unsafe_allow_html=True)
+
+# Sidebar with model and dataset information
+st.sidebar.title("About")
+st.sidebar.info("""
+This app uses a fine-tuned wav2vec2-base model to classify music genres.
+Model: juangtzi/wav2vec2-base-finetuned-gtzan
+Dataset: GTZAN
+""")
+
+# Upload file section
+uploaded_file = st.file_uploader("Choose an audio file", type=["wav", "mp3", "ogg"])
+
+if uploaded_file is not None:
+    # Display the uploaded audio file
+    st.audio(uploaded_file)
     
-    # Option to play the uploaded audio
-    st.audio(audio_file)
+    # Classify the uploaded audio
+    if st.button("Classify Genre"):
+        with st.spinner("Analyzing the music... 🎧"):
+            try:
+                results, pred_time = classify_audio(uploaded_file)
+                
+                # Get the top predicted genre
+                top_genre = max(results, key=results.get)
+                
+                # Display the top predicted genre
+                st.markdown(f"<h2 class='genre-result'>Detected Genre: {top_genre.capitalize()}</h2>", unsafe_allow_html=True)
+                st.markdown(f"<p class='prediction-time'>Prediction Time: {pred_time:.2f} seconds</p>", unsafe_allow_html=True)
+                
+                # Plot the genre probabilities as a bar chart
+                fig = go.Figure(data=[go.Bar(
+                    x=list(results.keys()),
+                    y=list(results.values()),
+                    marker_color='#1DB954'
+                )])
+                fig.update_layout(
+                    title="Genre Probabilities",
+                    xaxis_title="Genre",
+                    yaxis_title="Probability",
+                    paper_bgcolor='rgba(0,0,0,0)',
+                    plot_bgcolor='rgba(0,0,0,0)'
+                )
+                st.plotly_chart(fig, use_container_width=True)
+
+                # # Load the audio for displaying waveform
+                # y, sr = librosa.load(uploaded_file, sr=None)
+                
+                # # Plot the audio waveform
+                # st.subheader("Audio Waveform")
+                # fig_waveform = go.Figure(data=[go.Scatter(y=y, mode='lines', line=dict(color='#1DB954'))])
+                # fig_waveform.update_layout(
+                #     title="Audio Waveform",
+                #     xaxis_title="Time",
+                #     yaxis_title="Amplitude",
+                #     paper_bgcolor='rgba(0,0,0,0)',
+                #     plot_bgcolor='rgba(0,0,0,0)'
+                # )
+                # st.plotly_chart(fig_waveform, use_container_width=True)
+
+                # 🎈 Show balloons after successfully displaying the results
+                st.balloons()
+            
+            except Exception as e:
+                st.error(f"An error occurred while processing the audio: {str(e)}")
+                st.info("Please try uploading the file again or use a different audio file.")
 
-    # Clean up the temporary file
-    os.remove(temp_path)
+# Footer
+st.markdown("""
+<div style='text-align: center; margin-top: 2rem;'>
+    <p>Created with ❤️ by AI. Powered by Streamlit and Hugging Face Transformers.</p>
+</div>
+""", unsafe_allow_html=True)