import streamlit as st
import torch
import torchaudio
import joblib
from model import YourCNNModelClass
from utils import load_model, preprocess_audio, predict

# Load label encoder and get classes
label_encoder = joblib.load("label_encoder.joblib")
labels = label_encoder.classes_

# Load PyTorch model with correct number of classes
model = load_model("animal-sound-cnn.pth", YourCNNModelClass, num_classes=len(labels))

# Streamlit UI
st.title("Animal Sound Classification")

uploaded_file = st.file_uploader("Upload a WAV audio file", type=["wav"])

if uploaded_file:
    with open("temp.wav", "wb") as f:
        f.write(uploaded_file.read())

    st.audio("temp.wav")

    # Preprocess audio and predict
    input_tensor = preprocess_audio("temp.wav")
    pred_label = predict(model, input_tensor, labels)

    st.write(f"Predicted animal: **{pred_label}**")