text: # !/usr/bin/env python
  # -*- coding: UTF-8 -*-
  """
  @Project ：my.onnx
  @File    ：onnx_tools.py
  @Author  ：Linfeng Lee
  @Date    ：2021/7/7 下午3:37
  """
  import os
  import sys
  from typing import Optional, Sequence, Text, Union
  
  import numpy as np
  import onnx
  import torchvision.transforms as transforms
  import onnxruntime
  import torch
  import torch.onnx
  from torch.autograd import Variable
  
  sys.path.append(os.getcwd())
  
  Path = str
  TensorType = torch.Tensor
  NumpyType = np.ndarray
  ModelProtoType = onnx.ModelProto
  GraphProtoType = onnx.GraphProto
  NodeProtoType = onnx.NodeProto
  
  
  def write_onnx_data(model: Optional[Union[ModelProtoType, str]], path: Path):
      if type(model) == str:
          # 加载onnx模型
          model = onnx.load(model)
      print('write_onnx_data  In {}'.format(path))
      with open(path, 'w') as f:
          f.write(str(model))
  
  
  class ONNXModel:
      def __init__(self, onnx_path: Path):
          self.onnx_path = onnx_path
          assert os.path.exists(self.onnx_path) is True, 'Onnx path is not exists'
  
          onnx.checker.check_model(onnx_path)
  
          self.model = onnx.load(self.onnx_path)
          self.graph = self.model.graph
  
          self.onnx_session = onnxruntime.InferenceSession(onnx_path)
          self.input_name = self.get_input_name()
          self.output_name = self.get_output_name()
  
      def get_model(self, save_path: Optional[Path] = None) -> ModelProtoType:
          write_onnx_data(self.model, save_path) if save_path is not None else ...
          return self.model
  
      def get_graph(self, save_path: Optional[Path] = None) -> GraphProtoType:
          write_onnx_data(self.graph, save_path) if save_path is not None else ...
          return self.graph
  
      def get_nodes(self, save_path: Optional[Path] = None):
          nodes = self.graph.node
          write_onnx_data(nodes, save_path) if save_path is not None else ...
          return nodes
  
      def get_node_from_idx(self, idx: int, save_path: Optional[Path] = None) -> Optional[NodeProtoType]:
          nodes = self.get_nodes()
          for i, node in enumerate(nodes):
              if node.output[0] == str(idx):
                  write_onnx_data(node, save_path) if save_path is not None else ...
                  return node
          return None
  
      def get_node_from_name(self, name: str, save_path: Optional[Path] = None) -> Optional[NodeProtoType]:
          nodes = self.get_nodes()
          for i, node in enumerate(nodes):
              if node.name == name:
                  write_onnx_data(node, save_path) if save_path is not None else ...
                  return node
          return None
  
      # TODO
      def modify_node(self,
                      node_idx: int,
                      op_type: str,
                      inputs: Sequence[Text],
                      outputs: Sequence[Text],
                      name: Optional[Text] = None,
                      doc_string: Optional[Text] = None,
                      domain: Optional[Text] = None,
                      **kwargs) -> None:
  
          replace_node = self.get_node_from_idx(node_idx)
          new_node = onnx.helper.make_node(
              op_type,
              inputs,
              outputs,
              name,
              doc_string,
              domain,
              **kwargs
          )  # 新建新节点
          print(new_node)
          replace_node.input = new_node.input
          print(replace_node)
  
      # print(new_node)
      # nodes.remove(replace_node)  # 删除旧节点
      # nodes.insert(node_idx, new_node)  # 插入新节点
  
      def save_model(self, path: Path):
          onnx.checker.check_model(self.model)
          onnx.save(self.model, path)
  
      def get_initializers(self, save_path: Optional[Path] = None):
          datas = self.graph.initializer
          write_onnx_data(datas, save_path) if save_path is not None else ...
          return datas
  
      def get_initializer_from_name(self, idx: str, save_path: Optional[Path] = None) -> Optional[onnx.TensorProto]:
          initializers = self.graph.initializer
          for i, data in enumerate(initializers):
              if data.name == idx:
                  write_onnx_data(data, save_path) if save_path is not None else ...
                  return data
          return None
  
      def get_input(self, save_path: Optional[Path] = None):
          g_input = self.graph.input
          write_onnx_data(g_input, save_path) if save_path is not None else ...
          return g_input
  
      def get_output(self, save_path: Optional[Path] = None):
          output = self.graph.output
          write_onnx_data(output, save_path) if save_path is not None else ...
          return output
  
      def get_output_name(self) -> list:
          """
          output_name = onnx_session.get_outputs()[0].name
          :return:
          """
          output_name = []
          for node in self.onnx_session.get_outputs():
              output_name.append(node.name)
          return output_name
  
      def get_input_name(self) -> list:
          """
          input_name = onnx_session.get_inputs()[0].name
          :return:
          """
          input_name = []
          for node in self.onnx_session.get_inputs():
              input_name.append(node.name)
          return input_name
  
      def get_input_feed(self, image_numpy: NumpyType) -> dict:
          """
          input_feed={self.input_name: image_numpy}
          :param image_numpy:
          :return:
          """
          input_feed = {}
          for name in self.input_name:
              input_feed[name] = image_numpy
          return input_feed
  
      def forward(self, image_numpy: NumpyType):
          """
          # image_numpy = image.transpose(2, 0, 1)
          # image_numpy = image_numpy[np.newaxis, :]
          # onnx_session.run([output_name], {input_name: x})
          # :param image_numpy:
          # :return:
          """
          # 输入数据的类型必须与模型一致,以下三种写法都是可以的
          # scores, boxes = self.onnx_session.run(None, {self.input_name: image_numpy})
          # scores, boxes = self.onnx_session.run(self.output_name, input_feed={self.input_name: iimage_numpy})
          input_feed = self.get_input_feed(image_numpy)
          print(self.onnx_session.get_inputs()[0])
          print(self.onnx_session.get_outputs()[0])
          result = self.onnx_session.run(self.output_name, input_feed=input_feed)
          return result
  
  
  def to_numpy(tensor: TensorType):
      return tensor.detach().cpu().numpy() if tensor.requires_grad else tensor.cpu().numpy()
  
  
  def check_proto_file(file_path: str):
      # Load the ONNX model
      model = onnx.load(file_path)
      # Check that the IR is well formed
      onnx.checker.check_model(model)
      # Print a human readable representation of the graph
      ret = onnx.helper.printable_graph(model.graph)
      return ret
  
  
  def torch2onnx(model, save_path: Path, input_size: list, dummy_input: Optional[TensorType] = None):
      """
      :param dummy_input:
      :param model:
      :param save_path:  XXX/XXX.my.onnx
      :return: my.onnx model
      """
  
      if torch.cuda.is_available():
          print("Use GPU")
      else:
          print("Use CPU")
  
      model = model.cuda() if torch.cuda.is_available() else model
      model.eval()
      input_data = dummy_input if dummy_input is not None else torch.randn(input_size)
      input_data = Variable(input_data).cuda() if torch.cuda.is_available() else Variable(input_data)
      print('input data shape:', input_data.shape)
      input_names = ["input"]
      output_names = ["out"]
      torch.onnx.export(model,
                        input_data,
                        save_path,
                        export_params=True,
                        verbose=True,
                        opset_version=11,
                        input_names=input_names,
                        output_names=output_names)
      print("torch2onnx finish.")
  
  
  def torch2onnx_dynamic(model, save_path: Path, dummy_input: Optional[TensorType] = None):
      """
      :param dummy_input:
      :param model:
      :param save_path:  XXX/XXX.my.onnx
      :return: my.onnx model
      """
      model.eval()
      input_data = dummy_input if dummy_input is not None else torch.rand(1, 3, 224, 224)
  
      input_names = ["input"]  # ncnn需要
      output_names = ["out"]  # ncnn需要
      torch.onnx.export(model,
                        input_data,
                        save_path,
                        export_params=True,
                        opset_version=11,
                        input_names=input_names,
                        output_names=output_names,
                        dynamic_axes={'input': [2, 3], 'out': [2, 3]})
      print("torch2onnx dynamic finish.")
  
  
  # 加载网络权重
  def load_network(network, network_weights_path: Path):
      assert os.path.exists(network_weights_path) is True, 'Weight is not exists'
      network.load_state_dict(torch.load(network_weights_path))
      return network
  
  
  def get_transform():
      return transforms.Compose([
          transforms.Resize([224, 224]),
          transforms.ToTensor(),
          transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])])
  