leaf-counter / app.py
muskangoyal06's picture
Update app.py
4e9dcdb verified
raw
history blame
1.55 kB
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()