Spaces:
Running
Running
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 | |
# Whitelist safe globals (only do this if you trust the source of the model) | |
torch.serialization.add_safe_globals([DetectionModel, Sequential, Conv]) | |
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 | |
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() | |