opencv python加载YOLOv5出错，如何解决？

Question

我正在尝试使用 YOLOv5 在 opencv 中运行人员检测，但我在加载 model 时遇到问题。我的 opencv 版本是 4.6.0，按照此链接中的教程，我编写了这段代码

# hyper parameters
INPUT_WIDTH = 640  # width of the input for the YOLOv5 network
INPUT_HEIGHT = 640  # height of the  input for the YOLOv5 network
SCORE_THRESHOLD = 0.5  # filter low probability classes
NMS_THRESHOLD = 0.45  # remove overlapping bounding boxes (?)
CONFIDENCE_THRESHOLD = 0.45  # filter low probability detection
BLACK = (0, 0, 0)
BLUE = (255, 178, 50)
YELLOW = (0, 255, 255)
classes = []
person_found = False
# taking classes names
class_files = 'yolo/coco.names'
with open(class_files, 'r') as source:
classes = source.read().rstrip('\n').split('\n')
# weights path
weights = 'yolo/models/yolov5s.onnx'
net = cv2.dnn.readNet(weights)
# parameters: image, scale factor, size in form of tuple, mean subtraction, swapRB
blob = cv2.dnn.blobFromImage(frame, scalefactor=1/255, size=(INPUT_WIDTH, INPUT_HEIGHT), mean=[0, 0, 0], swapRB=True, crop=False)
net.setInput(blob)
outputs = net.forward(net.getUnconnectedOutLayersNames())

但是当我尝试加载 model 时出现此错误：

[ERROR:0@1.622] global D:\a\opencv-python\opencv-python\opencv\modules\dnn\src\onnx\onnx_importer.cpp (1021) cv::dnn::dnn4_v20220524::ONNXImporter::handleNode DNN/ONNX: ERROR during processing node with 1 inputs and 1 outputs: [Identity]:(onnx_node!Identity_0) from domain='ai.onnx'
Traceback (most recent call last):
  File "D:\Unitn\Corsi\Computer vision\project\synopsis\main.py", line 25, in <module>
    bg_frames = find_bg_frames(video, duration, fps)
  File "D:\Unitn\Corsi\Computer vision\project\synopsis\utils.py", line 152, in find_bg_frames
    people_found = find_people_with_yolo(frame)
  File "D:\Unitn\Corsi\Computer vision\project\synopsis\utils.py", line 102, in find_people_with_yolo
    net = cv2.dnn.readNetFromONNX(weights)
cv2.error: OpenCV(4.6.0) D:\a\opencv-python\opencv-python\opencv\modules\dnn\src\onnx\onnx_importer.cpp:1040: error: (-2:Unspecified error) in function 'cv::dnn::dnn4_v20220524::ONNXImporter::handleNode'
> Node [Identity@ai.onnx]:(onnx_node!Identity_0) parse error: OpenCV(4.6.0) D:\a\opencv-python\opencv-python\opencv\modules\dnn\src\layer.cpp:246: error: (-215:Assertion failed) inputs.size() in function 'cv::dnn::dnn4_v20220524::Layer::getMemoryShapes'
> 

谁能帮我找到问题所在？

Answer 1

我有同样的问题，我通过为导出添加“--simplify”来解决它

Answer 2

仔细观察outputs形状，发现它是 [1, 3, 80, 80, 85]。 而对于默认的 640x640 onnx 导出，它应该是 [25200x85]。 这似乎是从 PyTorch 到 ONNX 的错误转换。 在yolov5s.onnx模型文章中已更新。

此外，torch==1.12 有一些错误。 在修复之前，请在转换模型时使用 torch==1.11。

在 GitHub 上查看此问题。