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

# Load the YOLOv8 model (ensure the model file is correct and accessible)
try:
    model = YOLO('yolov8n.pt')
    print("Model loaded successfully.")
except Exception as e:
    print(f"Error loading model: {e}")

def identify_disease(image):
    # Convert the image to RGB if it's not already in RGB
    try:
        if image.mode != 'RGB':
            image = image.convert('RGB')
        print("Image converted to RGB successfully.")
    except Exception as e:
        print(f"Error converting image to RGB: {e}")
        return image, [{"Disease": "Error", "Confidence": "N/A"}]
    
    # Perform inference
    try:
        results = model(image)
        predictions = results[0]
        print("Model inference completed successfully.")
    except Exception as e:
        print(f"Error during model inference: {e}")
        return image, [{"Disease": "Error during inference", "Confidence": "N/A"}]
    
    # Check for detections
    if len(predictions.boxes) == 0:
        print("No detections found.")
        annotated_image = np.array(image)
        cv2.putText(annotated_image, "No disease detected", (10, 30), 
                    cv2.FONT_HERSHEY_SIMPLEX, 1, (255, 0, 0), 2)
        annotated_image = Image.fromarray(annotated_image)
        return annotated_image, [{"Disease": "None", "Confidence": "N/A"}]
    
    # Extract predictions and annotate image
    try:
        boxes = predictions.boxes
        labels = boxes.cls.cpu().numpy() if boxes.cls is not None else []
        scores = boxes.conf.cpu().numpy() if boxes.conf is not None else []
        class_names = model.names
        
        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)
        annotated_image = Image.fromarray(annotated_image)
        print("Image annotation completed.")
        
        # Prepare results list for output
        results_list = [{"Disease": class_names[int(label)], "Confidence": f"{score * 100:.2f}%"} for label, score in zip(labels, scores)]
        return annotated_image, results_list
    
    except Exception as e:
        print(f"Error during annotation: {e}")
        return image, [{"Disease": "Error during annotation", "Confidence": "N/A"}]

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

# Launch the app
interface.launch()