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

# Load the YOLOv8 model (replace 'best.pt' with the path to your model)
model = YOLO('best.pt')
# Function to perform object detection using YOLOv8
def detect_objects(image):
    # Convert the image from PIL format (used by Gradio) to a NumPy array
    image = np.array(image)

    # Convert the image to RGB (OpenCV loads images as BGR)
    image_rgb = cv2.cvtColor(image, cv2.COLOR_RGB2BGR)  # Convert from RGB (PIL) to BGR (OpenCV)

    # Perform inference on the image
    results = model.predict(image_rgb)

    # Get bounding boxes and class labels from the results
    annotated_image = results[0].plot()  # YOLOv8 has a plot method that returns an annotated image
    
    return annotated_image

# Gradio Interface (image upload instead of webcam)
demo = gr.Interface(
    fn=detect_objects,                    # The function to perform object detection
    inputs=gr.Image(type="pil", label="Upload an image"),  # Accept a file upload
    outputs=gr.Image(label="Detected Shelves"),
    title="Empty Shelf Detection using YOLOv8",
    description="Upload an image of shelves, and the model will identify empty spaces using YOLOv8. This tool is designed for recognizing empty shelves in retail or warehouse environments to assist with inventory management and restocking decisions."
)

# Start the app
if __name__ == "__main__":
    demo.launch(share=True)