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

css = '''
.gradio-container {
  width: 85% !important;
}
h1 {
  width: 100% !important;
} 
p {
  margin-left: 30px !important; 
  margin-right: 30px !important;
  font-size: 1.1rem !important;
}
'''

title = r"""
<h1>GDSC: PwnAI Image Classifier Demo</h1>
"""
description = r"""
<p>PwnAI is an educational event that explores adversarial machine learning. It aims to help students learn about LLM Prompt Injection and Fooling Image Classifiers.</p>
<p>In this demo, there are two example pictures. They may look very similar to us, but they are classified differently. Try it out yourself by submitting both pictures for inference and observing the results.</p>
<p>Try it out yourself by experimenting with the image on the left to see if you can cause a misclassification.</p>
<h2>What to Expect</h2>
<p>Participants will explore these concepts through interactive demos and hands-on workshops, gaining insights into applying adversarial attacks to machine learning models.</p>
<h2>NOTE</h2>
<p>This demo runs on a free-tier CPU, so its performance is slow.</p>
<h2>What can this AI Classify</h2>
<p>1. Bread</p>
<p>2. Dairy</p>
<p>3. Dessert</p>
<p>4. Egg</p>
<p>5. Fried Food</p>
<p>6. Fruit</p>
<p>7. Meat</p>
<p>8. Noodles</p>
<p>9. Rice</p>
<p>10. Seafood</p>
<p>11. Soup</p>
<p>12. Vegetable</p>
"""

Ending = r"""
<p>Woahhhh. Is this steganography???? Who knows? Come for the event to find out!</p>
"""

Footer = r"""
---
Challenge Created By Chai Cheng Xun
"""

with gr.Blocks(css=css, title="PwnAI: Image Classifier Demo") as demo:
  gr.Markdown(title)
  gr.Markdown(description)

  with gr.Row():
    with gr.Column():
      img_file = gr.Image(label="Upload a photo to be classfied")
      submit = gr.Button("Submit", variant="primary")


    with gr.Column():
      output = gr.Label()

    submit.click(
      fn=predict,
      inputs=img_file,
      outputs=output,
    )


  gr.Examples(
    examples=examples,
    inputs=[img_file],
  )
  
  gr.Markdown(Ending)
  gr.Markdown(Footer)

demo.launch()