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]

    # Assume "background with leaves" as the category and "N/A" as the disease name if label_parts only has one part
    if len(label_parts) == 1:
        disease_name = "Not applicable"
    else:
        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)}

    # Sort the percentage breakdown dictionary by values in descending order
    sorted_percentage = dict(sorted(percentage_breakdown.items(), key=lambda item: item[1], reverse=True))

    # Get the top 5 predicted diseases
    top_5_percentage = {disease: percentage for disease, percentage in list(sorted_percentage.items())[:5]}

    # Format the percentage breakdown for printing (top 5 only)
    formatted_percentage = '\n'.join([f"{disease.replace('___', ' Leaf - ').replace('_', ' ').capitalize()}: {percentage}%" for disease, percentage in top_5_percentage.items()])

    # Include the message in the leaf_category textbox if leaf identification fails
    if disease_name == "Not applicable":
        leaf_category = "Unfortunately, we are having trouble identifying a singular leaf in this image. \n(Try uploading a different image with your leaf on a white backgound)"
        formatted_percentage = "Sorry we couldn't figure it out :("
        

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

# Gradio interface setup with separate boxes for Leaf Type, Identified Disease, and Percentage Breakdown
interface = gr.Interface(
                        theme = 'gradio/dark',
                        fn=predict,
                        inputs=gr.Image(label="Upload the Image"),
                        outputs=[
                            gr.Textbox(label="Leaf Type:"),  # Modified label to include the message
                            gr.Textbox(label="Identified Disease:"),
                            gr.Textbox(label="Percentage Breakdown (top 5):")
                        ],
                        title="Plant Disease Identifier",
                        description="Quick Tip: If the image results do not match, consider taking a picture of the leaf against a clean white background for improved accuracy and better identification of plant diseases."
                    )

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