Update app.py

app.py

@@ -6,11 +6,12 @@ import numpy as np
 import plotly.graph_objects as go
 import tempfile
 import os
+import soundfile as sf
 
 # Set page config
 st.set_page_config(page_title="🎵 Music Genre Classifier", layout="wide")
 
-# Custom CSS (unchanged)
+# Custom CSS for UI
 st.markdown("""
 <style>
     .main-title {
@@ -51,28 +52,35 @@ def load_model():
 
 pipe = load_model()
 
+def convert_to_wav(audio_file):
+    """Converts uploaded audio file to WAV format."""
+    with tempfile.NamedTemporaryFile(delete=False, suffix=".wav") as tmp_wav:
+        y, sr = librosa.load(audio_file, sr=None)
+        sf.write(tmp_wav.name, y, sr)
+        return tmp_wav.name
+
 def classify_audio(audio_file):
+    """Classifies the audio file using the loaded model."""
     start_time = time.time()
-    
-    # Save the uploaded file to a temporary file
-    with tempfile.NamedTemporaryFile(delete=False, suffix=".mp3") as tmp_file:
-        tmp_file.write(audio_file.getvalue())
-        tmp_file_path = tmp_file.name
+
+    # Convert to WAV format before passing to the model
+    wav_file = convert_to_wav(audio_file)
     
     try:
-        y, sr = librosa.load(tmp_file_path, sr=None)
-        preds = pipe(y)
+        # 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, y, sr
+        return outputs, prediction_time
     finally:
-        # Make sure to remove the temporary file
-        os.unlink(tmp_file_path)
+        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.
@@ -80,23 +88,27 @@ 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)
     
+    # Classify the uploaded audio
     if st.button("Classify Genre"):
         with st.spinner("Analyzing the music... 🎧"):
             try:
-                results, pred_time, y, sr = classify_audio(uploaded_file)
+                results, pred_time = classify_audio(uploaded_file)
                 
-                # Get top genre
+                # 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)
                 
-                # Create a bar chart using Plotly
+                # Plot the genre probabilities as a bar chart
                 fig = go.Figure(data=[go.Bar(
                     x=list(results.keys()),
                     y=list(results.values()),
@@ -111,7 +123,10 @@ if uploaded_file is not None:
                 )
                 st.plotly_chart(fig, use_container_width=True)
 
-                # Display waveform
+                # 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(
@@ -127,8 +142,10 @@ if uploaded_file is not None:
                 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.")
 
+# 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)
+""", unsafe_allow_html=True)
+