import gradio as gr
from transformers import AutoImageProcessor, AutoModelForImageClassification
import torch
from PIL import Image

# Load the image processor and model
processor = AutoImageProcessor.from_pretrained("ozair23/swin-tiny-patch4-window7-224-finetuned-plantdisease")
model = AutoModelForImageClassification.from_pretrained("ozair23/swin-tiny-patch4-window7-224-finetuned-plantdisease")

# Define the function to process the image and make predictions
def classify_leaf_disease(image):
    # Preprocess the image
    inputs = processor(images=image, return_tensors="pt")
    
    # Run the model on the image
    with torch.no_grad():
        outputs = model(**inputs)
    
    # Get the predicted label and confidence score
    probs = torch.softmax(outputs.logits, dim=1)
    predicted_class_idx = probs.argmax(dim=1).item()
    predicted_label = model.config.id2label[predicted_class_idx]
    confidence_score = probs[0][predicted_class_idx].item()
    
    # Format the output
    return predicted_label, f"{confidence_score:.2f}", image

# Create Gradio Interface
interface = gr.Interface(
    fn=classify_leaf_disease,
    inputs=gr.Image(type="pil"),
    outputs=[
        gr.Textbox(label="Disease Name"),
        gr.Textbox(label="Confidence Score"),
        gr.Image(type="pil", label="Uploaded Image")
    ],
    title="Leaf Disease Identification",
    description="Upload an image of any plant leaf, and this model will identify the disease and show the confidence score."
)

# Launch the app
if __name__ == "__main__":
    interface.launch()