openaiZ/Animal_Classification_app

1import os2import joblib3import numpy as np4import streamlit as st5from PIL import Image6import tensorflow as tf7from tensorflow.keras.applications import ResNet508from tensorflow.keras.applications.resnet50 import preprocess_input9from tensorflow.keras.preprocessing.image import img_to_array10 11# -----------------------------12# Page config13# -----------------------------14st.set_page_config(15    page_title="Animal Classification App",16    page_icon="🦁",17    layout="centered",18)19 20st.title("🐾 Animal Classification App")21st.write("Upload an animal image and the app will predict its class using your trained KNN model.")22 23MODEL_PATH = "knn_model.pkl"24CLASS_NAMES_PATH = "class_names.pkl"25 26 27# -----------------------------28# Cached loaders29# -----------------------------30@st.cache_resource31def load_feature_extractor():32    return ResNet50(weights="imagenet", include_top=False, pooling="avg")33 34 35@st.cache_resource36def load_knn_and_classes():37    if not os.path.exists(MODEL_PATH):38        raise FileNotFoundError(f"Missing file: {MODEL_PATH}")39    if not os.path.exists(CLASS_NAMES_PATH):40        raise FileNotFoundError(f"Missing file: {CLASS_NAMES_PATH}")41 42    knn_model = joblib.load(MODEL_PATH)43    class_names = joblib.load(CLASS_NAMES_PATH)44 45    return knn_model, class_names46 47 48# -----------------------------49# Image preprocessing50# -----------------------------51def prepare_image(uploaded_image: Image.Image) -> np.ndarray:52    image = uploaded_image.convert("RGB")53    image = image.resize((224, 224))54    arr = img_to_array(image)55    arr = np.expand_dims(arr, axis=0)56    arr = preprocess_input(arr)57    return arr58 59 60def extract_features(feature_model, image_array: np.ndarray) -> np.ndarray:61    features = feature_model.predict(image_array, verbose=0)62    return features63 64 65# -----------------------------66# Main app67# -----------------------------68try:69    feature_model = load_feature_extractor()70    knn_model, class_names = load_knn_and_classes()71 72    uploaded_file = st.file_uploader(73        "Upload an image",74        type=["jpg", "jpeg", "png", "webp"]75    )76 77    if uploaded_file is not None:78        image = Image.open(uploaded_file)79        st.image(image, caption="Uploaded image", use_container_width=True)80 81        if st.button("Predict"):82            with st.spinner("Processing image and predicting..."):83                image_array = prepare_image(image)84                features = extract_features(feature_model, image_array)85                pred_idx = int(knn_model.predict(features)[0])86 87                predicted_label = class_names[pred_idx]88                st.success(f"Predicted animal: **{predicted_label}**")89 90                # Optional: show nearest-neighbor info if available91                if hasattr(knn_model, "kneighbors"):92                    distances, indices = knn_model.kneighbors(features, n_neighbors=5)93                    st.write("Nearest-neighbor distances:")94                    st.write([float(x) for x in distances[0]])95 96except Exception as e:97    st.error(f"App failed to load: {e}")98    st.info(99        "Make sure these files are present in the Space root folder:\n"100        "- app.py\n"101        "- requirements.txt\n"102        "- knn_model.pkl\n"103        "- class_names.pkl"104    )