caffe转onnx

pip install -i Simple Index

源码需要依赖caffe的python接口

代码路径：GitHub - inisis/caffe2onnx: caffe model to onnx

注意事项：caffe的relu和pool层需要修改名字，否则会报错，如下：

onnx.onnx_cpp2py_export.checker.ValidationError: Graph must be in single static assignment (SSA) form, however 'conv2_relu' has been used as output names multiple times.

测试代码：

import onnxruntime

import numpy as np

import PIL.Image as Image

import argparse

import cv2

def cv2_imread(path):

    img = cv2.imdecode(np.fromfile(path, dtype=np.uint8), 1)

    return img

def process_image(img_path,input_shape):

    #img = Image.open(img_path).convert("BGR")

    #img = img.resize(input_shape)

    img = cv2_imread(img_path)

    img = cv2.resize(img,(320,320))

    image = np.array(img, dtype=np.float32)

    image = image/255.0

    image = image.transpose((2,0,1))[np.newaxis, ...]

    return image

def main():

    parser = argparse.ArgumentParser()

    parser.add_argument('--input_shape', help="caffe's caffemodel file path", nargs='+', default=(320,320))

    parser.add_argument('--img_path', help="test image path", type=str, default="./220728103415203_QZ.jpg")

    parser.add_argument('--onnx_path', help="onnx model file path", type=str, default="./tobacco_jjh_mirror_320_new_iter_8000_0.8393.onnx")

    args = parser.parse_args()

    input_shape = [int(x) for x in args.input_shape] #模型输入尺寸

    img_path = args.img_path

    onnx_path = args.onnx_path

    print("image path:",img_path)

    print("onnx model path:",onnx_path)

    data_input = process_image(img_path,input_shape)

    session = onnxruntime.InferenceSession(onnx_path)

    inname = [input.name for input in session.get_inputs()]

    outname = [output.name for output in session.get_outputs()]

    print("inputs name:",inname,"|| outputs name:",outname)

    data_output = session.run(outname, {inname[0]: data_input})

    output = data_output[0]

    print("Label predict: ", output.argmax())

if __name__ == '__main__':

    main()