jatin1233232's picture
Update app.py
3dd8f3d verified
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()