Update app.py

app.py

@@ -1,28 +1,30 @@
 import gradio as gr
 import tensorflow as tf
 import numpy as np
-from tensorflow.keras.applications.mobilenet_v2 import preprocess_input
+import json
 from PIL import Image
-import os
 
-# Load the trained model (make sure the file is in the same directory)
+# Load model and class names from local repo files
 model = tf.keras.models.load_model("animal_classifier.keras")
 
-# Define class names (same order as training)
-train_dir = "split_animals/train"
-class_names = sorted(os.listdir(train_dir))
+with open("class_names.json", "r") as f:
+    class_names = json.load(f)
 
-# Prediction function
-def predict_image(img):
-    img = img.resize((224, 224))
-    img_array = tf.keras.utils.img_to_array(img)
+def predict(image: Image.Image):
+    image = image.resize((224, 224))
+    img_array = np.array(image)
     img_array = np.expand_dims(img_array, axis=0)
-    img_array = preprocess_input(img_array)
+    img_array = img_array / 255.0  # normalize if used in training
 
-    prediction = model.predict(img_array)
-    predicted_class = class_names[np.argmax(prediction)]
-    return f"Predicted Animal: {predicted_class}"
+    preds = model.predict(img_array)
+    pred_class = np.argmax(preds, axis=1)[0]
+    return class_names[pred_class]
 
-# Launch Gradio app
-app = gr.Interface(fn=predict_image, inputs=gr.Image(type="pil"), outputs="text", title="🧠🐾 TF-WildNet: Animal Classifier")
-app.launch()
+demo = gr.Interface(fn=predict, 
+                    inputs=gr.Image(type="pil"), 
+                    outputs="text",
+                    title="MobileNetV2 Animal Classifier",
+                    description="Upload an image to classify the animal.")
+
+if __name__ == "__main__":
+    demo.launch()