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)