Update app.py

app.py

@@ -4,7 +4,6 @@ from PIL import Image
 from model import predict
 from datasets import load_dataset
 
-# Streaming dataset (SAFE for large datasets)
 dataset = load_dataset(
     "AIOmarRehan/AnimalsDataset",
     split="train",
@@ -24,19 +23,16 @@ def classify_image(img: Image.Image):
         {k: round(v, 3) for k, v in probs.items()}
     )
 
-# Random example for streaming dataset
 def random_example():
 
-    # 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
+    return img, img, label_str
 
 
 demo = gr.Blocks()
@@ -65,7 +61,7 @@ with demo:
 
     rand_img.click(
         random_example,
-        outputs=[rand_display, rand_label]
+        outputs=[input_img, rand_display, rand_label]
     )
 
 if __name__ == "__main__":