Update app.py

app.py

@@ -4,11 +4,17 @@ from PIL import Image
 from model import predict
 from datasets import load_dataset
 
-# Load the HF dataset once
-dataset = load_dataset("AIOmarRehan/AnimalsDataset", split="train", streaming=True)
+# Load dataset (NO streaming → allows len() and indexing)
+dataset = load_dataset("AIOmarRehan/AnimalsDataset", split="train")
 
 def classify_image(img: Image.Image):
+
+    # Handle empty input safely
+    if img is None:
+        return "No image uploaded", 0, {}
+
     label, confidence, probs = predict(img)
+
     return (
         label,
         round(confidence, 3),
@@ -17,14 +23,17 @@ def classify_image(img: Image.Image):
 
 # Pick a random example
 def random_example():
-    # Choose a random row
+
     idx = random.randint(0, len(dataset) - 1)
     item = dataset[idx]
-    img = item["image"].convert("RGB")  # PIL Image
-    label = item["label"]  # numeric label from dataset
-    label_str = dataset.features["label"].int2str(label)  # class name
+
+    img = item["image"].convert("RGB")
+    label = item["label"]
+    label_str = dataset.features["label"].int2str(label)
+
     return img, label_str
 
+
 demo = gr.Blocks()
 
 with demo:
@@ -46,8 +55,16 @@ with demo:
     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])
+    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()