import gradio as gr
from transformers import AutoModelForImageClassification, AutoImageProcessor
import torch
import numpy as np

examples = [
  "shrimp.png",
  "adverarial.png"
]

hugging_face_model = "Kaludi/food-category-classification-v2.0"
model = AutoModelForImageClassification.from_pretrained(hugging_face_model)
processor = AutoImageProcessor.from_pretrained(hugging_face_model)

def predict(img):
  inputs = processor(images=img, return_tensors="pt")
  outputs = model(**inputs)
  logits = outputs.logits


  # ChatGPT Code: I have no idea what is going on
  probabilities = torch.softmax(logits, dim=1)[0].tolist()
  labels = model.config.id2label
  top_10_indices = np.argsort(probabilities)[::-1][:10]
  top_10_labels = [labels[i] for i in top_10_indices]
  top_10_probabilities = [probabilities[i] for i in top_10_indices]
  label_confidences = {label: prob for label, prob in zip(top_10_labels, top_10_probabilities)}
  return label_confidences

demo = gr.Interface(
  fn=predict,
  inputs=gr.Image(),
  outputs=gr.Label(),
  examples=examples
)

demo.launch()