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

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

转换代码：

import os
import sys
import yaml
import numpy as np
import cv2
import argparse
import paddle
from paddle import nn

from argparse import ArgumentParser, RawDescriptionHelpFormatter
import paddle.distributed as dist
from ppocr.postprocess import build_post_process
from ppocr.utils.save_load import init_model
from ppocr.modeling.architectures import build_model


class AttrDict(dict):
    """Single level attribute dict, NOT recursive"""

    def __init__(self, **kwargs):
        super(AttrDict, self).__init__()
        super(AttrDict, self).update(kwargs)

    def __getattr__(self, key):
        if key in self:
            return self[key]
        raise AttributeError("object has no attribute '{}'".format(key))

global_config = AttrDict()
default_config = {'Global': {'debug': False, }}

class ArgsParser(ArgumentParser):
    def __init__(self):
        super(ArgsParser, self).__init__(
            formatter_class=RawDescriptionHelpFormatter)
        # self.add_argument("-c", "--config", default='./configs/ch_PP-OCRv2_rec_idcard.yml',
        #                   help="configuration file to use")

        self.add_argument("-c", "--config", default='./configs/ch_PP-OCRv2_rec.yml',
                          help="configuration file to use")
        self.add_argument(
            "-o", "--opt", nargs='+', help="set configuration options")

    def parse_args(self, argv=None):
        args = super(ArgsParser, self).parse_args(argv)
        assert args.config is not None, \
            "Please specify --config=configure_file_path."
        args.opt = self._parse_opt(args.opt)
        return args

    def _parse_opt(self, opts):
        config = {}
        if not opts:
            return config
        for s in opts:
            s = s.strip()
            k, v = s.split('=')
            config[k] = yaml.load(v, Loader=yaml.Loader)
        return config

def merge_config(config):
    """
    Merge config into global config.
    Args:
        config (dict): Config to be merged.
    Returns: global config
    """
    for key, value in config.items():
        if "." not in key:
            if isinstance(value, dict) and key in global_config:
                global_config[key].update(value)
            else:
                global_config[key] = value
        else:
            sub_keys = key.split('.')
            assert (
                sub_keys[0] in global_config
            ), "the sub_keys can only be one of global_config: {}, but get: {}, please check your running command".format(
                global_config.keys(), sub_keys[0])
            cur = global_config[sub_keys[0]]
            for idx, sub_key in enumerate(sub_keys[1:]):
                if idx == len(sub_keys) - 2:
                    cur[sub_key] = value
                else:
                    cur = cur[sub_key]

def load_config(file_path):
    """
    Load config from yml/yaml file.
    Args:
        file_path (str): Path of the config file to be loaded.
    Returns: global config
    """
    merge_config(default_config)
    _, ext = os.path.splitext(file_path)
    assert ext in ['.yml', '.yaml'], "only support yaml files for now"
    merge_config(yaml.load(open(file_path, 'rb'), Loader=yaml.Loader))
    return global_config

def check_device(use_gpu, use_xpu=False):
    """
    Log error and exit when set use_gpu=true in paddlepaddle
    cpu version.
    """
    err = "Config {} cannot be set as true while your paddle " \
          "is not compiled with {} ! \nPlease try: \n" \
          "\t1. Install paddlepaddle to run model on {} \n" \
          "\t2. Set {} as false in config file to run " \
          "model on CPU"

    try:
        if use_gpu and use_xpu:
            print("use_xpu and use_gpu can not both be ture.")
        if use_gpu and not paddle.is_compiled_with_cuda():
            print(err.format("use_gpu", "cuda", "gpu", "use_gpu"))
            sys.exit(1)
        if use_xpu and not paddle.device.is_compiled_with_xpu():
            print(err.format("use_xpu", "xpu", "xpu", "use_xpu"))
            sys.exit(1)
    except Exception as e:
        pass

def getArgs(is_train=False):
    FLAGS = ArgsParser().parse_args()
    config = load_config(FLAGS.config)
    merge_config(FLAGS.opt)

    # check if set use_gpu=True in paddlepaddle cpu version
    use_gpu = config['Global']['use_gpu']

    use_xpu = False

    alg = config['Architecture']['algorithm']
    assert alg in [
        'EAST', 'DB', 'SAST', 'Rosetta', 'CRNN', 'STARNet', 'RARE', 'SRN',
        'CLS', 'PGNet', 'Distillation', 'NRTR', 'TableAttn', 'SAR', 'PSE',
        'SEED', 'SDMGR', 'LayoutXLM', 'LayoutLM', 'LayoutLMv2', 'PREN', 'FCE',
        'SVTR', 'ViTSTR', 'ABINet', 'DB++', 'TableMaster', 'SPIN', 'VisionLAN',
        'Gestalt', 'SLANet', 'RobustScanner'
    ]

    device = 'gpu:{}'.format(dist.ParallelEnv().dev_id) if use_gpu else 'cpu'
    check_device(use_gpu, use_xpu)

    device = paddle.set_device(device)

    config['Global']['distributed'] = dist.get_world_size() != 1

    return config, device


class CRNN(nn.Layer):
    def __init__(self, config, device):
        super(CRNN, self).__init__()
        # 定义预处理参数
        mean = (0.5, 0.5, 0.5)
        std = (0.5, 0.5, 0.5)
        self.mean = paddle.to_tensor(mean).reshape([1, 3, 1, 1])
        self.std = paddle.to_tensor(std).reshape([1, 3, 1, 1])

        self.config = config
        # build post process
        self.post_process_class = build_post_process(config['PostProcess'],
                                                     config['Global'])
        # build model
        if hasattr(self.post_process_class, 'character'):
            char_num = len(getattr(self.post_process_class, 'character'))
            if self.config['Architecture']["algorithm"] in ["Distillation",
                                                            ]:  # distillation model
                for key in self.config['Architecture']["Models"]:
                    if self.config['Architecture']['Models'][key]['Head'][
                        'name'] == 'MultiHead':  # for multi head
                        out_channels_list = {}
                        if self.config['PostProcess'][
                            'name'] == 'DistillationSARLabelDecode':
                            char_num = char_num - 2
                        out_channels_list['CTCLabelDecode'] = char_num
                        out_channels_list['SARLabelDecode'] = char_num + 2
                        self.config['Architecture']['Models'][key]['Head'][
                            'out_channels_list'] = out_channels_list
                    else:
                        self.config['Architecture']["Models"][key]["Head"][
                            'out_channels'] = char_num
            elif self.config['Architecture']['Head'][
                'name'] == 'MultiHead':  # for multi head
                out_channels_list = {}
                if self.config['PostProcess']['name'] == 'SARLabelDecode':
                    char_num = char_num - 2
                out_channels_list['CTCLabelDecode'] = char_num
                out_channels_list['SARLabelDecode'] = char_num + 2
                self.config['Architecture']['Head'][
                    'out_channels_list'] = out_channels_list
            else:  # base rec model
                self.config['Architecture']["Head"]['out_channels'] = char_num

        # 加载模型
        self.model = build_model(config['Architecture'])
        # load_model(config, self.model)
        init_model(self.config, self.model)
        self.model.eval()

    def forward(self, x):
        # x = paddle.transpose(x, [0,3,1,2])
        # x = x / 255.0
        # x = (x - self.mean) / self.std

        model_out = self.model(x)

        # return model_out
        preds_idx = model_out.argmax(axis=2, name='class').astype('float32')
        # preds_idx = model_out.argmax(axis=2, name='class')
        preds_prob = model_out.max(axis=2, name='score').astype('float32')
        return preds_idx, preds_prob

EXPORT_ONNX = True
DYNAMIC = False

if __name__ == '__main__':
    config, device = getArgs()
    model_crnn = CRNN(config, device=device)

    # 构建输入数据images:
    image_path = "1.jpg"
    img = cv2.imread(image_path)
    img = cv2.resize(img, (320, 32))
    print('input data:', img.shape)
    img = img.astype(np.float32)
    img = img.transpose((2, 0, 1)) / 255
    input_data = img[np.newaxis, :]
    print('input data:', input_data.shape)
    x = paddle.to_tensor(input_data)
    print('input data:', x.shape)

    output_idx, output_prob = model_crnn(x)
    print('output_idx: ', output_idx)
    print('output_prob: ', output_prob)

    input_spec = paddle.static.InputSpec.from_tensor(x,  name='input')
    onnx_save_path = "./export_onnx"
    if EXPORT_ONNX:
        onnx_model_name = onnx_save_path + "/char_recognize_20230526_v1"
        if DYNAMIC:
            input_spec = paddle.static.InputSpec(
                shape=[None, 32, 320, 3], dtype='float32',  name='input')

        # ONNX模型导出
        paddle.onnx.export(model_crnn, onnx_model_name, input_spec=[input_spec], opset_version=11,
                           enable_onnx_checker=True, output_spec=[output_idx, output_prob])

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

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

测试ONNX的代码：

'''
测试转出的onnx模型
'''
import cv2
import numpy as np

import torch
import onnxruntime as rt
import math
import os

class TestOnnx:
    def __init__(self, onnx_file, character_dict_path, use_space_char=True):
        self.sess = rt.InferenceSession(onnx_file)
        # 获取输入节点名称
        self.input_names = [input.name for input in self.sess.get_inputs()]
        # 获取输出节点名称
        self.output_names = [output.name for output in self.sess.get_outputs()]

        self.character = []
        self.character.append("blank")
        with open(character_dict_path, "rb") as fin:
            lines = fin.readlines()
            for line in lines:
                line = line.decode('utf-8').strip("\n").strip("\r\n")
                self.character.append(line)
        if use_space_char:
            self.character.append(" ")

    def resize_norm_img(self, img, image_shape=[3, 32, 320]):
        imgC, imgH, imgW = image_shape
        h = img.shape[0]
        w = img.shape[1]
        ratio = w / float(h)
        if math.ceil(imgH * ratio) > imgW:
            resized_w = imgW
        else:
            resized_w = int(math.ceil(imgH * ratio))
        resized_image = cv2.resize(img, (resized_w, imgH))
        resized_image = resized_image.astype('float32')
        if image_shape[0] == 1:
            resized_image = resized_image / 255
            resized_image = resized_image[np.newaxis, :]
        else:
            resized_image = resized_image.transpose((2, 0, 1)) / 255
        resized_image -= 0.5
        resized_image /= 0.5
        padding_im = np.zeros((imgC, imgH, imgW), dtype=np.float32)
        padding_im[:, :, 0:resized_w] = resized_image
        return padding_im

    # # 准备模型运行的feed_dict
    def process(self, input_names, image):
        feed_dict = dict()
        for input_name in input_names:
            feed_dict[input_name] = image

        return feed_dict

    def get_ignored_tokens(self):
        return [0]

    def decode(self, text_index, text_prob=None, is_remove_duplicate=False):
        """ convert text-index into text-label. """
        result_list = []
        ignored_tokens = self.get_ignored_tokens()
        batch_size = len(text_index)
        for batch_idx in range(batch_size):
            selection = np.ones(len(text_index[batch_idx]), dtype=bool)
            if is_remove_duplicate:
                selection[1:] = text_index[batch_idx][1:] != text_index[
                                                                 batch_idx][:-1]
            for ignored_token in ignored_tokens:
                selection &= text_index[batch_idx] != ignored_token

            char_list = [
                self.character[int(text_id)].replace('\n', '')
                for text_id in text_index[batch_idx][selection]
            ]
            if text_prob is not None:
                conf_list = text_prob[batch_idx][selection]
            else:
                conf_list = [1] * len(selection)
            if len(conf_list) == 0:
                conf_list = [0]

            text = ''.join(char_list)
            result_list.append((text, np.mean(conf_list).tolist()))

        return result_list

    def test(self, image_path):
        img_onnx = cv2.imread(image_path)
        # img_onnx = cv2.resize(img_onnx, (320, 32))
        # img_onnx = img_onnx.transpose((2, 0, 1)) / 255
        img_onnx = self.resize_norm_img(img_onnx)
        onnx_indata = img_onnx[np.newaxis, :, :, :]
        onnx_indata = torch.from_numpy(onnx_indata)
        # print('diff:', onnx_indata - input_data)
        print('image shape: ', onnx_indata.shape)
        onnx_indata = np.array(onnx_indata, dtype=np.float32)
        feed_dict = self.process(self.input_names, onnx_indata)

        output_onnx = self.sess.run(self.output_names, feed_dict)
        # print('output1 shape: ', output_onnx[0].shape)
        # print('output1: ', output_onnx[0])
        # print('output2 shape: ', output_onnx[1].shape)
        # print('output2: ', output_onnx[1])

        preds_idx = output_onnx[0]
        preds_prob = output_onnx[1]
        post_result = self.decode(preds_idx, preds_prob, is_remove_duplicate=True)

        if isinstance(post_result, dict):
            rec_info = dict()
            for key in post_result:
                if len(post_result[key][0]) >= 2:
                    rec_info[key] = {
                        "label": post_result[key][0][0],
                        "score": float(post_result[key][0][1]),
                    }
            print(image_path, rec_info)
        else:
            if len(post_result[0]) >= 2:
                # info = post_result[0][0] + "\t" + str(post_result[0][1])
                info = post_result[0][0]
            print(image_path, info)




if __name__=='__main__':
    image_dir = "./sample/img"
    onnx_file = './export_onnx/char_recognize_20230526_v1.onnx'
    character_dict_path = './all_label_num_20230517.txt'

    testobj = TestOnnx(onnx_file, character_dict_path)

    files = os.listdir(image_dir)
    for file in files:
        image_path = os.path.join(image_dir, file)
        result = testobj.test(image_path)

模型转换结束。