import streamlit as st
import numpy as np
import joblib
import tensorflow as tf
from tensorflow.keras.applications import ResNet50
from tensorflow.keras.applications.resnet50 import preprocess_input
from tensorflow.keras.preprocessing.image import img_to_array
import cv2

# Load KNN model and class labels
knn_model = joblib.load("knn_model.pkl")
class_labels = joblib.load("class_labels.pkl")

# Load ResNet50 model (same as used in training)
resnet_model = ResNet50(weights="imagenet", include_top=False, input_shape=(224, 224, 3))
resnet_model = tf.keras.Model(inputs=resnet_model.input, outputs=tf.keras.layers.GlobalAveragePooling2D()(resnet_model.output))

def extract_features(image):
    """Extract features from image using ResNet50 (same as training)."""
    image = cv2.resize(image, (224, 224))  # Resize image
    image = img_to_array(image)
    image = np.expand_dims(image, axis=0)
    image = preprocess_input(image)
    
    features = resnet_model.predict(image)  # Extract features
    return features.reshape(1, -1)  # Ensure shape matches KNN input

def predict_animal(image):
    """Predict the class of the input image."""
    processed_image = extract_features(image)  # Extract features
    prediction = knn_model.predict(processed_image)[0]  # Predict class index
    return list(class_labels.keys())[list(class_labels.values()).index(prediction)]  # Convert index to label

# Streamlit UI
st.title("🐾 Animal Image Classifier")
uploaded_file = st.file_uploader("📤 Upload an image", type=["jpg", "png", "jpeg"])

if uploaded_file:
    image = cv2.imdecode(np.frombuffer(uploaded_file.read(), np.uint8), 1)
    st.image(image, caption="📷 Uploaded Image", use_container_width=True)
    
    if st.button("🔍 Identify"):
        prediction = predict_animal(image)
        st.success(f"🎯 Predicted Animal: {prediction}")