import streamlit as st
import numpy as np
import tensorflow as tf
from PIL import Image
import json

# ============================================================
# 📦 LOAD MODEL
# ============================================================

MODEL_PATH = "animal_model.keras"
model = tf.keras.models.load_model(MODEL_PATH)

# ============================================================
# 📂 LOAD CLASS LABELS
# ============================================================

with open("class_labels.json", "r") as f:
    class_labels = json.load(f)

class_names = list(class_labels.keys())

# ============================================================
# 🖥️ STREAMLIT UI
# ============================================================

st.title("🐾 Animal Classification App")
st.write("Upload an image and the model will predict the animal.")

uploaded_file = st.file_uploader("Choose an image", type=["jpg", "png", "jpeg"])

IMG_SIZE = (160, 160)

if uploaded_file is not None:
    image = Image.open(uploaded_file)
    st.image(image, caption="Uploaded Image", use_container_width=True)

    # Preprocess
    img = image.resize(IMG_SIZE)
    img_array = np.array(img) / 255.0
    img_array = np.expand_dims(img_array, axis=0)

    # Prediction
    predictions = model.predict(img_array)
    predicted_class = class_names[np.argmax(predictions)]

    st.subheader("🔍 Prediction:")
    st.success(predicted_class)