import streamlit as st
import numpy as np
import pickle
from tensorflow.keras.applications.mobilenet_v2 import MobileNetV2, preprocess_input
from tensorflow.keras.preprocessing.image import img_to_array
from tensorflow.keras.models import Model
from PIL import Image

# Load saved model and class names
with open("knn_model.pkl", "rb") as f:
    knn = pickle.load(f)

with open("class_mapping.pkl", "rb") as f:
    classes = pickle.load(f)

# Load MobileNetV2 feature extractor
base_model = MobileNetV2(weights="imagenet", include_top=False, pooling="avg", input_shape=(224, 224, 3))

st.title("🐾 Animal Classifier using KNN & MobileNetV2")

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

if uploaded_file is not None:
    img = Image.open(uploaded_file).convert("RGB")
    st.image(img, caption="Uploaded Image", use_column_width=True)

    # Preprocess and extract features
    img = img.resize((224, 224))
    img_array = img_to_array(img)
    img_array = preprocess_input(img_array)
    features = base_model.predict(np.expand_dims(img_array, axis=0), verbose=0)
    
    # Predict using KNN
    pred = knn.predict(features)[0]
    predicted_class = classes[pred]

    st.markdown(f"### 🔍 Predicted Class: `{predicted_class}`")