import streamlit as st
from transformers import AutoModelForAudioClassification, AutoFeatureExtractor, pipeline
import os

# Title of the Streamlit app
st.title('Music Genre Classification')

# Load the model and feature extractor
model = AutoModelForAudioClassification.from_pretrained("sandychoii/distilhubert-finetuned-gtzan-audio-classification")
feature_extractor = AutoFeatureExtractor.from_pretrained("sandychoii/distilhubert-finetuned-gtzan-audio-classification")

# Create the audio classification pipeline
pipe = pipeline(task="audio-classification", model=model, feature_extractor=feature_extractor)

# File uploader to upload an audio file
audio_file = st.file_uploader("Upload an audio file", type=["wav", "mp3", "flac"])

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)
    
    # Display the results
    st.subheader("Classification Results:")
    for label in result:
        st.write(f"**Genre**: {label['label']} | **Confidence**: {label['score']:.4f}")
    
    # Option to play the uploaded audio
    st.audio(audio_file)

    # Clean up the temporary file
    os.remove(temp_path)