Update model.py

model.py

@@ -1,32 +1,70 @@
-import tensorflow as tf
-import numpy as np
+import random
+import gradio as gr
 from PIL import Image
-import os
+from model import predict
+from datasets import load_dataset
 
-MODEL_PATH = os.path.join(
-    os.path.dirname(__file__),
-    "saved_model",
-    "Inception_V3_Animals_Classification.h5"
-)
+# Load dataset (NO streaming → allows len() and indexing)
+dataset = load_dataset("AIOmarRehan/AnimalsDataset", split="train")
 
-model = tf.keras.models.load_model(MODEL_PATH)
+def classify_image(img: Image.Image):
 
-CLASS_NAMES = ["Cat", "Dog", "Snake"]
+    # Handle empty input safely
+    if img is None:
+        return "No image uploaded", 0, {}
 
-def preprocess_image(img: Image.Image, target_size=(256, 256)):
-    img = img.convert("RGB")
-    img = img.resize(target_size)
-    img = np.array(img).astype("float32") / 255.0
-    img = np.expand_dims(img, axis=0)
-    return img
+    label, confidence, probs = predict(img)
 
-def predict(img: Image.Image):
-    input_tensor = preprocess_image(img)
-    preds = model.predict(input_tensor)[0]
+    return (
+        label,
+        round(confidence, 3),
+        {k: round(v, 3) for k, v in probs.items()}
+    )
 
-    class_idx = int(np.argmax(preds))
-    confidence = float(np.max(preds))
+# Pick a random example
+def random_example():
 
-    prob_dict = {CLASS_NAMES[i]: float(preds[i]) for i in range(len(CLASS_NAMES))}
+    idx = random.randint(0, len(dataset) - 1)
+    item = dataset[idx]
 
-    return CLASS_NAMES[class_idx], confidence, prob_dict
+    img = item["image"].convert("RGB")
+    label = item["label"]
+    label_str = dataset.features["label"].int2str(label)
+
+    return img, label_str
+
+
+demo = gr.Blocks()
+
+with demo:
+    gr.Markdown("## Animal Image Classifier with Random Dataset Samples")
+    
+    with gr.Row():
+        input_img = gr.Image(type="pil", label="Upload an image")
+        rand_img = gr.Button("Random Dataset Image")
+    
+    with gr.Row():
+        pred_btn = gr.Button("Predict")
+    
+    output_label = gr.Label(label="Predicted Class")
+    output_conf = gr.Number(label="Confidence")
+    output_probs = gr.JSON(label="All Probabilities")
+    
+    # Display random dataset sample
+    rand_display = gr.Image(type="pil", label="Random Dataset Sample")
+    rand_label = gr.Textbox(label="Sample Label")
+    
+    # Actions
+    pred_btn.click(
+        classify_image,
+        inputs=input_img,
+        outputs=[output_label, output_conf, output_probs]
+    )
+
+    rand_img.click(
+        random_example,
+        outputs=[rand_display, rand_label]
+    )
+
+if __name__ == "__main__":
+    demo.launch()