import gradio as gr
import torch
from ultralytics.nn.tasks import DetectionModel
from torch.nn.modules.container import Sequential
from ultralytics.nn.modules import Conv  # Import Conv

# Whitelist the globals to bypass the pickle errors (only do this if you trust the model source!)
torch.serialization.add_safe_globals([DetectionModel, Sequential, Conv])

from ultralyticsplus import YOLO
from PIL import Image

# Load your custom YOLOv8 leaf detection model
model = YOLO('foduucom/plant-leaf-detection-and-classification')

def count_leaves(image):
    # Convert image to a PIL Image and ensure it's in RGB
    image = Image.open(image).convert("RGB")
    # Run inference
    results = model.predict(image)
    # Count the number of detected leaves
    num_leaves = len(results[0].boxes)
    return f"Number of leaves detected: {num_leaves}"

# Gradio UI
iface = gr.Interface(
    fn=count_leaves,
    inputs=gr.Image(type="filepath"),
    outputs="text",
    title="Leaf Counter",
    description="Upload an image of a plant, and the model will detect and count the number of leaves."
)

if __name__ == "__main__":
    iface.launch()