app.py

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

# Check for GPU availability
device = 'cuda' if torch.cuda.is_available() else 'cpu'
print(f"Using device: {device}")

# Load the YOLOv8 model from the 'model' directory
model = YOLO('best.pt')
model.to(device)

def detect_objects(image):
    """
    Performs object detection on the input image and returns the image with
    bounding boxes and labels drawn on it.
    """
    # Run inference on the image
    results = model(image)
    
    # The plot() method returns a BGR numpy array with detections
    annotated_image = results[0].plot()
    
    # Convert the annotated image from BGR (OpenCV format) to RGB (Gradio format)
    annotated_image_rgb = cv2.cvtColor(annotated_image, cv2.COLOR_BGR2RGB)
    
    return annotated_image_rgb

# Create a Gradio interface
iface = gr.Interface(
    fn=detect_objects,
    inputs=gr.Image(type="numpy", label="Upload Image"),
    outputs=gr.Image(type="numpy", label="Detected Objects"),
    title="YOLOv8 Object Detection",
    description="Upload an image and the YOLOv8 model will detect objects. Runs on CPU or GPU if available.",
)

if __name__ == "__main__":
    iface.launch()