Spaces:
Runtime error
Runtime error
File size: 2,053 Bytes
9637988 |
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 54 55 56 57 58 59 60 61 62 63 |
"""
This module provides an interface for classifying images using the ResNet-18 model.
The interface allows users to upload an image and receive the top 3 predicted labels.
"""
import gradio as gr
import spaces
import torch
from transformers import AutoImageProcessor, AutoModelForImageClassification
from datasets import load_dataset
from PIL import Image
# Load dataset and get test image
dataset = load_dataset("huggingface/cats-image", trust_remote_code=True)
test_image = dataset["test"]["image"][0]
# Initialize the image processor and model
image_processor = AutoImageProcessor.from_pretrained("microsoft/resnet-18")
model = AutoModelForImageClassification.from_pretrained("microsoft/resnet-18").to("cuda")
@spaces.GPU
def predict(image: Image, top_k: int = 3) -> dict:
"""
Predicts the top 'top_k' labels for an image using the ResNet-18 model.
Args:
image (Image): The input image as a PIL Image object.
top_k (int): The number of top predictions to return.
Returns:
dict: A dictionary with the top 'top_k' labels and their probabilities.
"""
inputs = image_processor(image, return_tensors="pt").to("cuda")
with torch.no_grad():
logits = model(**inputs).logits
# Apply softmax to logits to get probabilities
probabilities = torch.softmax(logits, dim=-1)
# Get the top 'top_k' probabilities and their corresponding indices
top_k_probs, top_k_indices = torch.topk(input=probabilities, k=top_k, dim=-1)
# Map the indices to labels and probabilities
predicted_labels = {
model.config.id2label[idx.item()]: prob.item()
for idx, prob in zip(top_k_indices[0], top_k_probs[0])
}
return predicted_labels
# Define the Gradio interface
demo = gr.Interface(
fn=predict,
inputs=gr.Image(type="pil"),
outputs=gr.Label(num_top_classes=3),
title="Classifying Images with ResNet-18",
description="Upload an image to predict the top 3 labels.",
examples=[test_image]
)
# Launch the Gradio interface
demo.launch()
|