Update onnx_inference.py

onnx_inference.py

@@ -127,8 +127,11 @@ if __name__ == '__main__':
 
     img0 = cv2.imread(path)
     img = pre_process(img0)
-    onnx_input = {onnx_model.get_inputs()[0].name: img}
+    # onnx_input = {onnx_model.get_inputs()[0].name: img}
+    onnx_input = {onnx_model.get_inputs()[0].name: np.transpose(img, (0, 2 ,3, 1))}
     onnx_output = onnx_model.run(None, onnx_input)
+    onnx_output = [np.transpose(out, (0, 3, 1, 2)) for out in onnx_output]
+
     pred = post_process(onnx_output, conf_thres,
                         iou_thres, multi_label=False,
                         classes=classes, agnostic=agnostic_nms)