class OrtBase:
    def __init__(self, onnx_path):
        sess_options = ort.SessionOptions()
        sess_options.intra_op_num_threads = int(os.environ.get("CPU_CORE_NUM", 2))
        sess_options.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_EXTENDED
        providers = [
            ("CUDAExecutionProvider", {
                "device_id": 0,
                "arena_extend_strategy": "kSameAsRequested",
                "gpu_mem_limit": 5 * 1024 * 1024 * 1024,
                "cudnn_conv_algo_search": "HEURISTIC",
                "do_copy_in_default_stream": True,
            }),
            "CPUExecutionProvider" 
        ]
        # providers = ['CUDAExecutionProvider', 'CPUExecutionProvider']
        # self.session = ort.InferenceSession(path_or_bytes=onnx_path)
        # self.session = ort.InferenceSession(path_or_bytes=onnx_path, providers=providers)
        # self.session = ort.InferenceSession(path_or_bytes=onnx_path, sess_options=sess_options)
        self.session = ort.InferenceSession(path_or_bytes=onnx_path, sess_options=sess_options, providers=providers)
        self.input_names = [node.name for node in self.session.get_inputs()]
        self.output_name = [node.name for node in self.session.get_outputs()]
        print("input_name:{}".format(self.input_names))
        print("output_name:{}".format(self.output_name))

    async def forward(self, inputs, *args, **kwargs):
        """
        image_tensor = image.transpose(2, 0, 1)
        image_tensor = image_tensor[np.newaxis, :]
        onnx_session.run([output_name], {input_name: x})
        :param inputs:
        :return:
        """
        if len(self.input_names) == 1:
            input_name = self.input_names[0]
            input_feed = {input_name: inputs}
        else:
            assert len(self.input_names) == len(inputs)
            input_feed = {input_name: input_data for input_name, input_data in zip(self.input_names, inputs)}
        predict = self.session.run(self.output_name, input_feed=input_feed)
        return predict

    async def single_predict(self, img, *args, **kwargs):
        data = {'image': img[0]} 
        inputs = transform(data, self.preprocess_op)
        image = np.array(inputs)
        outputs = await self.model.forward(image)
        post_result = self.postprocess_op(outputs[0])
        return post_result