import datasets
import gradio as gr
from transformers import AutoFeatureExtractor, AutoModelForImageClassification
import torch

dataset = datasets.load_dataset("beans")

feature_extractor = AutoFeatureExtractor.from_pretrained("saved_model_files")
model = AutoModelForImageClassification.from_pretrained("saved_model_files")

labels = dataset['train'].features['labels'].names
example_imgs = ["example_0.jpg", "example_1.jpg","example_2.jpg"]

def classify(im):
  features = feature_extractor(im, return_tensors='pt')
  logits = model(features["pixel_values"])[-1]
  probability = torch.nn.functional.softmax(logits, dim=-1)
  probs = probability[0].detach().numpy()
  confidences = {label: float(probs[i]) for i, label in enumerate(labels)} 
  return confidences

interface = gr.Interface(fn = classify, 
                         inputs="image",
                         outputs = "label",
                         title = "Plant Leaf Disease Classifier",
                         description = """Below is a simple app to detect Angular Leaf Spot and Bean Rust diseases on leaves.
                          Data was annotated by experts from the National Crops Resources Research Institute (NaCRRI) 
                          in Uganda and collected by the Makerere AI research lab.
                          The model being used is a fine-tuned Vision Transformer, specifically beginning with [google/vit-base-patch16-224]
                          (https://huggingface.co/google/vit-base-patch16-224) and trained using the [beans](https://huggingface.co/datasets/beans) dataset.
                          This app was created in Abubakar Abid's 'Building End-to-End Vision Applications' course through CoRise.
                          """,
                         examples = example_imgs)

interface.launch(debug=True)