import streamlit as st
import librosa
import numpy as np
import pandas as pd
import IPython.display as ipd
from sklearn.preprocessing import LabelEncoder
from tensorflow.keras.utils import to_categorical
from tensorflow.keras.models import load_model


model = load_model('model/audio_prediction_model')


df = pd.read_csv('UrbanSound8K.csv')

y = np.array(df['class'].tolist())

le = LabelEncoder()
y = to_categorical(le.fit_transform(y))


categories = df.groupby('classID')['class'].unique

def prediction_audio(audiofile, model):

    audio_data, sample_rate = librosa.load(audiofile, res_type='kaiser_fast')
    mel_frequency = np.mean(librosa.feature.melspectrogram(y = audio_data, sr = sample_rate).T, axis=0) # (128,)

    X = mel_frequency.reshape(1, -1) # (1, 128)

    predicted_label = np.argmax(model.predict(X), axis=-1)
    predicted_class = le.inverse_transform(predicted_label)

    return predicted_class

st.title("Gunshot Detection")

# Formulaire pour charger le fichier audio
audio_file = st.file_uploader("Select audio file", type=["wav", "mp3"])

if audio_file is not None:
    # Affichez le fichier audio
    st.audio(audio_file, format='audio/wav')

    # Bouton pour déclencher la prédiction
    if st.button("Predict"):
        # Effectuez la prédiction
        prediction = prediction_audio(audio_file, model)
        
        # Affichez le résultat de la prédiction
        st.write("Result : ", prediction)