ahmad-fakhar commited on
Commit
8ba87fb
β€’
1 Parent(s): 800a812

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +144 -32
app.py CHANGED
@@ -1,42 +1,154 @@
1
  import streamlit as st
2
- from transformers import TFAutoModelForAudioClassification, AutoFeatureExtractor, pipeline
 
 
 
 
 
3
  import os
 
4
 
5
- # Title of the Streamlit app
6
- st.title('Music Genre Classification')
7
 
8
- # Load the model and feature extractor
9
- model = TFAutoModelForAudioClassification.from_pretrained("sandychoii/distilhubert-finetuned-gtzan-audio-classificationyour-tensorflow-compatible-model")
10
- feature_extractor = AutoFeatureExtractor.from_pretrained("sandychoii/distilhubert-finetuned-gtzan-audio-classification")
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
11
 
12
- # Create the audio classification pipeline
13
- pipe = pipeline(task="audio-classification", model=model, feature_extractor=feature_extractor)
 
14
 
15
- # File uploader to upload an audio file
16
- audio_file = st.file_uploader("Upload an audio file", type=["wav", "mp3", "flac"])
17
 
18
- if audio_file is not None:
19
- # Save the uploaded file to a temporary directory
20
- if not os.path.exists("temp_audio"):
21
- os.makedirs("temp_audio")
22
-
23
- temp_path = os.path.join("temp_audio", audio_file.name)
24
- with open(temp_path, "wb") as f:
25
- f.write(audio_file.getbuffer())
26
-
27
- # Display a loading message
28
- st.text("Classifying the audio file... Please wait.")
29
-
30
- # Perform classification
31
- result = pipe(temp_path)
32
 
33
- # Display the results
34
- st.subheader("Classification Results:")
35
- for label in result:
36
- st.write(f"**Genre**: {label['label']} | **Confidence**: {label['score']:.4f}")
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
37
 
38
- # Option to play the uploaded audio
39
- st.audio(audio_file)
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
40
 
41
- # Clean up the temporary file
42
- os.remove(temp_path)
 
 
 
 
 
1
  import streamlit as st
2
+ import time
3
+ from transformers import pipeline
4
+ import librosa
5
+ import numpy as np
6
+ import plotly.graph_objects as go
7
+ import tempfile
8
  import os
9
+ import soundfile as sf
10
 
11
+ # Set page config
12
+ st.set_page_config(page_title="🎡 Music Genre Classifier", layout="wide")
13
 
14
+ # Custom CSS for UI
15
+ st.markdown("""
16
+ <style>
17
+ .main-title {
18
+ font-size: 3rem;
19
+ color: #1DB954;
20
+ text-align: center;
21
+ padding: 2rem 0;
22
+ text-shadow: 2px 2px 4px rgba(0,0,0,0.1);
23
+ }
24
+ .sub-title {
25
+ font-size: 1.5rem;
26
+ color: #191414;
27
+ text-align: center;
28
+ margin-bottom: 2rem;
29
+ }
30
+ .stAudio {
31
+ margin: 2rem auto;
32
+ display: block;
33
+ }
34
+ .genre-result {
35
+ font-size: 2rem;
36
+ font-weight: bold;
37
+ text-align: center;
38
+ color: #1DB954;
39
+ margin: 1rem 0;
40
+ }
41
+ .prediction-time {
42
+ font-size: 1.2rem;
43
+ color: #191414;
44
+ text-align: center;
45
+ }
46
+ </style>
47
+ """, unsafe_allow_html=True)
48
 
49
+ @st.cache_resource
50
+ def load_model():
51
+ return pipeline("audio-classification", model="juangtzi/wav2vec2-base-finetuned-gtzan")
52
 
53
+ pipe = load_model()
 
54
 
55
+ def convert_to_wav(audio_file):
56
+ """Converts uploaded audio file to WAV format."""
57
+ with tempfile.NamedTemporaryFile(delete=False, suffix=".wav") as tmp_wav:
58
+ # Use soundfile to load and save the audio file as WAV
59
+ audio_data, samplerate = sf.read(audio_file)
60
+ sf.write(tmp_wav.name, audio_data, samplerate)
61
+ return tmp_wav.name
62
+
63
+ def classify_audio(audio_file):
64
+ """Classifies the audio file using the loaded model."""
65
+ start_time = time.time()
66
+
67
+ # Convert to WAV format before passing to the model
68
+ wav_file = convert_to_wav(audio_file)
69
 
70
+ try:
71
+ # Use the wav file with the model
72
+ preds = pipe(wav_file)
73
+ outputs = {p["label"]: p["score"] for p in preds}
74
+ end_time = time.time()
75
+ prediction_time = end_time - start_time
76
+ return outputs, prediction_time
77
+ finally:
78
+ os.unlink(wav_file) # Remove the temp file
79
+
80
+ # Page title and subtitle
81
+ st.markdown("<h1 class='main-title'>🎡 Music Genre Classifier</h1>", unsafe_allow_html=True)
82
+ st.markdown("<p class='sub-title'>Upload a music file and let AI detect its genre!</p>", unsafe_allow_html=True)
83
+
84
+ # Sidebar with model and dataset information
85
+ st.sidebar.title("About")
86
+ st.sidebar.info("""
87
+ This app uses a fine-tuned wav2vec2-base model to classify music genres.
88
+ Model: juangtzi/wav2vec2-base-finetuned-gtzan
89
+ Dataset: GTZAN
90
+ """)
91
+
92
+ # Upload file section
93
+ uploaded_file = st.file_uploader("Choose an audio file", type=["wav", "mp3", "ogg"])
94
+
95
+ if uploaded_file is not None:
96
+ # Display the uploaded audio file
97
+ st.audio(uploaded_file)
98
 
99
+ # Classify the uploaded audio
100
+ if st.button("Classify Genre"):
101
+ with st.spinner("Analyzing the music... 🎧"):
102
+ try:
103
+ results, pred_time = classify_audio(uploaded_file)
104
+
105
+ # Get the top predicted genre
106
+ top_genre = max(results, key=results.get)
107
+
108
+ # Display the top predicted genre
109
+ st.markdown(f"<h2 class='genre-result'>Detected Genre: {top_genre.capitalize()}</h2>", unsafe_allow_html=True)
110
+ st.markdown(f"<p class='prediction-time'>Prediction Time: {pred_time:.2f} seconds</p>", unsafe_allow_html=True)
111
+
112
+ # Plot the genre probabilities as a bar chart
113
+ fig = go.Figure(data=[go.Bar(
114
+ x=list(results.keys()),
115
+ y=list(results.values()),
116
+ marker_color='#1DB954'
117
+ )])
118
+ fig.update_layout(
119
+ title="Genre Probabilities",
120
+ xaxis_title="Genre",
121
+ yaxis_title="Probability",
122
+ paper_bgcolor='rgba(0,0,0,0)',
123
+ plot_bgcolor='rgba(0,0,0,0)'
124
+ )
125
+ st.plotly_chart(fig, use_container_width=True)
126
+
127
+ # # Load the audio for displaying waveform
128
+ # y, sr = librosa.load(uploaded_file, sr=None)
129
+
130
+ # # Plot the audio waveform
131
+ # st.subheader("Audio Waveform")
132
+ # fig_waveform = go.Figure(data=[go.Scatter(y=y, mode='lines', line=dict(color='#1DB954'))])
133
+ # fig_waveform.update_layout(
134
+ # title="Audio Waveform",
135
+ # xaxis_title="Time",
136
+ # yaxis_title="Amplitude",
137
+ # paper_bgcolor='rgba(0,0,0,0)',
138
+ # plot_bgcolor='rgba(0,0,0,0)'
139
+ # )
140
+ # st.plotly_chart(fig_waveform, use_container_width=True)
141
+
142
+ # 🎈 Show balloons after successfully displaying the results
143
+ st.balloons()
144
+
145
+ except Exception as e:
146
+ st.error(f"An error occurred while processing the audio: {str(e)}")
147
+ st.info("Please try uploading the file again or use a different audio file.")
148
 
149
+ # Footer
150
+ st.markdown("""
151
+ <div style='text-align: center; margin-top: 2rem;'>
152
+ <p>Created with ❀️ by AI. Powered by Streamlit and Hugging Face Transformers.</p>
153
+ </div>
154
+ """, unsafe_allow_html=True)