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

image_processor = ViTImageProcessor.from_pretrained("google/vit-base-patch16-224")

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

labels = ['angular_leaf_spot', 'bean_rust', 'healthy'] 

def classify(image):
  features = image_processor(image, 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)} 
  print(confidences)
  return confidences


theme = gr.themes.Soft(
    primary_hue="green",
    secondary_hue="green",
    neutral_hue="green",
).set(
    block_background_fill_dark='*body_background_fill',
    button_border_width='*block_label_border_width',
    button_border_width_dark='*checkbox_label_border_width'
)

with gr.Blocks(theme=theme) as demo:

    inference = gr.Interface(fn=classify, inputs="image", outputs="label",
                        title="Plant leaves Classification",
                    description="Classify the leaves by uploading image",
                    examples=["images/1.png","images/2.png", "images/3.png"])

demo.launch()