Spaces:
Running
Running
File size: 2,215 Bytes
67a9b5d |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 |
import onnxruntime
import numpy as np
class OnnxModel(object):
def __init__(self, model_path):
sess_options = onnxruntime.SessionOptions()
# # Set graph optimization level to ORT_ENABLE_EXTENDED to enable bert optimization.
# sess_options.graph_optimization_level = onnxruntime.GraphOptimizationLevel.ORT_ENABLE_EXTENDED
# # Use OpenMP optimizations. Only useful for CPU, has little impact for GPUs.
# sess_options.intra_op_num_threads = multiprocessing.cpu_count()
onnx_gpu = (onnxruntime.get_device() == 'GPU')
providers = ['CUDAExecutionProvider', 'CPUExecutionProvider'] if onnx_gpu else ['CPUExecutionProvider']
self.sess = onnxruntime.InferenceSession(model_path, sess_options, providers=providers)
self._input_names = [item.name for item in self.sess.get_inputs()]
self._output_names = [item.name for item in self.sess.get_outputs()]
@property
def input_names(self):
return self._input_names
@property
def output_names(self):
return self._output_names
def forward(self, inputs):
to_list_flag = False
if not isinstance(inputs, (tuple, list)):
inputs = [inputs]
to_list_flag = True
input_feed = {name: input for name, input in zip(self.input_names, inputs)}
outputs = self.sess.run(self.output_names, input_feed)
if (len(self.output_names) == 1) and to_list_flag:
return outputs[0]
else:
return outputs
def check_image_dtype_and_shape(image):
if not isinstance(image, np.ndarray):
raise Exception(f'image is not np.ndarray!')
if isinstance(image.dtype, (np.uint8, np.uint16)):
raise Exception(f'Unsupported image dtype, only support uint8 and uint16, got {image.dtype}!')
if image.ndim not in {2, 3}:
raise Exception(f'Unsupported image dimension number, only support 2 and 3, got {image.ndim}!')
if image.ndim == 3:
num_channels = image.shape[-1]
if num_channels not in {1, 3, 4}:
raise Exception(f'Unsupported image channel number, only support 1, 3 and 4, got {num_channels}!')
|