app.py

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
from torch.nn.modules.conv import Conv2d  # Import Conv2d

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

from ultralyticsplus import YOLO, render_result
from PIL import Image

# Load the YOLOv8s Leaf Detection and Classification model from Hugging Face
model = YOLO('foduucom/plant-leaf-detection-and-classification')

# Set recommended model parameters as per the model card
model.overrides['conf'] = 0.25       # NMS confidence threshold
model.overrides['iou'] = 0.45        # NMS IoU threshold
model.overrides['agnostic_nms'] = False  # NMS class-agnostic setting
model.overrides['max_det'] = 1000    # Maximum detections per image

def count_leaves(image):
    # Convert the input to a PIL image (ensuring RGB)
    image = Image.open(image).convert("RGB")
    
    # Perform inference with the model
    results = model.predict(image)
    
    # Count the detected leaves using the bounding boxes from the first result
    num_leaves = len(results[0].boxes)
    
    return f"Number of leaves detected: {num_leaves}"

# Build a Gradio interface for the leaf counter
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()