Update app.py

app.py

@@ -1,95 +1,38 @@
-import os
-import numpy as np
-import joblib
-from PIL import Image
 import streamlit as st
+import joblib
+import numpy as np
 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)
+from tensorflow.keras.preprocessing import image
+import os
+from PIL import Image
 
-# 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
+# Load the pre-trained KNN model and class names
+knn = joblib.load('knn_model.pk1')
+class_names = joblib.load('class_names.pk1')
 
-# 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
+# Title of the app
+st.title("Animal Classification Using KNN Model")
 
-# 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
+# Description
+st.write("Upload an image of an animal and the model will predict which animal it is.")
 
-# Main Streamlit app
-def main():
-    st.title("Animal Classifier")
-    st.write("Upload an image of an animal to predict its category.")
+# Upload image
+uploaded_image = st.file_uploader("Choose an image...", type=["jpg", "jpeg", "png"])
 
-    # 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
+if uploaded_image is not None:
+    # Display image
+    img = Image.open(uploaded_image)
+    st.image(img, caption='Uploaded Image.', use_column_width=True)
 
-    knn = joblib.load(MODEL_PATH)
-    class_names = joblib.load(CLASS_NAMES_PATH)
+    # 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
 
-    # 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}")
+    # Make prediction
+    prediction = knn.predict(img_array)
+    predicted_class = class_names[prediction[0]]
 
-# Train and save model (Uncomment if running for the first time)
-# train_and_save_model(DATASET_PATH, TARGET_SIZE, MODEL_PATH, CLASS_NAMES_PATH)
+    # Display prediction
+    st.write(f"Prediction: {predicted_class}")
 
-# Run Streamlit app
-if __name__ == "__main__":
-    main()