shah1zil's picture
Update app.py
380eb1f verified
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}")