import gradio as gr
from transformers import ViTForImageClassification, ViTImageProcessor
import torch

# Define the pretrained model
model_name = "Treelar/vit-b16-plant_village"

# Load the ViT model and the image processor
model = ViTForImageClassification.from_pretrained(model_name)
image_processor = ViTImageProcessor.from_pretrained(model_name)

def predict(image):
    # Convert the image into the model's required format
    inputs = image_processor(images=image, return_tensors="pt")

    # Disable gradient calculation to make the process efficient
    with torch.no_grad():
        outputs = model(**inputs)  # Gets the model's output for the image
        logits = outputs.logits   # Output scores from the model

    # Convert the logits to a probability using the softmax function
    probability = torch.nn.functional.softmax(logits, dim=1)
    top_probability, top_index = probability.max(1)  # Gets the highest probability with its respective index

    # Gets the disease label from the model using the probability's index
    label_index = top_index.item()
    label = model.config.id2label[label_index]

    # Split the label into leaf category and disease name
    label_parts = label.split("___")
    leaf_category = label_parts[0]
    disease_name = label_parts[1]

    # Calculate the percentage breakdown of predicted diseases
    percentage_breakdown = {disease: round(float(probability[0, index]) * 100, 2) for index, disease in enumerate(model.config.label2id)}

    return leaf_category.capitalize(), disease_name.replace('_', ' ').capitalize(), percentage_breakdown

# Gradio interface setup with separate boxes for Leaf Type, Identified Disease, and Percentage Breakdown
interface = gr.Interface(
    fn=predict,
    inputs=gr.Image(label="Upload the Image"),
    outputs=[
        gr.Textbox(label="Leaf Type:"),
        gr.Textbox(label="Identified Disease:"),
        gr.JSON(label="Percentage Breakdown:")
    ],
    title="Plant Disease Identifier",
    description="Once the image has been uploaded and submitted, the disease of the plant will be determined. The percentage breakdown shows the probability of each disease."
)

interface.launch(debug=True)  # Start server and launch the UI