| |
| 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") |
|
|