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

dataset = datasets.load_dataset("beans")

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

labels = dataset['train'].features['labels'].names

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.""",
                         examples = example_imgs)

interface.launch(debug=True)