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)