Update app.py

app.py

@@ -1,196 +1,118 @@
 import streamlit as st
-import os
-import tempfile
-from transformers import pipeline, AutoProcessor, AutoModelForAudioClassification
+import time
+from transformers import pipeline
 import librosa
-import soundfile as sf
-from streamlit.components.v1 import html
-import base64
+import numpy as np
+import plotly.graph_objects as go
 
-# Set page configuration
-st.set_page_config(page_title="Music Genre Classifier", layout="wide")
+# Set page config
+st.set_page_config(page_title="🎵 Music Genre Classifier", layout="wide")
 
-# Custom CSS for styling (supports both light and dark modes)
-custom_css = """
+# Custom CSS
+st.markdown("""
 <style>
-    .stApp {
-        transition: background-color 0.3s ease;
-    }
     .main-title {
-        font-size: 3em;
-        font-weight: bold;
+        font-size: 3rem;
+        color: #1DB954;
         text-align: center;
-        margin-bottom: 30px;
+        padding: 2rem 0;
+        text-shadow: 2px 2px 4px rgba(0,0,0,0.1);
     }
     .sub-title {
-        font-size: 1.5em;
+        font-size: 1.5rem;
+        color: #191414;
         text-align: center;
-        margin-bottom: 20px;
+        margin-bottom: 2rem;
     }
-    .result-container {
-        border-radius: 10px;
-        padding: 20px;
-        margin-top: 20px;
+    .stAudio {
+        margin: 2rem auto;
+        display: block;
     }
     .genre-result {
-        font-size: 2em;
+        font-size: 2rem;
         font-weight: bold;
         text-align: center;
+        color: #1DB954;
+        margin: 1rem 0;
     }
-    .confidence-bar {
-        height: 30px;
-        border-radius: 15px;
-        transition: background-color 0.3s ease;
-    }
-    /* Light mode styles */
-    .light-mode .stApp {
-        background-color: #f0f0f5;
-    }
-    .light-mode .main-title, .light-mode .sub-title {
-        color: #1e1e1e;
-    }
-    .light-mode .result-container {
-        background-color: #ffffff;
-        box-shadow: 0 4px 6px rgba(0, 0, 0, 0.1);
-    }
-    .light-mode .genre-result {
-        color: #2c3e50;
-    }
-    .light-mode .confidence-bar {
-        background-color: #3498db;
-    }
-    /* Dark mode styles */
-    .dark-mode .stApp {
-        background-color: #1e1e1e;
-    }
-    .dark-mode .main-title, .dark-mode .sub-title {
-        color: #f0f0f5;
-    }
-    .dark-mode .result-container {
-        background-color: #2c2c2c;
-        box-shadow: 0 4px 6px rgba(255, 255, 255, 0.1);
-    }
-    .dark-mode .genre-result {
-        color: #3498db;
-    }
-    .dark-mode .confidence-bar {
-        background-color: #3498db;
+    .prediction-time {
+        font-size: 1.2rem;
+        color: #191414;
+        text-align: center;
     }
 </style>
-"""
-
-# Render custom CSS
-st.markdown(custom_css, unsafe_allow_html=True)
+""", unsafe_allow_html=True)
 
-# Function to load the model
 @st.cache_resource
 def load_model():
-    try:
-        processor = AutoProcessor.from_pretrained("sandychoii/distilhubert-finetuned-gtzan-audio-classification")
-        model = AutoModelForAudioClassification.from_pretrained("sandychoii/distilhubert-finetuned-gtzan-audio-classification")
-        pipe = pipeline("audio-classification", model=model, feature_extractor=processor)
-        return pipe
-    except Exception as e:
-        st.error(f"Error loading the model: {str(e)}")
-        st.info("Please check your internet connection and try again. If the problem persists, the model might be temporarily unavailable.")
-        return None
+    return pipeline("audio-classification", model="juangtzi/wav2vec2-base-finetuned-gtzan")
 
-# Load the model
 pipe = load_model()
 
-# Function to classify audio
 def classify_audio(audio_file):
-    try:
-        # Load audio file
-        y, sr = librosa.load(audio_file, sr=None)
-        
-        # Ensure the audio is at least 3 seconds long (model requirement)
-        if len(y) < 3 * sr:
-            y = librosa.util.fix_length(y, size=3 * sr)
-        
-        # Classification
-        result = pipe(y, sampling_rate=sr)
-        return result
-    except Exception as e:
-        st.error(f"Error during classification: {str(e)}")
-        return None
-
-# Function to toggle between light and dark mode
-def toggle_theme():
-    if 'theme' not in st.session_state:
-        st.session_state.theme = 'light'
-    if st.session_state.theme == 'light':
-        st.session_state.theme = 'dark'
-    else:
-        st.session_state.theme = 'light'
-    st.experimental_rerun()
-
-# Main app
-def main():
-    # Set theme class
-    theme_class = 'light-mode' if st.session_state.get('theme', 'light') == 'light' else 'dark-mode'
-    st.markdown(f'<div class="{theme_class}">', unsafe_allow_html=True)
-
-    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)
-
-    # Theme toggle button
-    if st.button("Toggle Light/Dark Mode"):
-        toggle_theme()
-
-    # File uploader
-    uploaded_file = st.file_uploader("Choose an audio file", type=["wav", "mp3", "ogg"])
-
-    if uploaded_file is not None:
-        # Display audio player
-        st.audio(uploaded_file)
-
-        # Classify button
-        if st.button("Classify Genre"):
-            if pipe is None:
-                st.error("Model is not loaded. Please check your internet connection and try again.")
-            else:
-                with st.spinner("Analyzing the music... 🎧"):
-                    # Save uploaded file temporarily
-                    with tempfile.NamedTemporaryFile(delete=False, suffix=os.path.splitext(uploaded_file.name)[1]) as tmp_file:
-                        tmp_file.write(uploaded_file.getvalue())
-                        tmp_file_path = tmp_file.name
-
-                    # Perform classification
-                    result = classify_audio(tmp_file_path)
-
-                    # Remove temporary file
-                    os.unlink(tmp_file_path)
-
-                    if result:
-                        # Display results
-                        st.markdown("<div class='result-container'>", unsafe_allow_html=True)
-                        st.markdown(f"<h2 class='genre-result'>Detected Genre: {result[0]['label'].capitalize()}</h2>", unsafe_allow_html=True)
-                        
-                        # Display confidence bar
-                        confidence = result[0]['score']
-                        st.markdown(f"<div class='confidence-bar' style='width: {confidence*100}%;'></div>", unsafe_allow_html=True)
-                        st.write(f"Confidence: {confidence:.2%}")
-                        
-                        # Display top 3 predictions
-                        st.write("Top 3 Predictions:")
-                        for r in result[:3]:
-                            st.write(f"- {r['label'].capitalize()}: {r['score']:.2%}")
-                        st.markdown("</div>", unsafe_allow_html=True)
-
-    # Add information about the model
-    st.sidebar.title("About")
-    st.sidebar.info("This app uses a fine-tuned DistilHuBERT model to classify music genres. It can identify genres like rock, pop, hip-hop, classical, and more!")
-
-    # Add a footer
-    footer_html = """
-    <div style="position: fixed; bottom: 0; width: 100%; text-align: center; padding: 10px;">
-        <p>Created with ❤️ by AI. Powered by Streamlit and Hugging Face Transformers.</p>
-    </div>
-    """
-    st.markdown(footer_html, unsafe_allow_html=True)
-
-    st.markdown('</div>', unsafe_allow_html=True)
-
-if __name__ == "__main__":
-    main()
+    start_time = time.time()
+    y, sr = librosa.load(audio_file, sr=None)
+    preds = pipe(y)
+    outputs = {p["label"]: p["score"] for p in preds}
+    end_time = time.time()
+    prediction_time = end_time - start_time
+    return outputs, prediction_time
+
+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)
+
+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
+""")
+
+uploaded_file = st.file_uploader("Choose an audio file", type=["wav", "mp3", "ogg"])
+
+if uploaded_file is not None:
+    st.audio(uploaded_file)
+    
+    if st.button("Classify Genre"):
+        with st.spinner("Analyzing the music... 🎧"):
+            results, pred_time = classify_audio(uploaded_file)
+            
+            # Get top genre
+            top_genre = max(results, key=results.get)
+            
+            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
+            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)
+
+            # Display waveform
+            st.subheader("Audio Waveform")
+            y, sr = librosa.load(uploaded_file, sr=None)
+            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)
+
+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)