Update app.py

app.py

@@ -4,12 +4,15 @@ from PIL import Image
 from model import predict
 from datasets import load_dataset
 
-# Load dataset (NO streaming → allows len() and indexing)
-dataset = load_dataset("AIOmarRehan/AnimalsDataset", split="train")
+# Streaming dataset (SAFE for large datasets)
+dataset = load_dataset(
+    "AIOmarRehan/AnimalsDataset",
+    split="train",
+    streaming=True
+)
 
 def classify_image(img: Image.Image):
 
-    # Handle empty input safely
     if img is None:
         return "No image uploaded", 0, {}
 
@@ -21,14 +24,16 @@ def classify_image(img: Image.Image):
         {k: round(v, 3) for k, v in probs.items()}
     )
 
-# Pick a random example
+# Random example for streaming dataset
 def random_example():
 
-    idx = random.randint(0, len(dataset) - 1)
-    item = dataset[idx]
+    # Shuffle streaming buffer then take first item
+    item = next(iter(dataset.shuffle(buffer_size=100)))
 
     img = item["image"].convert("RGB")
     label = item["label"]
+
+    # streaming dataset keeps features
     label_str = dataset.features["label"].int2str(label)
 
     return img, label_str
@@ -43,18 +48,15 @@ with demo:
         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")
+    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,