from ultralytics import YOLO
import gradio as gr
import numpy as np
from PIL import Image

# -------------------------
# Load detection model
# -------------------------
model = YOLO("buck_vs_doe_Detection_best.pt")

# -------------------------
# Inference function
# -------------------------
def predict(image):
    # Run inference (YOLO accepts numpy RGB directly)
    results = model(image)

    # Take first result (single image)
    r = results[0]

    # Plot results (BGR numpy array)
    im_bgr = r.plot()

    # Convert BGR → RGB for Gradio
    im_rgb = im_bgr[..., ::-1]

    return im_rgb


# -------------------------
# Gradio UI
# -------------------------
app = gr.Interface(
    fn=predict,
    inputs=gr.Image(type="numpy", label="Upload Image"),
    outputs=gr.Image(type="numpy", label="Detection Result"),
    title="Buck Tracker AI – Deer Detection",
    description="YOLO-based buck vs doe detection using Ultralytics native plotting."
)

# -------------------------
# Launch
# -------------------------
if __name__ == "__main__":
    app.launch()