Opt: use ONNX model converted from MXNet for faster OCR with MXNet fallback

module/ocr/al_ocr.py

@@ -176,7 +176,7 @@ class AlOcr(CnOcr):
         prefix = os.path.join(self._model_dir, self._model_file_prefix)
         data_names = ['data']
         data_shapes = [(data_names[0], (hp.batch_size, 1, hp.img_height, hp.img_width))]
-        logger.info('Loading OCR model: %s' % self._model_dir)  # Change log appearance.
+        logger.info('Loading OCR model (MXNET): %s' % self._model_dir)  # Change log appearance.
         mod = load_module(
             prefix,
             self._model_epoch,
@@ -235,3 +235,79 @@ class AlOcr(CnOcr):
         img_list, img_widths = self._pad_arrays(img_list)
         image = cv2.hconcat(img_list)[0, :, :]
         Image.fromarray(image).show()
+
+
+class _OnnxModule:
+    """
+    Wrapper presenting the MXNet Module `predict(sample)` interface,
+    but returning numpy arrays to avoid the MXNet NDArray round trip.
+    """
+
+    def __init__(self, session):
+        self._session = session
+        self._input_name = session.get_inputs()[0].name
+
+    def predict(self, sample):
+        import mxnet as mx
+
+        if isinstance(sample, mx.nd.NDArray):
+            sample = sample.asnumpy()
+        # sample: (batch, 1, 32, width)  ->  output: (seq_len * batch, num_classes)
+        out = self._session.run(None, {self._input_name: np.asarray(sample, dtype=np.float32)})[0]
+        return out
+
+
+class AlOcrOnnx(AlOcr):
+    """
+    Subclass of AlOcr that runs inference through ONNX Runtime instead of MXNet,
+    providing faster CPU inference with numerically identical output.
+    All pre/post-processing is inherited unchanged.
+
+    Requires onnxruntime installed and model.onnx in the model directory
+    (generated by dev_tools/convert_mxnet_to_onnx.py from the MXNet checkpoint).
+    Falls back to MXNet if either is missing.
+
+    NOTE: For fully dropping MXNet, 
+    should vendor the pure-Python components of cnocr 1.2.2,
+    and remove the mxnet/gluoncv dependencies.
+    """
+
+    def _get_module(self, context):
+        _network, self._hp = gen_network(self._model_name, self._hp, self._net_prefix)
+        onnx_path = os.path.join(self._model_dir, 'model.onnx')
+
+        if not os.path.exists(onnx_path):
+            logger.warning(f'ONNX model not found: {onnx_path}, '
+                           'fall back to MXNet')
+            return AlOcr._get_module(self, context)
+
+        try:
+            import onnxruntime as ort
+        except ImportError:
+            logger.warning('onnxruntime not installed, fall back to MXNet')
+            return AlOcr._get_module(self, context)
+
+        logger.info(f'Loading OCR model (ONNX): {onnx_path}')
+        so = ort.SessionOptions()
+        from module.webui.setting import State
+        threads = State.deploy_config.IntraOpThreads
+        so.intra_op_num_threads = threads if isinstance(threads, int) and threads > 0 else 0
+        session = ort.InferenceSession(onnx_path, so, providers=['CPUExecutionProvider'])
+        return _OnnxModule(session)
+
+    def _predict(self, sample):
+        """
+        Args:
+            sample (np.ndarray or mx.nd.NDArray): (batch, 1, 32, width), float32.
+
+        Returns:
+            np.ndarray: (seq_len * batch, num_classes)
+        """
+        if isinstance(self._mod, _OnnxModule):
+            return self._mod.predict(sample)
+        else:
+            import mxnet as mx
+            # MXNet module expects NDArray
+            if not isinstance(sample, mx.nd.NDArray):
+                sample = mx.nd.array(sample)
+            return super()._predict(sample)

module/ocr/models.py

@@ -2,64 +2,76 @@ from module.base.decorator import cached_property
 
 
 class OcrModel:
+    @cached_property
+    def _OCR(self):
+        from module.webui.setting import State
+        backend = State.deploy_config.OcrBackend
+        if backend == 'onnx':
+            from module.ocr.al_ocr import AlOcrOnnx as _OCR
+        else:
+            from module.ocr.al_ocr import AlOcr as _OCR
+        return _OCR
+
     @cached_property
     def azur_lane(self):
         # Folder: ./bin/cnocr_models/azur_lane
-        # Size: 3.25MB
+        # Size: 3.25MB (MXNet) / 3.34MB (ONNX)
         # Model: densenet-lite-gru
         # Epoch: 15
         # Validation accuracy: 99.43%
         # Font: Impact, AgencyFB-Regular, MStiffHeiHK-UltraBold
         # Charset: 0123456789ABCDEFGHIJKLMNPQRSTUVWXYZ:/- (Letter 'O' and <space> is not included)
         # _num_classes: 39
-        from module.ocr.al_ocr import AlOcr
-        return AlOcr(model_name='densenet-lite-gru', model_epoch=15, root='./bin/cnocr_models/azur_lane',
-                     name='azur_lane')
+        return self._OCR(model_name='densenet-lite-gru', model_epoch=15,
+                         root='./bin/cnocr_models/azur_lane', name='azur_lane')
 
     @cached_property
     def azur_lane_jp(self):
         # Folder: ./bin/cnocr_models/azur_lane_jp
-        # Size: 3.25MB
+        # Size: 3.25MB (MXNet) / 3.34MB (ONNX)
         # Model: densenet-lite-gru
         # Epoch: 20
         # Validation accuracy: 99.01%
         # Font: Impact, VibeMO Compressed Pro Thin, Folk R, Source Han Serif JP
         # Charset: 0123456789ABCDEFGHIJKLMNPQRSTUVWXYZ:/- (Letter 'O' and <space> is not included)
         # _num_classes: 39
-        from module.ocr.al_ocr import AlOcr
-        return AlOcr(model_name='densenet-lite-gru', model_epoch=20, root='./bin/cnocr_models/azur_lane_jp',
-                     name='azur_lane_jp')
+        return self._OCR(model_name='densenet-lite-gru', model_epoch=20,
+                         root='./bin/cnocr_models/azur_lane_jp', name='azur_lane_jp')
 
     @cached_property
     def cnocr(self):
         # Folder: ./bin/cnocr_models/cnocr
-        # Size: 9.51MB
+        # Size: 9.51MB (MXNet) / 9.75MB (ONNX)
         # Model: densenet-lite-gru
         # Epoch: 39
         # Validation accuracy: 99.04%
         # Font: Various
         # Charset: Number, English character, Chinese character, symbols, <space>
         # _num_classes: 6426
-        from module.ocr.al_ocr import AlOcr
-        return AlOcr(model_name='densenet-lite-gru', model_epoch=39, root='./bin/cnocr_models/cnocr', name='cnocr')
+        return self._OCR(model_name='densenet-lite-gru', model_epoch=39,
+                         root='./bin/cnocr_models/cnocr', name='cnocr')
 
     @cached_property
     def jp(self):
-        from module.ocr.al_ocr import AlOcr
-        return AlOcr(model_name='densenet-lite-gru', model_epoch=125, root='./bin/cnocr_models/jp', name='jp')
+        # Folder: ./bin/cnocr_models/jp
+        # Size: 6.36MB (ONNX)
+        # Model: densenet-lite-gru
+        # Epoch: 125
+        return self._OCR(model_name='densenet-lite-gru', model_epoch=125,
+                         root='./bin/cnocr_models/jp', name='jp')
 
     @cached_property
     def tw(self):
         # Folder: ./bin/cnocr_models/tw
-        # Size: 8.43MB
+        # Size: 8.43MB (MXNet) / 8.64MB (ONNX)
         # Model: densenet-lite-gru
         # Epoch: 63
         # Validation accuracy: 99.24%
         # Font: Various, 6 kinds
         # Charset: Numbers, Upper english characters, Chinese traditional characters
         # _num_classes: 5322
-        from module.ocr.al_ocr import AlOcr
-        return AlOcr(model_name='densenet-lite-gru', model_epoch=63, root='./bin/cnocr_models/tw', name='tw')
+        return self._OCR(model_name='densenet-lite-gru', model_epoch=63,
+                         root='./bin/cnocr_models/tw', name='tw')
 
 
 OCR_MODEL = OcrModel()