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

# Load YOLOv8 model
model = YOLO("best.pt")  # Ensure best.pt is in the same directory

# Preprocess and run inference
def predict(image):
    # Perform prediction
    results = model.predict(source=image, conf=0.5)
    
    # Annotate the image with bounding boxes
    annotated_image = results[0].plot()
    
    # Convert to PIL Image
    return Image.fromarray(annotated_image)

# Gradio interface
gr.Interface(
    fn=predict,
    inputs=gr.Image(type="pil"),
    outputs="image",
    title="Hippo or Rhino Detection",
    description="Upload an image for object detection with YOLOv8."
).launch()