import datasets
import torch
from transformers import AutoFeatureExtractor, AutoModelForImageClassification
dataset = 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 = image_processor(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

import gradio as gr

interface = gr.Interface(
   fn = classify,
   inputs = "image",
   outputs = "label",
   interpretation = "default",
   # interpretation ="shap", Shapley didn't work for me but default does
   # num_shap = 5,
   title= "Bean Image Classifier",
   description =  "A simple image classifier for bean diseases. Upload an image of a bean leaf to get started.",
   examples = [
       ["https://media.istockphoto.com/id/472954806/photo/bean-leaf-heart-shape.jpg?s=170667a&w=0&k=20&c=es-jmKQSZLwKLU8NtVZ8KyBVoMNx6rHhW7NBw93EqJw="],
       ["https://d3qz1qhhp9wxfa.cloudfront.net/growingproduce/wp-content/uploads/2020/11/common_bacterial_blight_of_beans_featured.jpg"],
       ["https://3.bp.blogspot.com/-1FSsbPueH5Y/U8IyG9LY3VI/AAAAAAAADu0/Y5HfcKuJ5-w/s1600/8040277776_036e43f2fd_z.jpg"],
       ["https://plantwiseplusknowledgebank.org/cms/10.1079/pwkb.species.40010/asset/dbf6c9b1-c370-4a46-85db-4da543a13e98/assets/graphic/angular%20leaf%20spot%20(phaeoisariopsis%20griseola)%20on%20pole%20bean%20%20leaves.jpg"]
      
       ]

)

interface.launch(debug=True)