leaf-counter / app.py
muskangoyal06's picture
Update app.py
2dd2d70 verified
raw
history blame
1.11 kB
import gradio as gr
import torch
from ultralytics.nn.tasks import DetectionModel
from torch.nn.modules.container import Sequential # Allow Sequential if needed
# Whitelist the globals to bypass the pickle error (only do this if you trust the model source!)
torch.serialization.add_safe_globals([DetectionModel, Sequential])
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()