leaf-counter / app.py
muskangoyal06's picture
Update app.py
3b7e08f verified
raw
history blame
1.79 kB
import gradio as gr
from ultralyticsplus import YOLO, render_result
import cv2
import torch
# Verify torch version
print(f"Using torch version: {torch.__version__}")
# Load model with compatibility fix
def load_model():
try:
model = YOLO('foduucom/plant-leaf-detection-and-classification')
model.overrides['conf'] = 0.25
model.overrides['iou'] = 0.45
model.overrides['agnostic_nms'] = False
model.overrides['max_det'] = 1000
return model
except Exception as e:
raise RuntimeError("Error loading model. Please check the requirements versions.") from e
model = load_model()
def detect_leaves(image):
# Convert image format
image = cv2.cvtColor(image, cv2.COLOR_RGB2BGR)
cv2.imwrite('temp_image.jpg', image)
# Perform prediction
results = model.predict('temp_image.jpg')
# Process results
num_leaves = len(results[0].boxes)
render = render_result(model=model, image='temp_image.jpg', result=results[0])
return render, num_leaves
# Create Gradio interface
with gr.Blocks(theme=gr.themes.Soft(), title="Leaf Detection") as demo:
gr.Markdown("## πŸƒ Plant Leaf Detection & Counter")
gr.Markdown("Upload a plant image to analyze leaf count and species")
with gr.Row():
input_image = gr.Image(label="Input Image", type="numpy")
output_image = gr.Image(label="Detected Leaves", interactive=False)
leaf_count = gr.Number(label="Total Leaves Detected", precision=0)
submit_btn = gr.Button("Analyze Image", variant="primary")
submit_btn.click(
fn=detect_leaves,
inputs=[input_image],
outputs=[output_image, leaf_count]
)
if __name__ == "__main__":
demo.launch(server_port=7860, show_error=True)