leaf-counter / app.py
muskangoyal06's picture
Update app.py
7f97dd6 verified
raw
history blame
1.68 kB
import gradio as gr
from ultralyticsplus import YOLO, render_result
import cv2
# Load the YOLO model
model = YOLO('foduucom/plant-leaf-detection-and-classification')
# Set 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
model.overrides['max_det'] = 1000 # maximum detections per image
def detect_leaves(image):
# Convert from Gradio's numpy array to image file
image = cv2.cvtColor(image, cv2.COLOR_RGB2BGR)
cv2.imwrite('temp_image.jpg', image)
# Perform prediction
results = model.predict('temp_image.jpg')
# Count number of leaves detected
num_leaves = len(results[0].boxes)
# Render results on image
render = render_result(model=model, image='temp_image.jpg', result=results[0])
return render, num_leaves
# Create Gradio interface
with gr.Blocks(title="Leaf Detection & Classification") as demo:
gr.Markdown("# πŸƒ Plant Leaf Detection & Classification")
gr.Markdown("Upload a plant image to detect and count leaves")
with gr.Row():
with gr.Column():
input_image = gr.Image(label="Upload Plant Image")
submit_btn = gr.Button("Detect Leaves")
with gr.Column():
output_image = gr.Image(label="Detection Results")
leaf_count = gr.Number(label="Number of Leaves Detected")
submit_btn.click(
fn=detect_leaves,
inputs=[input_image],
outputs=[output_image, leaf_count]
)
if __name__ == "__main__":
demo.launch(server_port=7860, share=False)