File size: 987 Bytes
8e1d2f0 227f288 8e1d2f0 227f288 8e1d2f0 227f288 8e1d2f0 227f288 8e1d2f0 227f288 8e1d2f0 |
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 |
import gradio as gr
from transformers import AutoImageProcessor, AutoModelForImageClassification
import torch
from PIL import Image
MODEL_ID = "jaqen79/retail_images_classification_v1"
processor = AutoImageProcessor.from_pretrained(MODEL_ID)
model = AutoModelForImageClassification.from_pretrained(MODEL_ID)
def predict(image: Image.Image):
inputs = processor(images=image, return_tensors="pt")
with torch.no_grad():
outputs = model(**inputs)
probs = outputs.logits.softmax(dim=-1).tolist()[0]
labels = [model.config.id2label[i] for i in range(len(probs))]
return {labels[i]: float(probs[i]) for i in range(len(probs))}
demo = gr.Interface(
fn=predict,
inputs=gr.components.Image(type="pil"),
outputs=gr.components.Label(num_top_classes=5),
title="Retail Image classification using fine-tuned ViT",
description="Upload an image and the model returns the classes with probabilities."
)
if __name__ == "__main__":
demo.launch() |