mjolnir1122's picture
Update app.py
19f3b33 verified
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}")