import torch import gradio as gr from transformers import AutoModelForImageClassification, AutoFeatureExtractor model_id = "chanelcolgate/vit-base-patch16-224-finetuned-flower" labels = ["daisy", "dandelion", "roses", "sunflowers", "tulips"] def classify_image(image): model = AutoModelForImageClassification.from_pretrained(model_id) feature_extractor = AutoFeatureExtractor.from_pretrained(model_id) inp = feature_extractor(image, return_tensors="pt") outp = model(**inp) pred = torch.nn.functional.softmax(outp.logits, dim=-1) preds = pred[0].cpu().detach().numpy() confidence = {label: float(preds[i]) for i, label in enumerate(labels)} return confidence interface = gr.Interface( fn=classify_image, inputs="image", examples=["flower-1.jpeg", "flower-2.jpeg"], outputs="label", ).launch()