Using an onnx model to make predictions on an image using opencv

Viewed 1016

I'm trying to make real-time predictions on video frames using a pre-trained model (YOLOv5) saved in .onnx format through OpenCV. I keep getting the following error

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

My current code looks like this:

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

I haven't included the code to display the bounding boxes yet as I just want to get the model.forward() working first. The blob and models expected input size are the same, I'm not sure what else to try.

0 Answers
Related