Update app.py

app.py

@@ -1,40 +1,44 @@
 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
-from tensorflow.keras.applications import VGG16
-from tensorflow.keras.applications.vgg16 import preprocess_input
 
-# Load model and class labels
+# Load KNN model and class labels
 knn_model = joblib.load("knn_model.pkl")
 class_labels = joblib.load("class_labels.pkl")
 
-# Load VGG16 feature extractor (same as used in training)
-feature_extractor = VGG16(weights="imagenet", include_top=False, input_shape=(224, 224, 3))
+# 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):
-    """Manually reshape the image into the expected feature vector size."""
-    image = cv2.resize(image, (64, 32))  # Resize to the shape used in training
-    features = image.flatten().reshape(1, -1)  # Flatten to match model input size
-    return features
+    """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):
-    """Preprocess the image and make a prediction."""
+    """Predict the class of the input image."""
     processed_image = extract_features(image)  # Extract features
-    prediction = knn_model.predict(processed_image)[0]  # Get predicted index
-    predicted_label = list(class_labels.keys())[list(class_labels.values()).index(prediction)]
-    return predicted_label
-
-st.title("Animal Image Classifier")
-
-uploaded_file = st.file_uploader("Upload an image...", type=["jpg", "png", "jpeg"])
-
-if uploaded_file is not None:
-    image = np.array(cv2.imdecode(np.frombuffer(uploaded_file.read(), np.uint8), cv2.IMREAD_COLOR))
-
-    # Display the image
-    st.image(image, caption="Uploaded Image", use_container_width=True)
-
-    if st.button("Identify"):
+    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}**")
+        st.success(f"🎯 Predicted Animal: {prediction}")