mjolnir1122's picture
Update app.py
19f3b33 verified
Raw
History Blame Contribute Delete
1.86 kB
import streamlit as st
import numpy as np
import joblib
import tensorflow as tf
from tensorflow.keras.applications import ResNet50
from tensorflow.keras.applications.resnet50 import preprocess_input
from tensorflow.keras.preprocessing.image import img_to_array
import cv2
# Load KNN model and class labels
knn_model = joblib.load("knn_model.pkl")
class_labels = joblib.load("class_labels.pkl")
# Load ResNet50 model (same as used in training)
resnet_model = ResNet50(weights="imagenet", include_top=False, input_shape=(224, 224, 3))
resnet_model = tf.keras.Model(inputs=resnet_model.input, outputs=tf.keras.layers.GlobalAveragePooling2D()(resnet_model.output))
def extract_features(image):
"""Extract features from image using ResNet50 (same as training)."""
image = cv2.resize(image, (224, 224)) # Resize image
image = img_to_array(image)
image = np.expand_dims(image, axis=0)
image = preprocess_input(image)
features = resnet_model.predict(image) # Extract features
return features.reshape(1, -1) # Ensure shape matches KNN input
def predict_animal(image):
"""Predict the class of the input image."""
processed_image = extract_features(image) # Extract features
prediction = knn_model.predict(processed_image)[0] # Predict class index
return list(class_labels.keys())[list(class_labels.values()).index(prediction)] # Convert index to label
# Streamlit UI
st.title("🐾 Animal Image Classifier")
uploaded_file = st.file_uploader("📤 Upload an image", type=["jpg", "png", "jpeg"])
if uploaded_file:
image = cv2.imdecode(np.frombuffer(uploaded_file.read(), np.uint8), 1)
st.image(image, caption="📷 Uploaded Image", use_container_width=True)
if st.button("🔍 Identify"):
prediction = predict_animal(image)
st.success(f"🎯 Predicted Animal: {prediction}")