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

# Dictionary containing disease names and remedies
disease_remedies = {
    "bacterial spot": "Remove infected plant debris, use copper-based fungicides.",
    "early blight": "Apply fungicides, practice crop rotation.",
    "healthy": "No action needed.",
    "late blight": "Remove infected plants, use fungicides.",
    "leaf miner": "Use insecticidal sprays, remove affected leaves.",
    "leaf mold": "Improve air circulation, use fungicides.",
    "mosaic virus": "Remove infected plants, control aphids.",
    "septoria": "Remove infected leaves, use fungicides.",
    "spider mites": "Use miticides, introduce beneficial insects.",
    "yellow leaf curl virus": "Remove infected plants, control whiteflies."
}

def custom_images(img, target_size):
    """
    Verifies the shape of the input image and reshapes it if necessary.

    Args:
        img (numpy.ndarray): The input image.
        target_size (tuple): The target size for the image.

    Returns:
        img (numpy.ndarray): The reshaped image.
    """
    img = cv2.resize(img, target_size)
    return img

def plant_disease_detect(img, model, original_size):
    """
    Performs plant disease detection on the given image using the YOLO model.

    Args:
        img (numpy.ndarray): The input image.
        model (YOLO): The YOLO model to use for detection.
        original_size (tuple): The original size of the input image.

    Returns:
        detect_img (numpy.ndarray): The image with the detected objects and their classes.
        classes (list): A list of the predicted class names.
    """
    # Pass the image through the detection model and get the result
    detect_result = model(img)

    # Plot the detections
    detect_img = detect_result[0].plot()
    detections = detect_result[0].boxes.data.tolist()
    classes = [model.names[int(detection[5])] for detection in detections]

    # Resize the output image to the original input size
    detect_img = cv2.resize(detect_img, original_size)

    return detect_img, classes

def detect_and_display(img, width, height):
    model = YOLO('best_95.pt')
    original_size = (img.shape[1], img.shape[0])
    resized_image = custom_images(img, (width, height))
    detect_img, classes = plant_disease_detect(resized_image, model, original_size)
    unique_classes = list(set(classes))
    class_table = [[cls, disease_remedies.get(cls.lower(), "No remedy available")] for cls in unique_classes]
    return detect_img, class_table

def image_to_numpy(image):
    """
    Converts an uploaded image to a numpy array.

    Args:
        image (PIL.Image): The uploaded image.

    Returns:
        np.ndarray: The image as a numpy array.
    """
    image_array = np.array(image)
    if image_array.ndim == 2:  # grayscale
        image_array = cv2.cvtColor(image_array, cv2.COLOR_GRAY2BGR)
    elif image_array.shape[2] == 4:  # RGBA
        image_array = cv2.cvtColor(image_array, cv2.COLOR_RGBA2BGR)
    return image_array

with gr.Blocks() as demo:
    gr.Markdown("# DRONE-INSPECTOR")
    gr.Markdown("Upload an image of a plant to detect diseases.")

    with gr.Row():
        with gr.Column():
            img_input = gr.Image(type="numpy", label="Upload Image")
            width_slider = gr.Slider(256, 1024, value=512, label="Select Image Width")
            height_slider = gr.Slider(256, 1024, value=512, label="Select Image Height")
            submit_btn = gr.Button("Submit")
        with gr.Column():
            img_output = gr.Image(label="Processed Image with Detections")
            class_table = gr.DataFrame(headers=["Disease", "Remedy"], label="Detected Diseases and Remedies")

    submit_btn.click(fn=detect_and_display, inputs=[img_input, width_slider, height_slider], outputs=[img_output, class_table])
    gr.Markdown("Author: Stéphane KPOVIESSI")
    gr.Markdown("LinkedIn: [Amiche](https://www.linkedin.com/in/stephanekpoviessi/)")

demo.launch()