## https://github.com/BlueMirrors/Yolov5-ONNX/blob/main/detect.py
import gradio as gr
import cv2
import numpy as np
from cvu.detector.yolov5 import Yolov5 as Yolov5Onnx
from vidsz.opencv import Reader, Writer

def process_input(input_file, weight='./jtbz_opt.onnx', device='cpu'):
    # Determine if the input is an image or a video
    if input_file.name.endswith(('.jpg', '.jpeg', '.png')):
        # load model
        model = Yolov5Onnx(classes="coco",
                           backend="onnx",
                           weight=weight,
                           device=device)

        # read image
        image = cv2.imread(input_file.name)

        # inference
        preds = model(image)

        # draw image
        preds.draw(image)

        # return image
        return image

    else:
        model = Yolov5Onnx(classes="coco",
                           backend="onnx",
                           weight=weight,
                           device=device)
        reader = Reader(input_file.name)
        writer = Writer(reader, name="output.mp4")

        # warmup
        warmup = np.random.randint(0, 255, reader.read().shape).astype("float")
        for i in range(100):
            model(warmup)

        inference_time = 0
        for frame in reader:
            # inference
            start = time.time()
            preds = model(frame)
            inference_time += time.time() - start

            # draw on frame
            preds.draw(frame)
            writer.write(frame)

        writer.release()
        reader.release()

        # return video
        return gr.outputs.Video("output.mp4")