import streamlit as st
from transformers import pipeline
from pydub import AudioSegment
import io

# Load the audio classification pipeline
audio_classification_pipeline = pipeline("audio-classification", model="MIT/ast-finetuned-audioset-10-10-0.4593")

def convert_mp3_to_wav(audio_bytes):
    # Convert MP3 bytes to WAV format
    audio = AudioSegment.from_file(io.BytesIO(audio_bytes), format="mp3")
    wav_bytes = audio.export(format="wav").read()
    return wav_bytes

def classify_audio(audio_file):
    # Perform audio classification
    results = audio_classification_pipeline(audio=audio_file)
    return results

def main():
    st.title('Hugging Face Audio Classification')

    # File uploader for audio file
    st.subheader('Upload Audio File:')
    audio_file = st.file_uploader("Choose an MP3 file", type=["mp3"])

    # Check if audio file is uploaded
    if audio_file is not None:
        st.audio(audio_file, format='audio/mp3')

        # Button to classify audio
        if st.button('Classify'):
            with st.spinner('Classifying...'):
                # Convert MP3 to WAV and classify
                wav_bytes = convert_mp3_to_wav(audio_file.read())
                results = classify_audio(wav_bytes)
            st.success('Classification complete!')
            st.write("Prediction:", results)

if __name__ == '__main__':
    main()