Update app.py

app.py

@@ -2,7 +2,7 @@ import streamlit as st
 import joblib
 import numpy as np
 from tensorflow.keras.preprocessing import image
-from tensorflow.keras.applications.resnet50 import preprocess_input
+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
@@ -11,9 +11,11 @@ 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 (without the top layer)
+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
@@ -23,11 +25,13 @@ if uploaded_file is not None:
     # Process the image
     img = load_img(uploaded_file, target_size=(224, 224))  # Resize image to match model input
     img = img_to_array(img)  # Convert to array
+    img = np.expand_dims(img, axis=0)  # Add batch dimension
     img = preprocess_input(img)  # Preprocess image for ResNet50
 
-    # Make prediction
-    img = np.expand_dims(img, axis=0)  # Add batch dimension
-    features = model.predict(img)  # Extract features using the model
+    # Extract features using ResNet50
+    features = resnet_model.predict(img)  # Extract features using ResNet50
+
+    # Make prediction with the KNN model
     prediction = model.predict(features)  # Get prediction
 
     # Show the result