Spaces:
Runtime error
Runtime error
import gradio as gr | |
from PIL import Image | |
import tempfile | |
import torch | |
from torchvision.io import read_image | |
from transformers import ViTForImageClassification, ViTFeatureExtractor,ViTImageProcessor | |
# With ViTImageProcessor we have error so i comment it | |
# model = ViTImageProcessor.from_pretrained('SeyedAli/Food-Image-Classification-VIT') | |
model = ViTForImageClassification.from_pretrained('SeyedAli/Food-Image-Classification-VIT') | |
feature_extractor = ViTFeatureExtractor.from_pretrained('SeyedAli/Food-Image-Classification-VIT') | |
def FoodClassification(image): | |
with tempfile.NamedTemporaryFile(suffix=".png") as temp_image_file: | |
# Copy the contents of the uploaded image file to the temporary file | |
Image.fromarray(image).save(temp_image_file.name) | |
# Load the image file using torchvision | |
image = read_image(temp_image_file.name) | |
# Preprocess the image using the ViT feature extractor | |
inputs = feature_extractor(images=image, return_tensors="pt") | |
# Use the ViT model for image classification | |
outputs = model(**inputs) | |
predicted_class_idx = torch.argmax(outputs.logits) | |
predicted_class = model.config.id2label[predicted_class_idx.item()] | |
return predicted_class | |
iface = gr.Interface(fn=FoodClassification, inputs="image", outputs="text") | |
iface.launch(share=False) |