import gradio as gr
from ultralytics import YOLO
from PIL import Image
import numpy as np
import cv2

# Load the pre-trained YOLOv5 model
# Use 'yolov5s' or any available YOLO model from ultralytics
model = YOLO('yolov5s')  # Change to a leaf disease detection model if available

def identify_disease(image):
    # Convert the image to RGB if it's not
    if image.mode != 'RGB':
        image = image.convert('RGB')
    
    # Perform inference
    results = model(image)
    
    # Extract predictions
    predictions = results[0]
    boxes = predictions.boxes
    labels = boxes.cls.cpu().numpy()
    scores = boxes.conf.cpu().numpy()
    class_names = model.names
    
    # Annotate image with bounding boxes and labels
    annotated_image = np.array(image)
    for box, label, score in zip(boxes.xyxy.cpu().numpy(), labels, scores):
        x1, y1, x2, y2 = map(int, box)
        class_name = class_names[int(label)]
        confidence = f"{score * 100:.2f}%"
        annotated_image = cv2.rectangle(annotated_image, (x1, y1), (x2, y2), (0, 255, 0), 2)
        annotated_image = cv2.putText(annotated_image, f"{class_name} {confidence}", (x1, y1 - 10),
                                      cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 0), 2)
    
    # Convert annotated image back to PIL format
    annotated_image = Image.fromarray(annotated_image)
    
    # Prepare results for display
    results_list = [{"Disease": class_names[int(label)], "Confidence": f"{score * 100:.2f}%"} for label, score in zip(labels, scores)]
    
    return annotated_image, results_list

# Define Gradio interface
interface = gr.Interface(
    fn=identify_disease,
    inputs=gr.inputs.Image(type="pil"),
    outputs=[
        gr.outputs.Image(type="pil", label="Annotated Image"),
        gr.outputs.Dataframe(headers=["Disease", "Confidence"], label="Predictions")
    ],
    title="Leaf Disease Identifier with YOLOv5",
    description="Upload an image of a leaf, and this tool will identify the disease with confidence scores."
)

# Launch the app
interface.launch()