app.py

import gradio as gr
import torch
import cv2

model = torch.hub.load('ultralytics/yolov5', 'custom', 'model/best.onnx')


CLASS_COLORS = {
    0: [148, 0, 211],  # class 1 (violet)
    1: [255, 0, 0],    # class 2 (red)
    2: [255, 127, 0],  # class 3 (orange)
    3: [255, 255, 0],  # class 4 (yellow)
    4: [0, 255, 0],    # class 5 (green)
    5: [0, 0, 255],    # class 6 (blue)
}


def object_detection(image):
    image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)

    results = model(image)

    bboxes = results.xyxy[0].tolist()
    labels = results.xyxy[0][:, -1].long().tolist()
    scores = results.xyxy[0][:, -2].tolist()

    for bbox, label, score in zip(bboxes, labels, scores):
        label_name = results.names[label]
        color = CLASS_COLORS[label]

        cv2.rectangle(image, (int(bbox[0]), int(bbox[1])), (int(bbox[2]), int(bbox[3])), color, 2)
        text = f"{label_name} ({score:.2f})"
        cv2.putText(image, text, (int(bbox[0]), int(bbox[1]) - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.5, color, 2)

    image = cv2.cvtColor(image, cv2.COLOR_RGB2BGR)

    return image


inputs = gr.components.Image(shape=(640, 640))
outputs = gr.components.Image(label='Input Image', shape=(640, 640))

iface = gr.Interface(fn=object_detection, inputs=inputs, outputs=outputs,
                     examples=['examples/india.jpeg', 'examples/new-york.jpeg', 'examples/pedestrian-bikes.jpeg'])

iface.launch()