import gradio as gr
from ultralyticsplus import YOLO, render_result
import cv2
import torch
import ultralytics
import ultralyticsplus

# Check versions
print(f"Torch version: {torch.__version__}")
print(f"Ultralytics version: {ultralytics.__version__}")
print(f"UltralyticsPlus version: {ultralyticsplus.__version__}")

# Load model
model = YOLO('foduucom/plant-leaf-detection-and-classification')

# Model configuration
model.overrides['conf'] = 0.25  # Confidence threshold
model.overrides['iou'] = 0.45   # IoU threshold
model.overrides['agnostic_nms'] = False
model.overrides['max_det'] = 1000

def detect_leaves(image):
    # Convert image format
    img = cv2.cvtColor(image, cv2.COLOR_RGB2BGR)
    
    # Perform prediction
    results = model.predict(img)
    
    # Get results
    num_leaves = len(results[0].boxes)
    rendered_img = render_result(model=model, image=img, result=results[0])
    
    # Convert back to RGB for Gradio
    return cv2.cvtColor(rendered_img, cv2.COLOR_BGR2RGB), num_leaves

# Create Gradio interface
interface = gr.Interface(
    fn=detect_leaves,
    inputs=gr.Image(label="Upload Plant Image"),
    outputs=[
        gr.Image(label="Detected Leaves"),
        gr.Number(label="Number of Leaves Found")
    ],
    title="🍃 Plant Leaf Detection & Counting",
    description="Upload an image of a plant to detect and count its leaves"
)

if __name__ == "__main__":
    interface.launch(server_port=7860)