From 860cdf5f48d02190209de9d0d90aeacc87855414 Mon Sep 17 00:00:00 2001 From: "yuanzhi.zyz" Date: Mon, 25 Sep 2023 11:34:53 +0800 Subject: [PATCH] add onnx exporter for ocr recognition model MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 1. 添加ocr recognition相关的exporter，支持现有三类模型转onnx 2. 更新lightweight模型 Link: https://code.alibaba-inc.com/Ali-MaaS/MaaS-lib/codereview/14135723 * add ocr recognition export and update lightweight model * fix --- .../exporters/cv/ocr_recognition_exporter.py | 40 ++++++++++++++++ modelscope/models/cv/ocr_recognition/model.py | 4 +- .../modules/LightweightEdge/main_model.py | 36 ++++++-------- .../LightweightEdge/nas_block/proxyless.py | 2 +- tests/export/test_export_ocr_recognition.py | 47 +++++++++++++++++++ tests/pipelines/test_ocr_recognition.py | 4 +- 6 files changed, 107 insertions(+), 26 deletions(-) create mode 100644 modelscope/exporters/cv/ocr_recognition_exporter.py create mode 100644 tests/export/test_export_ocr_recognition.py diff --git a/modelscope/exporters/cv/ocr_recognition_exporter.py b/modelscope/exporters/cv/ocr_recognition_exporter.py new file mode 100644 index 00000000..56c5977c --- /dev/null +++ b/modelscope/exporters/cv/ocr_recognition_exporter.py @@ -0,0 +1,40 @@ +# Copyright (c) Alibaba, Inc. and its affiliates. +import os +from functools import partial +from typing import Mapping + +import numpy as np +import onnx +import torch + +from modelscope.exporters.builder import EXPORTERS +from modelscope.exporters.torch_model_exporter import TorchModelExporter +from modelscope.metainfo import Models +from modelscope.utils.constant import ModelFile, Tasks + + +@EXPORTERS.register_module( + Tasks.ocr_recognition, module_name=Models.ocr_recognition) +class OCRRecognitionExporter(TorchModelExporter): + + def export_onnx(self, + output_dir: str, + opset=11, + input_shape=(1, 3, 32, 640)): + onnx_file = os.path.join(output_dir, ModelFile.ONNX_MODEL_FILE) + dummy_input = torch.randn(*input_shape) + self.model.onnx_export = True + self.model.eval() + _ = self.model(dummy_input) + torch.onnx._export( + self.model, + dummy_input, + onnx_file, + input_names=[ + 'images', + ], + output_names=[ + 'pred', + ], + opset_version=opset) + return {'model', onnx_file} diff --git a/modelscope/models/cv/ocr_recognition/model.py b/modelscope/models/cv/ocr_recognition/model.py index 3510de45..4c5aa362 100644 --- a/modelscope/models/cv/ocr_recognition/model.py +++ b/modelscope/models/cv/ocr_recognition/model.py @@ -109,8 +109,8 @@ class OCRRecognition(TorchModel): with open(dict_path, 'r', encoding='utf-8') as f: lines = f.readlines() cnt = 1 - # ConvNextViT model start from index=2 - if self.do_chunking: + # ConvNextViT and LightweightEdge model start from index=2 + if cfgs.model.recognizer == 'ConvNextViT' or cfgs.model.recognizer == 'LightweightEdge': cnt += 1 for line in lines: line = line.strip('\n') diff --git a/modelscope/models/cv/ocr_recognition/modules/LightweightEdge/main_model.py b/modelscope/models/cv/ocr_recognition/modules/LightweightEdge/main_model.py index 08584b5c..ce5159a3 100644 --- a/modelscope/models/cv/ocr_recognition/modules/LightweightEdge/main_model.py +++ b/modelscope/models/cv/ocr_recognition/modules/LightweightEdge/main_model.py @@ -2,6 +2,7 @@ from collections import OrderedDict +import torch import torch.nn as nn from .nas_block import plnas_linear_mix_se @@ -16,27 +17,20 @@ class LightweightEdge(nn.Module): def __init__(self): super(LightweightEdge, self).__init__() - self.FeatureExtraction = plnas_linear_mix_se(3, 123) - self.AdaptiveAvgPool = nn.AdaptiveAvgPool2d( - (None, 1)) # Transform final (imgH/16-1) -> 1 - self.dropout = nn.Dropout(0.3) - self.Prediction = nn.Sequential( - OrderedDict([ - ('fc1', nn.Linear(123, 120)), - ('bn', nn.BatchNorm1d(120)), - ('fc2', nn.Linear(120, 7642)), - ])) + self.our_nas_model = plnas_linear_mix_se(1, 128) + self.embed_dim = 128 + self.head = nn.Linear(self.embed_dim, 7644) def forward(self, input): - visual_feature = self.FeatureExtraction(input) - visual_feature = self.AdaptiveAvgPool( - visual_feature.permute(0, 3, 1, 2)) # [b, c, h, w] -> [b, w, c, h] - visual_feature = visual_feature.squeeze(3) - visual_feature = self.dropout(visual_feature) - prediction = self.Prediction.fc1(visual_feature.contiguous()) - b, t, c = prediction.shape - prediction = self.Prediction.bn(prediction.view(b * t, - c)).view(b, t, c) - prediction = self.Prediction.fc2(prediction) - + # RGB2GRAY + input = input[:, 0: + 1, :, :] * 0.2989 + input[:, 1: + 2, :, :] * 0.5870 + input[:, 2: + 3, :, :] * 0.1140 + x = self.our_nas_model(input) + x = torch.squeeze(x, 2) + x = torch.transpose(x, 1, 2) + b, s, e = x.size() + x = x.reshape(b * s, e) + prediction = self.head(x).view(b, s, -1) return prediction diff --git a/modelscope/models/cv/ocr_recognition/modules/LightweightEdge/nas_block/proxyless.py b/modelscope/models/cv/ocr_recognition/modules/LightweightEdge/nas_block/proxyless.py index 4f525639..c438c6e1 100644 --- a/modelscope/models/cv/ocr_recognition/modules/LightweightEdge/nas_block/proxyless.py +++ b/modelscope/models/cv/ocr_recognition/modules/LightweightEdge/nas_block/proxyless.py @@ -126,7 +126,7 @@ def plnas_linear_mix_se(input_channel, output_channel): stride_stages = [(2, 2), (2, 1), (2, 1), (2, 1)] n_cell_stages = [5, 5, 5, 5] - width_stages = [32, 64, 96, 123] + width_stages = [32, 64, 96, 128] conv_op_ids = [ 2, 23, 24, 26, 2, 2, 11, 27, 27, 27, 27, 2, 0, 2, 16, 10, 27, 2, 2, 2, 22, 10, 27, 3 diff --git a/tests/export/test_export_ocr_recognition.py b/tests/export/test_export_ocr_recognition.py new file mode 100644 index 00000000..303275e9 --- /dev/null +++ b/tests/export/test_export_ocr_recognition.py @@ -0,0 +1,47 @@ +# Copyright (c) Alibaba, Inc. and its affiliates. +import os +import shutil +import tempfile +import unittest +from collections import OrderedDict + +from modelscope.exporters import Exporter +from modelscope.models import Model +from modelscope.utils.constant import Tasks +from modelscope.utils.test_utils import test_level + + +class TestExportOCRRecognition(unittest.TestCase): + + def setUp(self): + print(('Testing %s.%s' % (type(self).__name__, self._testMethodName))) + self.tmp_dir = tempfile.TemporaryDirectory().name + if not os.path.exists(self.tmp_dir): + os.makedirs(self.tmp_dir) + self.model_id = 'damo/cv_LightweightEdge_ocr-recognitoin-general_damo' + + @unittest.skipUnless(test_level() >= 0, 'skip test in current test level') + def test_export_ocr_detection(self): + model = Model.from_pretrained( + 'damo/cv_LightweightEdge_ocr-recognitoin-general_damo', + model_revision='v2.4.1') + Exporter.from_model(model).export_onnx( + input_shape=(1, 3, 32, 640), output_dir=self.tmp_dir) + + @unittest.skipUnless(test_level() >= 0, 'skip test in current test level') + def test_export_ocr_detection_crnn(self): + model = Model.from_pretrained( + 'damo/cv_crnn_ocr-recognition-general_damo') + Exporter.from_model(model).export_onnx( + input_shape=(1, 3, 32, 640), output_dir=self.tmp_dir) + + @unittest.skipUnless(test_level() >= 0, 'skip test in current test level') + def test_export_ocr_detection_cvit(self): + model = Model.from_pretrained( + 'damo/cv_convnextTiny_ocr-recognition-general_damo') + Exporter.from_model(model).export_onnx( + input_shape=(3, 3, 32, 300), output_dir=self.tmp_dir) + + +if __name__ == '__main__': + unittest.main() diff --git a/tests/pipelines/test_ocr_recognition.py b/tests/pipelines/test_ocr_recognition.py index 27870b10..a59612ab 100644 --- a/tests/pipelines/test_ocr_recognition.py +++ b/tests/pipelines/test_ocr_recognition.py @@ -88,7 +88,7 @@ class OCRRecognitionTest(unittest.TestCase): ocr_recognition = pipeline( Tasks.ocr_recognition, model='damo/cv_LightweightEdge_ocr-recognitoin-general_damo', - model_revision='v1.0.0') + model_revision='v2.4.1') self.pipeline_inference(ocr_recognition, self.test_image) @unittest.skipUnless(test_level() >= 1, 'skip test in current test level') @@ -165,7 +165,7 @@ class OCRRecognitionTest(unittest.TestCase): ocr_recognition = pipeline( Tasks.ocr_recognition, model='damo/cv_LightweightEdge_ocr-recognitoin-general_damo', - model_revision='v1.0.0', + model_revision='v2.4.1', device='cpu') self.pipeline_inference(ocr_recognition, self.test_image)