nhosseini's picture
share=True
dfcc379 verified
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)