cvjtbz / app-vid.py
cxeep's picture
Create app-vid.py
7a4e164 verified
raw
history blame
1.63 kB
## 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")