import numpy as np
import gradio as gr
from transformers import AutoImageProcessor, AutoModelForImageClassification
from PIL import Image
import torch
import pandas as pd



processor = AutoImageProcessor.from_pretrained("Moreza009/HF_CVcourse_FoodClassifier")
model = AutoModelForImageClassification.from_pretrained("Moreza009/HF_CVcourse_FoodClassifier")

def classifier(image):
    if isinstance(image, np.ndarray):
        image = Image.fromarray(image)
    #image = Image.open(image)
    inputs = processor(images=image, return_tensors="pt")
    outputs = model(**inputs)
    logits = outputs.logits
    probabilities = torch.nn.functional.softmax(logits, dim=-1)
    predicted_class_idxs = probabilities.topk(5, dim=-1)[1].tolist()[0]
    probabilities = sorted(probabilities.tolist()[0],  reverse=True)[:5]
    classes = [model.config.id2label[idx] for idx in predicted_class_idxs]
    df = pd.DataFrame({'food':classes , 'posibility': probabilities})
    return df.to_html(index=False)


food = gr.Interface(
    fn=classifier,
    inputs=gr.Image(type="pil"),
    outputs="html",
    title = "what's your eating?",
    description = " :)  "
)


food.launch()