import streamlit as st
import joblib
import numpy as np
from tensorflow.keras.preprocessing import image
from tensorflow.keras.applications.resnet50 import 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]

# 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=(224, 224))  # Resize image to match model input
    img = img_to_array(img)  # Convert to array
    img = preprocess_input(img)  # Preprocess image for ResNet50

    # Make prediction
    img = np.expand_dims(img, axis=0)  # Add batch dimension
    features = model.predict(img)  # Extract features using the model
    prediction = model.predict(features)  # Get prediction

    # Show the result
    predicted_class = class_names[prediction[0]]  # Get the class name
    st.image(uploaded_file, caption='Uploaded Image.', use_column_width=True)
    st.write(f"Predicted Class: {predicted_class}")