• 对paddleOCR中的字符识别模型转ONNX

    对paddle OCR中的模型转换成ONNX。

    转换代码：

    1.  
    2.  
    3. import os
    4. import sys
    5. import yaml
    6. import numpy as np
    7. import cv2
    8. import argparse
    9. import paddle
    10. from paddle import nn
    11.  
    12. from argparse import ArgumentParser, RawDescriptionHelpFormatter
    13. import paddle.distributed as dist
    14. from ppocr.postprocess import build_post_process
    15. from ppocr.utils.save_load import init_model
    16. from ppocr.modeling.architectures import build_model
    17.  
    18.  
    19. class AttrDict(dict):
    20.     """Single level attribute dict, NOT recursive"""
    21.  
    22.     def __init__(self, **kwargs):
    23.         super(AttrDict, self).__init__()
    24.         super(AttrDict, self).update(kwargs)
    25.  
    26.     def __getattr__(self, key):
    27.         if key in self:
    28.             return self[key]
    29.         raise AttributeError("object has no attribute '{}'".format(key))
    30.  
    31. global_config = AttrDict()
    32. default_config = {'Global': {'debug': False, }}
    33.  
    34. class ArgsParser(ArgumentParser):
    35.     def __init__(self):
    36.         super(ArgsParser, self).__init__(
    37.             formatter_class=RawDescriptionHelpFormatter)
    38.         # self.add_argument("-c", "--config", default='./configs/ch_PP-OCRv2_rec_idcard.yml',
    39.         #                   help="configuration file to use")
    40.  
    41.         self.add_argument("-c", "--config", default='./configs/ch_PP-OCRv2_rec.yml',
    42.                           help="configuration file to use")
    43.         self.add_argument(
    44.             "-o", "--opt", nargs='+', help="set configuration options")
    45.  
    46.     def parse_args(self, argv=None):
    47.         args = super(ArgsParser, self).parse_args(argv)
    48.         assert args.config is not None, \
    49.             "Please specify --config=configure_file_path."
    50.         args.opt = self._parse_opt(args.opt)
    51.         return args
    52.  
    53.     def _parse_opt(self, opts):
    54.         config = {}
    55.         if not opts:
    56.             return config
    57.         for s in opts:
    58.             s = s.strip()
    59.             k, v = s.split('=')
    60.             config[k] = yaml.load(v, Loader=yaml.Loader)
    61.         return config
    62.  
    63. def merge_config(config):
    64.     """
    65.     Merge config into global config.
    66.     Args:
    67.         config (dict): Config to be merged.
    68.     Returns: global config
    69.     """
    70.     for key, value in config.items():
    71.         if "." not in key:
    72.             if isinstance(value, dict) and key in global_config:
    73.                 global_config[key].update(value)
    74.             else:
    75.                 global_config[key] = value
    76.         else:
    77.             sub_keys = key.split('.')
    78.             assert (
    79.                 sub_keys[0] in global_config
    80.             ), "the sub_keys can only be one of global_config: {}, but get: {}, please check your running command".format(
    81.                 global_config.keys(), sub_keys[0])
    82.             cur = global_config[sub_keys[0]]
    83.             for idx, sub_key in enumerate(sub_keys[1:]):
    84.                 if idx == len(sub_keys) - 2:
    85.                     cur[sub_key] = value
    86.                 else:
    87.                     cur = cur[sub_key]
    88.  
    89. def load_config(file_path):
    90.     """
    91.     Load config from yml/yaml file.
    92.     Args:
    93.         file_path (str): Path of the config file to be loaded.
    94.     Returns: global config
    95.     """
    96.     merge_config(default_config)
    97.     _, ext = os.path.splitext(file_path)
    98.     assert ext in ['.yml', '.yaml'], "only support yaml files for now"
    99.     merge_config(yaml.load(open(file_path, 'rb'), Loader=yaml.Loader))
    100.     return global_config
    101.  
    102. def check_device(use_gpu, use_xpu=False):
    103.     """
    104.     Log error and exit when set use_gpu=true in paddlepaddle
    105.     cpu version.
    106.     """
    107.     err = "Config {} cannot be set as true while your paddle " \
    108.           "is not compiled with {} ! \nPlease try: \n" \
    109.           "\t1. Install paddlepaddle to run model on {} \n" \
    110.           "\t2. Set {} as false in config file to run " \
    111.           "model on CPU"
    112.  
    113.     try:
    114.         if use_gpu and use_xpu:
    115.             print("use_xpu and use_gpu can not both be ture.")
    116.         if use_gpu and not paddle.is_compiled_with_cuda():
    117.             print(err.format("use_gpu", "cuda", "gpu", "use_gpu"))
    118.             sys.exit(1)
    119.         if use_xpu and not paddle.device.is_compiled_with_xpu():
    120.             print(err.format("use_xpu", "xpu", "xpu", "use_xpu"))
    121.             sys.exit(1)
    122.     except Exception as e:
    123.         pass
    124.  
    125. def getArgs(is_train=False):
    126.     FLAGS = ArgsParser().parse_args()
    127.     config = load_config(FLAGS.config)
    128.     merge_config(FLAGS.opt)
    129.  
    130.     # check if set use_gpu=True in paddlepaddle cpu version
    131.     use_gpu = config['Global']['use_gpu']
    132.  
    133.     use_xpu = False
    134.  
    135.     alg = config['Architecture']['algorithm']
    136.     assert alg in [
    137.         'EAST', 'DB', 'SAST', 'Rosetta', 'CRNN', 'STARNet', 'RARE', 'SRN',
    138.         'CLS', 'PGNet', 'Distillation', 'NRTR', 'TableAttn', 'SAR', 'PSE',
    139.         'SEED', 'SDMGR', 'LayoutXLM', 'LayoutLM', 'LayoutLMv2', 'PREN', 'FCE',
    140.         'SVTR', 'ViTSTR', 'ABINet', 'DB++', 'TableMaster', 'SPIN', 'VisionLAN',
    141.         'Gestalt', 'SLANet', 'RobustScanner'
    142.     ]
    143.  
    144.     device = 'gpu:{}'.format(dist.ParallelEnv().dev_id) if use_gpu else 'cpu'
    145.     check_device(use_gpu, use_xpu)
    146.  
    147.     device = paddle.set_device(device)
    148.  
    149.     config['Global']['distributed'] = dist.get_world_size() != 1
    150.  
    151.     return config, device
    152.  
    153.  
    154. class CRNN(nn.Layer):
    155.     def __init__(self, config, device):
    156.         super(CRNN, self).__init__()
    157.         # 定义预处理参数
    158.         mean = (0.5, 0.5, 0.5)
    159.         std = (0.5, 0.5, 0.5)
    160.         self.mean = paddle.to_tensor(mean).reshape([1, 3, 1, 1])
    161.         self.std = paddle.to_tensor(std).reshape([1, 3, 1, 1])
    162.  
    163.         self.config = config
    164.         # build post process
    165.         self.post_process_class = build_post_process(config['PostProcess'],
    166.                                                      config['Global'])
    167.         # build model
    168.         if hasattr(self.post_process_class, 'character'):
    169.             char_num = len(getattr(self.post_process_class, 'character'))
    170.             if self.config['Architecture']["algorithm"] in ["Distillation",
    171.                                                             ]:  # distillation model
    172.                 for key in self.config['Architecture']["Models"]:
    173.                     if self.config['Architecture']['Models'][key]['Head'][
    174.                         'name'] == 'MultiHead':  # for multi head
    175.                         out_channels_list = {}
    176.                         if self.config['PostProcess'][
    177.                             'name'] == 'DistillationSARLabelDecode':
    178.                             char_num = char_num - 2
    179.                         out_channels_list['CTCLabelDecode'] = char_num
    180.                         out_channels_list['SARLabelDecode'] = char_num + 2
    181.                         self.config['Architecture']['Models'][key]['Head'][
    182.                             'out_channels_list'] = out_channels_list
    183.                     else:
    184.                         self.config['Architecture']["Models"][key]["Head"][
    185.                             'out_channels'] = char_num
    186.             elif self.config['Architecture']['Head'][
    187.                 'name'] == 'MultiHead':  # for multi head
    188.                 out_channels_list = {}
    189.                 if self.config['PostProcess']['name'] == 'SARLabelDecode':
    190.                     char_num = char_num - 2
    191.                 out_channels_list['CTCLabelDecode'] = char_num
    192.                 out_channels_list['SARLabelDecode'] = char_num + 2
    193.                 self.config['Architecture']['Head'][
    194.                     'out_channels_list'] = out_channels_list
    195.             else:  # base rec model
    196.                 self.config['Architecture']["Head"]['out_channels'] = char_num
    197.  
    198.         # 加载模型
    199.         self.model = build_model(config['Architecture'])
    200.         # load_model(config, self.model)
    201.         init_model(self.config, self.model)
    202.         self.model.eval()
    203.  
    204.     def forward(self, x):
    205.         # x = paddle.transpose(x, [0,3,1,2])
    206.         # x = x / 255.0
    207.         # x = (x - self.mean) / self.std
    208.  
    209.         model_out = self.model(x)
    210.  
    211.         # return model_out
    212.         preds_idx = model_out.argmax(axis=2, name='class').astype('float32')
    213.         # preds_idx = model_out.argmax(axis=2, name='class')
    214.         preds_prob = model_out.max(axis=2, name='score').astype('float32')
    215.         return preds_idx, preds_prob
    216.  
    217. EXPORT_ONNX = True
    218. DYNAMIC = False
    219.  
    220. if __name__ == '__main__':
    221.     config, device = getArgs()
    222.     model_crnn = CRNN(config, device=device)
    223.  
    224.     # 构建输入数据images:
    225.     image_path = "1.jpg"
    226.     img = cv2.imread(image_path)
    227.     img = cv2.resize(img, (320, 32))
    228.     print('input data:', img.shape)
    229.     img = img.astype(np.float32)
    230.     img = img.transpose((2, 0, 1)) / 255
    231.     input_data = img[np.newaxis, :]
    232.     print('input data:', input_data.shape)
    233.     x = paddle.to_tensor(input_data)
    234.     print('input data:', x.shape)
    235.  
    236.     output_idx, output_prob = model_crnn(x)
    237.     print('output_idx: ', output_idx)
    238.     print('output_prob: ', output_prob)
    239.  
    240.     input_spec = paddle.static.InputSpec.from_tensor(x,  name='input')
    241.     onnx_save_path = "./export_onnx"
    242.     if EXPORT_ONNX:
    243.         onnx_model_name = onnx_save_path + "/char_recognize_20230526_v1"
    244.         if DYNAMIC:
    245.             input_spec = paddle.static.InputSpec(
    246.                 shape=[None, 32, 320, 3], dtype='float32',  name='input')
    247.  
    248.         # ONNX模型导出
    249.         paddle.onnx.export(model_crnn, onnx_model_name, input_spec=[input_spec], opset_version=11,
    250.                            enable_onnx_checker=True, output_spec=[output_idx, output_prob])

    转换后的网络结构绘制出来，绘制使用的工具Netron

     绘制出的起始和末尾的网络结构：

    测试ONNX的代码：

    1. '''
    2. 测试转出的onnx模型
    3. '''
    4. import cv2
    5. import numpy as np
    6.  
    7. import torch
    8. import onnxruntime as rt
    9. import math
    10. import os
    11.  
    12. class TestOnnx:
    13.     def __init__(self, onnx_file, character_dict_path, use_space_char=True):
    14.         self.sess = rt.InferenceSession(onnx_file)
    15.         # 获取输入节点名称
    16.         self.input_names = [input.name for input in self.sess.get_inputs()]
    17.         # 获取输出节点名称
    18.         self.output_names = [output.name for output in self.sess.get_outputs()]
    19.  
    20.         self.character = []
    21.         self.character.append("blank")
    22.         with open(character_dict_path, "rb") as fin:
    23.             lines = fin.readlines()
    24.             for line in lines:
    25.                 line = line.decode('utf-8').strip("\n").strip("\r\n")
    26.                 self.character.append(line)
    27.         if use_space_char:
    28.             self.character.append(" ")
    29.  
    30.     def resize_norm_img(self, img, image_shape=[3, 32, 320]):
    31.         imgC, imgH, imgW = image_shape
    32.         h = img.shape[0]
    33.         w = img.shape[1]
    34.         ratio = w / float(h)
    35.         if math.ceil(imgH * ratio) > imgW:
    36.             resized_w = imgW
    37.         else:
    38.             resized_w = int(math.ceil(imgH * ratio))
    39.         resized_image = cv2.resize(img, (resized_w, imgH))
    40.         resized_image = resized_image.astype('float32')
    41.         if image_shape[0] == 1:
    42.             resized_image = resized_image / 255
    43.             resized_image = resized_image[np.newaxis, :]
    44.         else:
    45.             resized_image = resized_image.transpose((2, 0, 1)) / 255
    46.         resized_image -= 0.5
    47.         resized_image /= 0.5
    48.         padding_im = np.zeros((imgC, imgH, imgW), dtype=np.float32)
    49.         padding_im[:, :, 0:resized_w] = resized_image
    50.         return padding_im
    51.  
    52.     # # 准备模型运行的feed_dict
    53.     def process(self, input_names, image):
    54.         feed_dict = dict()
    55.         for input_name in input_names:
    56.             feed_dict[input_name] = image
    57.  
    58.         return feed_dict
    59.  
    60.     def get_ignored_tokens(self):
    61.         return [0]
    62.  
    63.     def decode(self, text_index, text_prob=None, is_remove_duplicate=False):
    64.         """ convert text-index into text-label. """
    65.         result_list = []
    66.         ignored_tokens = self.get_ignored_tokens()
    67.         batch_size = len(text_index)
    68.         for batch_idx in range(batch_size):
    69.             selection = np.ones(len(text_index[batch_idx]), dtype=bool)
    70.             if is_remove_duplicate:
    71.                 selection[1:] = text_index[batch_idx][1:] != text_index[
    72.                                                                  batch_idx][:-1]
    73.             for ignored_token in ignored_tokens:
    74.                 selection &= text_index[batch_idx] != ignored_token
    75.  
    76.             char_list = [
    77.                 self.character[int(text_id)].replace('\n', '')
    78.                 for text_id in text_index[batch_idx][selection]
    79.             ]
    80.             if text_prob is not None:
    81.                 conf_list = text_prob[batch_idx][selection]
    82.             else:
    83.                 conf_list = [1] * len(selection)
    84.             if len(conf_list) == 0:
    85.                 conf_list = [0]
    86.  
    87.             text = ''.join(char_list)
    88.             result_list.append((text, np.mean(conf_list).tolist()))
    89.  
    90.         return result_list
    91.  
    92.     def test(self, image_path):
    93.         img_onnx = cv2.imread(image_path)
    94.         # img_onnx = cv2.resize(img_onnx, (320, 32))
    95.         # img_onnx = img_onnx.transpose((2, 0, 1)) / 255
    96.         img_onnx = self.resize_norm_img(img_onnx)
    97.         onnx_indata = img_onnx[np.newaxis, :, :, :]
    98.         onnx_indata = torch.from_numpy(onnx_indata)
    99.         # print('diff:', onnx_indata - input_data)
    100.         print('image shape: ', onnx_indata.shape)
    101.         onnx_indata = np.array(onnx_indata, dtype=np.float32)
    102.         feed_dict = self.process(self.input_names, onnx_indata)
    103.  
    104.         output_onnx = self.sess.run(self.output_names, feed_dict)
    105.         # print('output1 shape: ', output_onnx[0].shape)
    106.         # print('output1: ', output_onnx[0])
    107.         # print('output2 shape: ', output_onnx[1].shape)
    108.         # print('output2: ', output_onnx[1])
    109.  
    110.         preds_idx = output_onnx[0]
    111.         preds_prob = output_onnx[1]
    112.         post_result = self.decode(preds_idx, preds_prob, is_remove_duplicate=True)
    113.  
    114.         if isinstance(post_result, dict):
    115.             rec_info = dict()
    116.             for key in post_result:
    117.                 if len(post_result[key][0]) >= 2:
    118.                     rec_info[key] = {
    119.                         "label": post_result[key][0][0],
    120.                         "score": float(post_result[key][0][1]),
    121.                     }
    122.             print(image_path, rec_info)
    123.         else:
    124.             if len(post_result[0]) >= 2:
    125.                 # info = post_result[0][0] + "\t" + str(post_result[0][1])
    126.                 info = post_result[0][0]
    127.             print(image_path, info)
    128.  
    129.  
    130.  
    131.  
    132. if __name__=='__main__':
    133.     image_dir = "./sample/img"
    134.     onnx_file = './export_onnx/char_recognize_20230526_v1.onnx'
    135.     character_dict_path = './all_label_num_20230517.txt'
    136.  
    137.     testobj = TestOnnx(onnx_file, character_dict_path)
    138.  
    139.     files = os.listdir(image_dir)
    140.     for file in files:
    141.         image_path = os.path.join(image_dir, file)
    142.         result = testobj.test(image_path)
    143.  
    144.  
    145.  
    146.

    模型转换结束。 

  • 相关阅读:
    Mybatis--动态SQL
    Gartner“客户之声”最高分，用户体验成中国数据库一大突破口
    标准IDOC同步物料
    21-SpringSecurity
    深度学习框架【MxNet】的安装
    蓝桥杯每日一题0223.10.23
    淘宝数据采集接口
    【物联网】802.15.4简介
    FFmpeg入门详解之103：FFmpeg Nginx VLC打造M3U8直播点播
    【软件工程】【第一章概述】d2
  • 原文地址：https://blog.csdn.net/qq_22764813/article/details/133787584