使用 onnx 模型使用 opencv 对图像进行预测

时间:2021-05-19 10:08:34

标签: python opencv computer-vision onnx

我正在尝试使用通过 OpenCV 以 .onnx 格式保存的预训练模型 (YOLOv5) 对视频帧进行实时预测。我不断收到以下错误

cv2.error: OpenCV(4.5.2) C:\Users\runneradmin\AppData\Local\Temp\pip-req-build-dn5w5exm\opencv\modules\dnn\include\opencv2/dnn/shape_utils.hpp:171: error: (-215:Assertion failed) start <= (int)shape.size() && end <= (int)shape.size() && 
start <= end in function 'cv::dnn::dnn4_v20210301::total

我当前的代码如下所示:

import cv2 as cv
import numpy as np

cap = cv.VideoCapture('PXL_20210517_094123898_2.mp4')

width = int(cap.get(cv.CAP_PROP_FRAME_WIDTH) / 3)
height = int(cap.get(cv.CAP_PROP_FRAME_HEIGHT) / 3)

model = cv.dnn.readNetFromONNX("yolov5s.onnx")
ln = model.getLayerNames()
ln = [ln[i[0] - 1] for i in model.getUnconnectedOutLayers()]
output_layers = ln[-1]

while True:
    # Capture frame-by-frame
    ret, frame = cap.read()
    # if frame is read correctly ret is True
    if not ret:
        print("Can't receive frame (stream end?). Exiting ...")
        break
    # Resize the frame to fit on my screen
    frame_rs = cv.resize(frame, (width, height))
    #blob = cv.dnn.blobFromImage(frame_rs, 1/255.0, (416, 416), swapRB=True, crop=False)
    blob = cv.dnn.blobFromImage(image=frame_rs, scalefactor=1, size=(416, 416)) #, mean=(104, 117, 123))
    model.setInput(blob)
    outputs = model.forward()
    
    # Display the resulting frame
    cv.imshow("frame", frame_rs)

    key = cv.waitKey(1)
    if key == 27:
        break

cap.release()
cv.destroyAllWindows()

我还没有包含显示边界框的代码,因为我只想先让 model.forward() 工作。 blob 和模型的预期输入大小相同,我不确定还可以尝试什么。

0 个答案:

没有答案
相关问题