import cv2
import gradio as gr
from PIL import Image
import numpy as np

with open('coco.names', 'r') as f:
    classes = f.read().splitlines()

net = cv2.dnn.readNetFromDarknet('yolov4.cfg', 'yolov4.weights')

model = cv2.dnn_DetectionModel(net)
model.setInputParams(scale=1 / 255, size=(416, 416), swapRB=True)

def detect(image):
    img = Image.fromarray(image)
    img = cv2.cvtColor(np.array(img), cv2.COLOR_BGR2RGB)
    classIds, scores, boxes = model.detect(img, confThreshold=0.6, nmsThreshold=0.4)

    for (classId, score, box) in zip(classIds, scores, boxes):
        cv2.rectangle(img, (box[0], box[1]), (box[0] + box[2], box[1] + box[3]),
                      color=(0, 255, 0), thickness=2)
        text = '%s: %.2f' % (classes[classId], score)
        cv2.putText(img, text, (box[0], box[1] - 5), cv2.FONT_HERSHEY_SIMPLEX, 1,
                    color=(0, 255, 0), thickness=2)

    return img

image_in = gr.Image()
image_out = gr.components.Image()

Iface = gr.Interface(
    fn=detect,
    inputs=image_in,
    outputs=image_out,
    title="Object Detection with YOLOS",
).launch()