Update app.py

app.py

@@ -1,30 +1,29 @@
 import gradio as gr
 import tensorflow as tf
 import numpy as np
-import json
 from PIL import Image
 
-# Load model and class names from local repo files
+# Load your model from Hugging Face repo (or local file if testing)
 model = tf.keras.models.load_model("animal_classifier.keras")
 
-with open("class_names.json", "r") as f:
-    class_names = json.load(f)
+# Dynamically get class names from model output layer
+class_names = list(model.class_names) if hasattr(model, 'class_names') else ["dog", "lion", "tiger"]  # fallback if not stored in model
 
-def predict(image: Image.Image):
-    image = image.resize((224, 224))
-    img_array = np.array(image)
+def predict_image(image):
+    img = image.resize((224, 224))
+    img_array = np.array(img) / 255.0
     img_array = np.expand_dims(img_array, axis=0)
-    img_array = img_array / 255.0  # normalize if used in training
 
     preds = model.predict(img_array)
-    pred_class = np.argmax(preds, axis=1)[0]
-    return class_names[pred_class]
+    confidence = np.max(preds)
+    predicted_index = np.argmax(preds)
 
-demo = gr.Interface(fn=predict, 
-                    inputs=gr.Image(type="pil"), 
-                    outputs="text",
-                    title="MobileNetV2 Animal Classifier",
-                    description="Upload an image to classify the animal.")
+    threshold = 0.5
+    if confidence < threshold:
+        return "Image not recognized as any animal in the dataset"
+    else:
+        return class_names[predicted_index]
 
-if __name__ == "__main__":
-    demo.launch()
+demo = gr.Interface(fn=predict_image, inputs=gr.Image(type="pil"), outputs="text", title="MobileNetV2 Animal Classifier")
+
+demo.launch()