| 4 | |
| 5 | |
| 6 | class ONNXEngine: |
| 7 | |
| 8 | def __init__(self, onnx_path, use_gpu): |
| 9 | """ |
| 10 | :param onnx_path: |
| 11 | """ |
| 12 | if not os.path.exists(onnx_path): |
| 13 | raise Exception(f'{onnx_path} is not exists') |
| 14 | |
| 15 | providers = ['CPUExecutionProvider'] |
| 16 | if use_gpu: |
| 17 | providers = ([ |
| 18 | 'TensorrtExecutionProvider', |
| 19 | 'CUDAExecutionProvider', |
| 20 | 'CPUExecutionProvider', |
| 21 | ], ) |
| 22 | self.onnx_session = onnxruntime.InferenceSession(onnx_path, |
| 23 | providers=providers) |
| 24 | self.input_name = self.get_input_name(self.onnx_session) |
| 25 | self.output_name = self.get_output_name(self.onnx_session) |
| 26 | |
| 27 | def get_output_name(self, onnx_session): |
| 28 | """ |
| 29 | output_name = onnx_session.get_outputs()[0].name |
| 30 | :param onnx_session: |
| 31 | :return: |
| 32 | """ |
| 33 | output_name = [] |
| 34 | for node in onnx_session.get_outputs(): |
| 35 | output_name.append(node.name) |
| 36 | return output_name |
| 37 | |
| 38 | def get_input_name(self, onnx_session): |
| 39 | """ |
| 40 | input_name = onnx_session.get_inputs()[0].name |
| 41 | :param onnx_session: |
| 42 | :return: |
| 43 | """ |
| 44 | input_name = [] |
| 45 | for node in onnx_session.get_inputs(): |
| 46 | input_name.append(node.name) |
| 47 | return input_name |
| 48 | |
| 49 | def get_input_feed(self, input_name, image_numpy): |
| 50 | """ |
| 51 | input_feed={self.input_name: image_numpy} |
| 52 | :param input_name: |
| 53 | :param image_numpy: |
| 54 | :return: |
| 55 | """ |
| 56 | input_feed = {} |
| 57 | for name in input_name: |
| 58 | input_feed[name] = image_numpy |
| 59 | return input_feed |
| 60 | |
| 61 | def run(self, image_numpy): |
| 62 | # 输入数据的类型必须与模型一致,以下三种写法都是可以的 |
| 63 | input_feed = self.get_input_feed(self.input_name, image_numpy) |