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

# Load the trained KNN model and class names
model = joblib.load('knn_model.joblib')
with open('class_names.txt', 'r') as f:
    class_names = f.readlines()
class_names = [x.strip() for x in class_names]

# Load pre-trained ResNet50 model for feature extraction
resnet_model = ResNet50(weights='imagenet', include_top=False, pooling='avg')

# Streamlit app
st.title('Animal Image Classifier')
st.write('Upload an image to classify it.')

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

if uploaded_file is not None:
    # Process the image
    img = load_img(uploaded_file, target_size=(32, 32))
    img = img_to_array(img)
    img = np.expand_dims(img, axis=0)
    img = preprocess_input(img)

    # Extract features
    features = resnet_model.predict(img)

    # Make prediction
    prediction = model.predict(features)
    predicted_class = class_names[prediction[0]]

    # Display result
    st.image(uploaded_file, caption='Uploaded Image.', use_column_width=True)
    st.write(f"Predicted Class: {predicted_class}")