Update app.py

app.py

@@ -1,38 +1,95 @@
-import streamlit as st
-import joblib
-import numpy as np
-from sklearn.neighbors import KNeighborsClassifier
-from tensorflow.keras.preprocessing import image
 import os
+import numpy as np
+import joblib
 from PIL import Image
+import streamlit as st
+from sklearn.neighbors import KNeighborsClassifier
+from sklearn.model_selection import train_test_split
+from sklearn.metrics import accuracy_score
+
+# Paths and global constants
+#DATASET_PATH = "path_to_your_dataset"  # Update with your dataset path
+MODEL_PATH = "knn_model.pkl"
+CLASS_NAMES_PATH = "class_names.pkl"
+TARGET_SIZE = (64, 64)
 
-# Load the pre-trained KNN model and class names
-knn = joblib.load('knn_model.pk1')
-class_names = joblib.load('class_names.pk1')
+# Helper function to load and preprocess dataset
+def load_dataset(dataset_path, target_size):
+    images = []
+    labels = []
+    class_names = sorted(os.listdir(dataset_path))  # Sort for consistent class indexing
+    for idx, class_name in enumerate(class_names):
+        class_path = os.path.join(dataset_path, class_name)
+        if not os.path.isdir(class_path):
+            continue
+        for image_name in os.listdir(class_path):
+            image_path = os.path.join(class_path, image_name)
+            try:
+                img = Image.open(image_path).convert('RGB')
+                img = img.resize(target_size)
+                img_array = np.array(img).flatten()  # Flatten to a single array
+                images.append(img_array)
+                labels.append(idx)
+            except Exception as e:
+                print(f"Error loading image {image_path}: {e}")
+    return np.array(images), np.array(labels), class_names
 
-# Title of the app
-st.title("Animal Classification Using KNN Model")
+# Function to train and save KNN model
+def train_and_save_model(dataset_path, target_size, model_path, class_names_path):
+    # Load data
+    X, y, class_names = load_dataset(dataset_path, target_size)
+    # Split dataset
+    X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42, stratify=y)
+    # Train KNN model
+    knn = KNeighborsClassifier(n_neighbors=3)
+    knn.fit(X_train, y_train)
+    # Evaluate model
+    y_pred = knn.predict(X_test)
+    accuracy = accuracy_score(y_test, y_pred)
+    print(f"KNN Model Accuracy: {accuracy * 100:.2f}%")
+    # Save model and class names
+    joblib.dump(knn, model_path)
+    joblib.dump(class_names, class_names_path)
+    print(f"Model and class names saved to {model_path} and {class_names_path}")
+    return knn, class_names
 
-# Description
-st.write("Upload an image of an animal and the model will predict which animal it is.")
+# Preprocess image for prediction
+def preprocess_image(image_path, target_size):
+    img = Image.open(image_path).convert('RGB')
+    img = img.resize(target_size)
+    img_array = np.array(img).flatten()
+    return img_array.reshape(1, -1)  # Add batch dimension
 
-# Upload image
-uploaded_image = st.file_uploader("Choose an image...", type=["jpg", "jpeg", "png"])
+# Main Streamlit app
+def main():
+    st.title("Animal Classifier")
+    st.write("Upload an image of an animal to predict its category.")
 
-if uploaded_image is not None:
-    # Display image
-    img = Image.open(uploaded_image)
-    st.image(img, caption='Uploaded Image.', use_column_width=True)
+    # Load model and class names
+    if not os.path.exists(MODEL_PATH) or not os.path.exists(CLASS_NAMES_PATH):
+        st.error("Model or class names file not found. Please train the model first.")
+        return
 
-    # Preprocess the image for prediction
-    img = img.resize((64, 64))  # Resize the image to match the model's expected size (adjust if needed)
-    img_array = np.array(img)  # Convert the image to numpy array
-    img_array = img_array.flatten().reshape(1, -1)  # Flatten the image and reshape it to match the input for KNN model
+    knn = joblib.load(MODEL_PATH)
+    class_names = joblib.load(CLASS_NAMES_PATH)
 
-    # Make prediction
-    prediction = knn.predict(img_array)
-    predicted_class = class_names[prediction[0]]
+    # Upload image
+    uploaded_file = st.file_uploader("Upload an image", type=["jpg", "jpeg", "png"])
+    if uploaded_file:
+        st.image(uploaded_file, caption="Uploaded Image", use_container_width=True)
+        try:
+            img_array = preprocess_image(uploaded_file, TARGET_SIZE)
+            if img_array.shape[1] != knn.n_features_in_:
+                st.error(f"Feature size mismatch: Model expects {knn.n_features_in_} features, but input has {img_array.shape[1]}.")
+            else:
+                prediction = knn.predict(img_array)
+                st.success(f"The predicted animal is: {class_names[prediction[0]]}")
+        except Exception as e:
+            st.error(f"Error processing image: {e}")
 
-    # Display prediction
-    st.write(f"Prediction: {predicted_class}")
+# Train and save model (Uncomment if running for the first time)
+# train_and_save_model(DATASET_PATH, TARGET_SIZE, MODEL_PATH, CLASS_NAMES_PATH)
 
+# Run Streamlit app
+if __name__ == "__main__":
+    main()