Update app.py

app.py

@@ -1,54 +1,42 @@
 import streamlit as st
-import pickle
-import cv2
 import numpy as np
-from skimage.transform import resize
-from PIL import Image
-
-import pickle
-
-try:
-    with open("knn_model.pkl", "rb") as f:
-        knn_model = pickle.load(f)
-    print("Model loaded successfully!")
-except Exception as e:
-    print(f"Error loading model: {e}")
-
-# Load the trained KNN model
-with open("knn_model.pkl", "rb") as f:
-    knn_model = pickle.load(f)
-
-# Load class labels (fixing the filename)
-with open("class_labels.pkl", "rb") as f:
-    class_labels = pickle.load(f)  # Ensure it is a list or dictionary
-
-# Function to preprocess image
-def preprocess_image(image):
-    """Resizes and flattens the image for model prediction."""
-    image = np.array(image)  # Convert PIL image to numpy array
-    image = cv2.cvtColor(image, cv2.COLOR_RGB2BGR)  # Convert RGB to BGR
-    image_resized = resize(image, (224, 224))  # Resize to 224x224
-    image_flattened = image_resized.flatten().reshape(1, -1)  # Flatten
-    return image_flattened
-
-# Prediction Function
+import joblib
+import cv2
+from tensorflow.keras.applications import VGG16
+from tensorflow.keras.applications.vgg16 import preprocess_input
+
+# Load 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))
+
+def extract_features(image):
+    """Extracts features using the same model used during training (VGG16 in this case)."""
+    image = cv2.resize(image, (224, 224))  # Resize image to 224x224
+    image = np.expand_dims(image, axis=0)  # Add batch dimension
+    image = preprocess_input(image)  # Apply VGG16 preprocessing
+    features = feature_extractor.predict(image)
+    return features.flatten().reshape(1, -1)  # Reshape to match model input
+
 def predict_animal(image):
-    """Predicts the class of the uploaded image."""
-    processed_image = preprocess_image(image)
+    """Preprocess the image and make a prediction."""
+    processed_image = extract_features(image)  # Extract features
     prediction = knn_model.predict(processed_image)[0]  # Get predicted index
-    return class_labels[prediction]  # Get class name from labels
+    predicted_label = list(class_labels.keys())[list(class_labels.values()).index(prediction)]
+    return predicted_label
+
+st.title("Animal Image Classifier")
 
-# Streamlit UI
-st.title("🐾 Animal Image Classifier")
-st.write("Upload an image of an animal and click **Identify** to predict the species.")
+uploaded_file = st.file_uploader("Upload an image...", type=["jpg", "png", "jpeg"])
 
-# Upload Image
-uploaded_image = st.file_uploader("Choose 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))
 
-if uploaded_image is not None:
-    image = Image.open(uploaded_image)
-    st.image(image, caption="Uploaded Image", use_column_width=True)
+    # Display the image
+    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}**")