Spaces:
Running
Running
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]]}") | |