import gradio as gr
import cv2
import numpy as np
from ultralytics import YOLO

# Load the YOLOv8 model
model = YOLO("best.pt")  # Ensure this file is in the same directory as app.py on Hugging Face

# Define the inference function
def predict(image):
    # Convert the input image from RGB to BGR (OpenCV format)
    image_bgr = cv2.cvtColor(image, cv2.COLOR_RGB2BGR)
    
    # Run the model on the input image
    results = model(image_bgr)
    
    # Extract the result image with detections
    annotated_image = results[0].plot()  # Returns a BGR image with annotations
    
    # Convert the image back to RGB for displaying in Gradio
    annotated_image_rgb = cv2.cvtColor(annotated_image, cv2.COLOR_BGR2RGB)
    
    return annotated_image_rgb

# Define the Gradio interface
interface = gr.Interface(
    fn=predict,
    inputs=gr.Image(type="numpy", label="Upload an Image"),
    outputs=gr.Image(type="numpy", label="Detected Objects"),
    title="YOLOv8 Object Detection",
    description="Upload an image to detect objects with YOLOv8 model."
)

# Launch the app
if __name__ == "__main__":
    interface.launch(share=True)