add onnx exporter for ocr recognition model

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

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}

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')

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

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

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()

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)