Update app.py

app.py

@@ -1,40 +1,51 @@
-import torch
-from ultralytics.nn.tasks import DetectionModel
-from torch.nn.modules.container import Sequential
-from ultralytics.nn.modules import Conv  # Import Conv
-
-# Whitelist safe globals for unpickling (only if you trust the model checkpoint!)
-torch.serialization.add_safe_globals([DetectionModel, Sequential, Conv])
-
-from huggingface_hub import hf_hub_download
-from ultralytics import YOLO
 import gradio as gr
-from PIL import Image
+from ultralyticsplus import YOLO, render_result
+import cv2
 
-# Download the weights file (ensure the filename is correct)
-weights_path = hf_hub_download(
-    repo_id="foduucom/plant-leaf-detection-and-classification",
-    filename="best.pt"
-)
+# Load the YOLO model
+model = YOLO('foduucom/plant-leaf-detection-and-classification')
 
-# Load the model using the local weights file
-model = YOLO(weights_path)
+# 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 predict_leaves(image_path):
-    # Run prediction on the provided image
-    results = model.predict(source=image_path, save=True)
-    # Count the number of detected leaves (using the first result)
-    count = len(results[0].boxes)
-    return f"Detected leaves: {count}"
+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 a Gradio interface for your app
-iface = gr.Interface(
-    fn=predict_leaves,
-    inputs=gr.Image(type="filepath"),
-    outputs="text",
-    title="Leaf Detection & Classification",
-    description="Upload an image to detect and count leaves using the YOLO model."
-)
+# 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__":
-    iface.launch()
+    demo.launch(server_port=7860, share=False)