| 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() |