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