File size: 1,627 Bytes
7a4e164 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 |
## 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")
|