File size: 1,661 Bytes
9f84078 | 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 | import gradio as gr
import torch
from transformers import AutoImageProcessor, AutoModelForImageClassification
from PIL import Image
# 1. Load a modern SOTA model (DINOv2)
# This model is faster and more accurate than the original ViT
model_name = "facebook/dinov2-base-imagenet1k-1-layer"
processor = AutoImageProcessor.from_pretrained(model_name)
model = AutoModelForImageClassification.from_pretrained(model_name)
def classify_image(img):
# 2. Pre-process image
inputs = processor(images=img, return_tensors="pt")
# 3. Inference
with torch.no_grad():
outputs = model(**inputs)
# 4. Convert logits to probabilities (0% to 100%)
logits = outputs.logits
probabilities = torch.nn.functional.softmax(logits, dim=-1)[0]
# 5. Extract top 5 results for a better UI
top5_prob, top5_indices = torch.topk(probabilities, 5)
# Create a dictionary of {Label: Probability} for Gradio's Label component
confidences = {
model.config.id2label[idx.item()]: float(prob)
for prob, idx in zip(top5_prob, top5_indices)
}
return confidences
# 6. Build a modern UI
demo = gr.Interface(
fn=classify_image,
inputs=gr.Image(type="pil", label="Upload Image"),
# gr.Label automatically creates a beautiful bar chart for probabilities
outputs=gr.Label(num_top_classes=5, label="Predictions"),
title="Next-Gen Image Classification",
description="Running on **Meta's DINOv2** foundation model. Upload any image to see the top 5 predicted categories.",
theme="soft"
)
if __name__ == "__main__":
demo.launch() |