File size: 1,931 Bytes
e989ba3 |
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 59 60 61 62 63 64 65 66 67 68 69 |
from tools.infer_onnx import ONNX_Predict
import cv2
import time
import collections
onnx_model = "edge_s_640_resize.onnx"
providers = ["CPUExecutionProvider"] #or ["CUDAExecutionProvider"]
predict = ONNX_Predict(onnx_model, providers=providers, use_letterbox=False) #Initialize Onnx predict class
#Settings for ONNX_Predict class
img_size = 640 #Image size for onnx model
conf = 0.5
iou = 0.3
max_det = 300
cap = cv2.VideoCapture("intersection.mp4")
max_w = 800 #Display width
#FPS-smoothing
fps_history = collections.deque(maxlen=30)
while True:
ret, frame = cap.read()
start_time = time.time()
if not ret:
break
h, w = frame.shape[:2]
if w > max_w:
scale = max_w / w
frame = cv2.resize(frame, (max_w, int(h * scale)), interpolation=cv2.INTER_LINEAR)
#Measure only inference time
start_time_0 = time.time()
boxes, scores, classes = predict.infer_image(
frame, img_size=img_size, conf=conf, iou=iou, max_det=max_det
)
end_time_0 = time.time()
#Draw
for (x1, y1, x2, y2), score, cls in zip(boxes, scores, classes):
x1, y1, x2, y2 = map(int, (x1, y1, x2, y2))
cv2.rectangle(frame, (x1, y1), (x2, y2), (0,255,0), 2)
cv2.putText(frame, f"{int(cls)} {score:.2f}", (x1, max(0, y1-5)),
cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0,255,0), 1)
dt = time.time() - start_time
fps = 1.0 / dt if dt > 0 else 0.0
fps_history.append(fps)
fps_avg = sum(fps_history) / len(fps_history)
cv2.putText(frame, f"edge_s 640x640 CPU FPS: {fps_avg:5.1f}",
(10, 25), cv2.FONT_HERSHEY_SIMPLEX, 0.7, (0,255,0), 2)
cv2.imshow('test', frame)
key = cv2.waitKey(1)
print(f'Total: {fps_avg:5.1f} Inference:{(end_time_0-start_time_0)*1000:3.1f}ms')
if key == 13: # Enter
break
|