File size: 1,038 Bytes
9f320da |
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 |
import gradio as gr
from transformers import AutoModelForImageClassification, AutoFeatureExtractor
from PIL import Image
import torch
# Load your Hugging Face model
model_id = "KabeerAmjad/food_classification_model" # Replace with your actual model ID
model = AutoModelForImageClassification.from_pretrained(model_id)
feature_extractor = AutoFeatureExtractor.from_pretrained(model_id)
# Define the prediction function
def classify_image(img):
inputs = feature_extractor(images=img, return_tensors="pt")
with torch.no_grad():
outputs = model(**inputs)
probs = torch.softmax(outputs.logits, dim=-1)
# Get the label with the highest probability
top_label = model.config.id2label[probs.argmax().item()]
return top_label
# Create the Gradio interface
iface = gr.Interface(
fn=classify_image,
inputs=gr.Image(type="pil"),
outputs="text",
title="Food Image Classification",
description="Upload an image to classify if it’s an apple pie, etc."
)
# Launch the app
iface.launch()
|