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

# Whitelist safe globals for unpickling (only if you trust the model checkpoint!)
torch.serialization.add_safe_globals([DetectionModel, Sequential, Conv])

from huggingface_hub import hf_hub_download
from ultralytics import YOLO
import gradio as gr
from PIL import Image

# Download the weights file (ensure the filename is correct)
weights_path = hf_hub_download(
    repo_id="foduucom/plant-leaf-detection-and-classification",
    filename="best.pt"
)

# Load the model using the local weights file
model = YOLO(weights_path)

def predict_leaves(image_path):
    # Run prediction on the provided image
    results = model.predict(source=image_path, save=True)
    # Count the number of detected leaves (using the first result)
    count = len(results[0].boxes)
    return f"Detected leaves: {count}"

# Create a Gradio interface for your app
iface = gr.Interface(
    fn=predict_leaves,
    inputs=gr.Image(type="filepath"),
    outputs="text",
    title="Leaf Detection & Classification",
    description="Upload an image to detect and count leaves using the YOLO model."
)

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