KabeerAmjad's picture
Update app.py
9f1c117 verified
raw
history blame
1.22 kB
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):
# Preprocess the image and extract features
inputs = feature_extractor(images=img, return_tensors="pt")
# Run the model in evaluation mode
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"), # Image input type as PIL (Pillow Image)
outputs="text", # Text output will display the predicted label
title="Food Image Classification",
description="Upload an image to classify if it’s an apple pie, etc."
)
# Launch the app
iface.launch()