Spaces:
Sleeping
Sleeping
import gradio as gr | |
import torch | |
from transformers import AutoModelForImageClassification, AutoImageProcessor | |
from PIL import Image | |
# Load model and processor | |
model = AutoModelForImageClassification.from_pretrained("jazzmacedo/fruits-and-vegetables-detector-36") | |
processor = AutoImageProcessor.from_pretrained("jazzmacedo/fruits-and-vegetables-detector-36") | |
labels = list(model.config.id2label.values()) | |
def classify_image(image): | |
# Preprocess the image | |
inputs = processor(images=image, return_tensors="pt") | |
with torch.no_grad(): | |
outputs = model(**inputs) | |
predicted_idx = torch.argmax(outputs.logits, dim=1).item() | |
predicted_label = labels[predicted_idx] | |
return predicted_label | |
# Create Gradio interface | |
interface = gr.Interface( | |
fn=classify_image, | |
inputs=gr.Image(type="pil"), | |
outputs=gr.Text(label="Detected Label"), | |
title="Fruit & Vegetable Detector", | |
description="Upload an image of a fruit or vegetable and get the predicted label." | |
) | |
if __name__ == "__main__": | |
interface.launch() | |