File size: 1,632 Bytes
eb4c49e 390c589 eb4c49e 390c589 eb4c49e 390c589 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 |
import random
import gradio as gr
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")
def classify_image(img: Image.Image):
label, confidence, probs = predict(img)
return (
label,
round(confidence, 3),
{k: round(v, 3) for k, v in probs.items()}
)
# 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
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() |