Update app.py

app.py

@@ -4,37 +4,29 @@ from PIL import Image
 from model import predict
 from datasets import load_dataset
 
-dataset = load_dataset(
-    "AIOmarRehan/AnimalsDataset",
-    split="train",
-    streaming=True
-)
+# Load the full dataset (non-streaming)
+dataset = load_dataset("AIOmarRehan/AnimalsDataset", split="train")
 
 def classify_image(img: Image.Image):
-
     if img is None:
         return "No image uploaded", 0, {}
-
+    
     label, confidence, probs = predict(img)
-
     return (
         label,
         round(confidence, 3),
         {k: round(v, 3) for k, v in probs.items()}
     )
 
+# Random example from the dataset
 def random_example():
-
-    item = next(iter(dataset.shuffle(buffer_size=1500)))
-
+    item = random.choice(dataset)
     img = item["image"].convert("RGB")
-    label = item["label"]
-
-    label_str = dataset.features["label"].int2str(label)
-
-    return img, img, label_str
-
+    label = dataset.features["label"].int2str(item["label"])
+    # Return image twice: once for input_img (for prediction), once for display
+    return img, img, label
 
+# Gradio UI
 demo = gr.Blocks()
 
 with demo:
@@ -53,12 +45,14 @@ with demo:
     rand_display = gr.Image(type="pil", label="Random Dataset Sample")
     rand_label = gr.Textbox(label="Sample Label")
     
+    # Predict button uses whatever image is currently in input_img
     pred_btn.click(
         classify_image,
         inputs=input_img,
         outputs=[output_label, output_conf, output_probs]
     )
 
+    # Random button picks a dataset image
     rand_img.click(
         random_example,
         outputs=[input_img, rand_display, rand_label]