|
|
|
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'): |
|
|
|
if input_file.name.endswith(('.jpg', '.jpeg', '.png')): |
|
|
|
model = Yolov5Onnx(classes="coco", |
|
backend="onnx", |
|
weight=weight, |
|
device=device) |
|
|
|
|
|
image = cv2.imread(input_file.name) |
|
|
|
|
|
preds = model(image) |
|
|
|
|
|
preds.draw(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 = np.random.randint(0, 255, reader.read().shape).astype("float") |
|
for i in range(100): |
|
model(warmup) |
|
|
|
inference_time = 0 |
|
for frame in reader: |
|
|
|
start = time.time() |
|
preds = model(frame) |
|
inference_time += time.time() - start |
|
|
|
|
|
preds.draw(frame) |
|
writer.write(frame) |
|
|
|
writer.release() |
|
reader.release() |
|
|
|
|
|
return gr.outputs.Video("output.mp4") |
|
|