Opt: use ONNX model converted from MXNet for faster OCR

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,62 @@ class AlOcr(CnOcr):
         img_list, img_widths = self._pad_arrays(img_list)
         image = cv2.hconcat(img_list)[0, :, :]
         Image.fromarray(image).show()
+
+
+class _OnnxModule:
+    """
+    Thin wrapper that presents the same `predict(sample) -> NDArray`
+    interface as an MXNet Module, so that CnOcr._predict works unchanged.
+    """
+
+    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 mx.nd.array(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).
+
+    Parameters are identical to AlOcr.
+
+    Enable by UseOcrOnnx in deploy config.
+    """
+
+    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}, '
+                           'falling back to MXNet')
+            return AlOcr._get_module(self, context)
+
+        try:
+            import onnxruntime as ort
+        except ImportError:
+            logger.warning('onnxruntime not installed, falling back to MXNet')
+            return AlOcr._get_module(self, context)
+
+        logger.info(f'Loading OCR model (ONNX): {onnx_path}')
+        so = ort.SessionOptions()
+        # intra_op_num_threads=1 avoids thread-scheduling overhead on the small GRU model (hidden=128). 
+        # Using all cores (default) is measurably slower
+        # because the per-op parallelisation gain is outweighed by thread-pool synchronisation cost.
+        so.intra_op_num_threads = 1
+        session = ort.InferenceSession(onnx_path, so, providers=['CPUExecutionProvider'])
+        return _OnnxModule(session)

module/ocr/models.py

@@ -1,65 +1,76 @@
 from module.base.decorator import cached_property
+from module.webui.setting import State
+
+
+def _ocr_class():
+    if not State.deploy_config.UseOcrOnnx:
+        from module.ocr.al_ocr import AlOcr
+        return AlOcr
+    from module.ocr.al_ocr import AlOcrOnnx
+    return AlOcrOnnx
 
 
 class OcrModel:
     @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 _ocr_class()(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 _ocr_class()(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 _ocr_class()(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 _ocr_class()(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 _ocr_class()(model_name='densenet-lite-gru', model_epoch=63,
+                            root='./bin/cnocr_models/tw', name='tw')
 
 
 OCR_MODEL = OcrModel()