import streamlit as st
import joblib
import numpy as np
from PIL import Image
from tensorflow.keras.preprocessing.image import load_img, img_to_array
from tensorflow.keras.applications.resnet50 import ResNet50, preprocess_input

# Load the trained KNN model and class names
knn = joblib.load('knn_model.pkl')  # Replace with the correct path
class_names = joblib.load('class_names.pkl')  # Replace with the correct path

# Function to extract features using ResNet50
def extract_features(img):
    # Load the pre-trained ResNet50 model (without the top layer)
    model = ResNet50(weights='imagenet', include_top=False, pooling='avg')
    # Extract features from the image
    features = model.predict(img)
    return features

# Streamlit app title
st.title("Animal Classification App")

# Description of the app
st.write("This app classifies animals based on uploaded images using a trained KNN model.")

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

# Process the uploaded image and predict
if uploaded_file is not None:
    # Open the image and display it
    image = Image.open(uploaded_file)
    st.image(image, caption="Uploaded Image", use_column_width=True)

    # Resize and preprocess the image
    image = image.resize((224, 224))  # Resize the image to 224x224
    img_array = img_to_array(image)  # Convert image to array
    img_array = np.expand_dims(img_array, axis=0)  # Expand dimensions to match ResNet50 input
    img_array = preprocess_input(img_array)  # Preprocess for ResNet50

    # Extract features using ResNet50
    features = extract_features(img_array)

    # Predict using the trained KNN model
    prediction = knn.predict(features)

    # Display the predicted class
    st.write(f"Prediction: {class_names[prediction[0]]}")