Feat: adopt upstream standalone ONNX OCR backend

Adopt the standalone OnnxOcr from upstream PR #5893 (381bb8d05,
by Horizon101011), which vendors cnocr's pure-Python pipeline so
the ONNX path loads no mxnet at runtime. Backend is selected via
deploy config, keeping onnxruntime a standard dependency and onnx
the default.

Co-authored-by: Horizon101011 <43370844+Horizon101011@users.noreply.github.com>

module/base/base.py

@@ -66,7 +66,7 @@ class ModuleBase:
 
     def early_ocr_import(self):
         """
-        Start a thread to import cnocr and mxnet while the Alas instance just starting to take screenshots
+        Start a thread to import OCR dependencies while the Alas instance just starting to take screenshots
         The import is paralleled since taking screenshot is I/O-bound while importing is CPU-bound,
         thus would speed up the startup 0.5 ~ 1.0s and even 5s on slow PCs.
         """
@@ -79,6 +79,11 @@ class ModuleBase:
             logger.info('No ocr in daemon task, skip early_ocr_import')
             return
 
+        from module.webui.setting import State
+        if State.deploy_config.UseOcrServer:
+            logger.info('UseOcrServer enabled, skip early_ocr_import')
+            return
+
         def do_ocr_import():
             # Wait first image
             import time
@@ -88,8 +93,8 @@ class ModuleBase:
                 time.sleep(0.01)
 
             logger.info('early_ocr_import start')
-            from module.ocr.al_ocr import AlOcr
-            _ = AlOcr
+            from module.ocr.models import OCR_MODEL
+            OCR_MODEL.load()
             logger.info('early_ocr_import finish')
 
         logger.info('early_ocr_import call')

module/ocr/al_ocr.py

@@ -235,79 +235,3 @@ 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

@@ -4,18 +4,31 @@ from module.base.decorator import cached_property
 class OcrModel:
     @cached_property
     def _OCR(self):
+        from module.exception import ScriptError
         from module.webui.setting import State
         backend = State.deploy_config.OcrBackend
-        if backend == 'onnx':
-            from module.ocr.al_ocr import AlOcrOnnx as _OCR
-        else:
+        if backend == 'mxnet':
             from module.ocr.al_ocr import AlOcr as _OCR
+        elif backend == 'onnx':
+            try:
+                import onnxruntime
+            except ImportError:
+                from module.logger import logger
+                logger.warning('ONNX Runtime is not available, fallback to MXNet')
+                from module.ocr.al_ocr import AlOcr as _OCR
+            else:
+                from module.ocr.onnx_ocr import OnnxOcr as _OCR
+        else:
+            raise ScriptError(f'Unsupported OCR backend: {backend}')
         return _OCR
 
+    def load(self):
+        _ = self._OCR
+
     @cached_property
     def azur_lane(self):
         # Folder: ./bin/cnocr_models/azur_lane
-        # Size: 3.25MB (MXNet) / 3.34MB (ONNX)
+        # Size: 3.25MB (MXNet) / 3.26MB (ONNX)
         # Model: densenet-lite-gru
         # Epoch: 15
         # Validation accuracy: 99.43%
@@ -28,7 +41,7 @@ class OcrModel:
     @cached_property
     def azur_lane_jp(self):
         # Folder: ./bin/cnocr_models/azur_lane_jp
-        # Size: 3.25MB (MXNet) / 3.34MB (ONNX)
+        # Size: 3.25MB (MXNet) / 3.26MB (ONNX)
         # Model: densenet-lite-gru
         # Epoch: 20
         # Validation accuracy: 99.01%
@@ -41,7 +54,7 @@ class OcrModel:
     @cached_property
     def cnocr(self):
         # Folder: ./bin/cnocr_models/cnocr
-        # Size: 9.51MB (MXNet) / 9.75MB (ONNX)
+        # Size: 9.52MB (MXNet) / 9.52MB (ONNX)
         # Model: densenet-lite-gru
         # Epoch: 39
         # Validation accuracy: 99.04%
@@ -54,16 +67,17 @@ class OcrModel:
     @cached_property
     def jp(self):
         # Folder: ./bin/cnocr_models/jp
-        # Size: 6.36MB (ONNX)
+        # Size: 6.21MB (MXNet) / 6.22MB (ONNX)
         # Model: densenet-lite-gru
         # Epoch: 125
+        # _num_classes: 3052
         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 (MXNet) / 8.64MB (ONNX)
+        # Size: 8.43MB (MXNet) / 8.44MB (ONNX)
         # Model: densenet-lite-gru
         # Epoch: 63
         # Validation accuracy: 99.24%

module/ocr/onnx_ocr.py

@@ -0,0 +1,240 @@
+import os
+
+import cv2
+import numpy as np
+import onnxruntime as ort
+from PIL import Image
+
+from module.exception import RequestHumanTakeover
+from module.logger import logger
+from module.webui.setting import State
+
+
+class OnnxOcr:
+    def __init__(
+            self,
+            model_name='densenet-lite-gru',
+            model_epoch=None,
+            cand_alphabet=None,
+            root=None,
+            context='cpu',
+            name=None,
+    ):
+        self._args = (model_name, model_epoch, cand_alphabet, root, context, name)
+        self._model_loaded = False
+
+    def init(
+            self,
+            model_name='densenet-lite-gru',
+            model_epoch=None,
+            cand_alphabet=None,
+            root=None,
+            context='cpu',
+            name=None,
+    ):
+        self._model_name = model_name
+        self._model_epoch = model_epoch
+        self._model_dir = root
+        self._assert_and_prepare_model_files()
+        self._alphabet, self._inv_alph_dict = self._read_charset(
+            os.path.join(self._model_dir, 'label_cn.txt')
+        )
+        self._cand_alph_idx = None
+
+        options = ort.SessionOptions()
+        threads = State.deploy_config.OnnxIntraOpThreads
+        if isinstance(threads, int) and threads > 0:
+            options.intra_op_num_threads = threads
+
+        logger.info('Loading OCR model: %s' % self._model_dir)
+        model = os.path.join(self._model_dir, 'model.onnx')
+        self._session = ort.InferenceSession(
+            model,
+            sess_options=options,
+            providers=['CPUExecutionProvider'],
+        )
+        self._input_name = self._session.get_inputs()[0].name
+        self._output_name = self._session.get_outputs()[0].name
+
+    @staticmethod
+    def _read_charset(charset_fp):
+        alphabet = [None]
+        with open(charset_fp, encoding='utf-8') as fp:
+            for line in fp:
+                alphabet.append(line.rstrip('\n'))
+        try:
+            alphabet[alphabet.index('<space>')] = ' '
+        except ValueError:
+            pass
+        inv_alph_dict = {_char: idx for idx, _char in enumerate(alphabet)}
+        return alphabet, inv_alph_dict
+
+    def _assert_and_prepare_model_files(self):
+        model_files = ['label_cn.txt', 'model.onnx']
+        for file in model_files:
+            if not os.path.exists(os.path.join(self._model_dir, file)):
+                logger.warning(f'Ocr model not prepared: {self._model_dir}')
+                logger.warning(f'Required files: {model_files}')
+                logger.critical('Please check if required files of pre-trained OCR model exist')
+                raise RequestHumanTakeover
+
+    def _ensure_loaded(self):
+        if not self._model_loaded:
+            self.init(*self._args)
+            self._model_loaded = True
+
+    def set_cand_alphabet(self, cand_alphabet):
+        self._ensure_loaded()
+        if cand_alphabet is None:
+            self._cand_alph_idx = None
+        else:
+            self._cand_alph_idx = [0] + [self._inv_alph_dict[word] for word in cand_alphabet]
+            self._cand_alph_idx.sort()
+
+    def ocr(self, img_fp):
+        self._ensure_loaded()
+        if isinstance(img_fp, str):
+            if not os.path.isfile(img_fp):
+                raise FileNotFoundError(img_fp)
+            img = np.array(Image.open(img_fp).convert('RGB'))
+        elif isinstance(img_fp, np.ndarray):
+            img = img_fp
+        else:
+            raise TypeError('Inappropriate argument type.')
+        if min(img.shape[0], img.shape[1]) < 2:
+            return ''
+        if img.mean() < 145:
+            img = 255 - img
+        line_imgs = self._line_split(img)
+        return self.ocr_for_single_lines(line_imgs)
+
+    def ocr_for_single_line(self, img_fp):
+        self._ensure_loaded()
+        if isinstance(img_fp, str):
+            if not os.path.isfile(img_fp):
+                raise FileNotFoundError(img_fp)
+            img = np.array(Image.open(img_fp).convert('L'))
+        elif isinstance(img_fp, np.ndarray):
+            img = img_fp
+        else:
+            raise TypeError('Inappropriate argument type.')
+        return self.ocr_for_single_lines([img])[0]
+
+    def ocr_for_single_lines(self, img_list):
+        self._ensure_loaded()
+        if len(img_list) == 0:
+            return []
+        img_list = [self._preprocess_img_array(img) for img in img_list]
+        batch_size = len(img_list)
+        img_list, img_widths = self._pad_arrays(img_list)
+
+        prob = self._predict(np.array(img_list, dtype=np.float32))
+        prob = np.reshape(prob, (-1, batch_size, prob.shape[1]))
+        if self._cand_alph_idx is not None:
+            prob = prob * self._gen_mask(prob.shape)
+
+        max_width = max(img_widths)
+        res = []
+        for i in range(batch_size):
+            res.append(self._gen_line_pred_chars(prob[:, i, :], img_widths[i], max_width))
+        return res
+
+    def atomic_ocr(self, img_fp, cand_alphabet=None):
+        self.set_cand_alphabet(cand_alphabet)
+        return self.ocr(img_fp)
+
+    def atomic_ocr_for_single_line(self, img_fp, cand_alphabet=None):
+        self.set_cand_alphabet(cand_alphabet)
+        return self.ocr_for_single_line(img_fp)
+
+    def atomic_ocr_for_single_lines(self, img_list, cand_alphabet=None):
+        self.set_cand_alphabet(cand_alphabet)
+        return self.ocr_for_single_lines(img_list)
+
+    @staticmethod
+    def _preprocess_img_array(img):
+        if len(img.shape) == 3 and img.shape[2] == 3:
+            if img.dtype != np.dtype('uint8'):
+                img = img.astype('uint8')
+            img = np.array(Image.fromarray(img).convert('L'))
+        new_width = int(round(32 / img.shape[0] * img.shape[1]))
+        img = cv2.resize(img, (new_width, 32))
+        img = np.expand_dims(img, 0).astype('float32') / 255.0
+        return img
+
+    @staticmethod
+    def _pad_arrays(img_list):
+        img_widths = [img.shape[2] for img in img_list]
+        if len(img_list) <= 1:
+            return img_list, img_widths
+        max_width = max(img_widths)
+        pad_width = [(0, 0), (0, 0), (0, 0)]
+        padded_img_list = []
+        for img in img_list:
+            if img.shape[2] < max_width:
+                pad_width[2] = (0, max_width - img.shape[2])
+                img = np.pad(img, pad_width, 'constant', constant_values=0.0)
+            padded_img_list.append(img)
+        return padded_img_list, img_widths
+
+    def _predict(self, sample):
+        return self._session.run([self._output_name], {self._input_name: sample})[0]
+
+    def _gen_mask(self, prob_shape):
+        mask_shape = list(prob_shape)
+        mask_shape[1] = 1
+        mask = np.zeros(mask_shape, dtype='int8')
+        mask[:, :, self._cand_alph_idx] = 1
+        return mask
+
+    def _gen_line_pred_chars(self, line_prob, img_width, max_img_width):
+        class_ids = np.argmax(line_prob, axis=-1)
+        class_ids *= np.max(line_prob, axis=-1) > 0.5
+        if img_width < max_img_width:
+            end_idx = img_width // 4
+            if end_idx < len(class_ids):
+                class_ids[end_idx:] = 0
+        prediction = self._ctc_label(class_ids.tolist())
+        return [self._alphabet[p] for p in prediction]
+
+    @staticmethod
+    def _ctc_label(class_ids):
+        prediction = []
+        previous = 0
+        for current in class_ids:
+            if current != 0 and current != previous:
+                prediction.append(current)
+            previous = current
+        return prediction
+
+    @staticmethod
+    def _line_split(img):
+        image = Image.fromarray(img)
+        gray = np.array(image.convert('L'))
+        binary = gray < 145
+        project = np.sum(binary, axis=1)
+        blank = np.where(project == 0)[0]
+        if len(blank) == 0:
+            return [np.array(image)]
+
+        borders = np.concatenate(([-1], blank, [gray.shape[0]]))
+        spans = []
+        for start, end in zip(borders[:-1], borders[1:]):
+            if end - start > 10:
+                spans.append((start + 1, end))
+        if not spans:
+            return [np.array(image)]
+
+        result = []
+        for start, end in spans:
+            start = max(0, start - 2)
+            end = min(gray.shape[0], end + 2)
+            result.append(np.array(image.crop((0, start, gray.shape[1], end))))
+        return result
+
+    def debug(self, img_list):
+        self._ensure_loaded()
+        img_list = [(self._preprocess_img_array(img) * 255.0).astype(np.uint8) for img in img_list]
+        img_list, _ = self._pad_arrays(img_list)
+        image = cv2.hconcat(img_list)[0, :, :]
+        Image.fromarray(image).show()

module/ocr/rpc.py

@@ -185,7 +185,6 @@ class ModelProxyFactory:
 def start_ocr_server(port=22268):
     import zerorpc
     import zmq
-    from module.ocr.al_ocr import AlOcr
     from module.ocr.models import OcrModel
 
     class OCRServer(OcrModel):
@@ -194,42 +193,42 @@ def start_ocr_server(port=22268):
 
         def ocr(self, lang, img_fp):
             img_fp = pickle.loads(img_fp)
-            cnocr: AlOcr = self.__getattribute__(lang)
-            return cnocr.ocr(img_fp)
+            model = self.__getattribute__(lang)
+            return model.ocr(img_fp)
 
         def ocr_for_single_line(self, lang, img_fp):
             img_fp = pickle.loads(img_fp)
-            cnocr: AlOcr = self.__getattribute__(lang)
-            return cnocr.ocr_for_single_line(img_fp)
+            model = self.__getattribute__(lang)
+            return model.ocr_for_single_line(img_fp)
 
         def ocr_for_single_lines(self, lang, img_list):
             img_list = [pickle.loads(img_fp) for img_fp in img_list]
-            cnocr: AlOcr = self.__getattribute__(lang)
-            return cnocr.ocr_for_single_lines(img_list)
+            model = self.__getattribute__(lang)
+            return model.ocr_for_single_lines(img_list)
 
         def set_cand_alphabet(self, lang, cand_alphabet):
-            cnocr: AlOcr = self.__getattribute__(lang)
-            return cnocr.set_cand_alphabet(cand_alphabet)
+            model = self.__getattribute__(lang)
+            return model.set_cand_alphabet(cand_alphabet)
 
         def atomic_ocr(self, lang, img_fp, cand_alphabet):
             img_fp = pickle.loads(img_fp)
-            cnocr: AlOcr = self.__getattribute__(lang)
-            return cnocr.atomic_ocr(img_fp, cand_alphabet)
+            model = self.__getattribute__(lang)
+            return model.atomic_ocr(img_fp, cand_alphabet)
 
         def atomic_ocr_for_single_line(self, lang, img_fp, cand_alphabet):
             img_fp = pickle.loads(img_fp)
-            cnocr: AlOcr = self.__getattribute__(lang)
-            return cnocr.atomic_ocr_for_single_line(img_fp, cand_alphabet)
+            model = self.__getattribute__(lang)
+            return model.atomic_ocr_for_single_line(img_fp, cand_alphabet)
 
         def atomic_ocr_for_single_lines(self, lang, img_list, cand_alphabet):
             img_list = [pickle.loads(img_fp) for img_fp in img_list]
-            cnocr: AlOcr = self.__getattribute__(lang)
-            return cnocr.atomic_ocr_for_single_lines(img_list, cand_alphabet)
+            model = self.__getattribute__(lang)
+            return model.atomic_ocr_for_single_lines(img_list, cand_alphabet)
 
         def debug(self, lang, img_list):
             img_list = [pickle.loads(img_fp) for img_fp in img_list]
-            cnocr: AlOcr = self.__getattribute__(lang)
-            return cnocr.debug(img_list)
+            model = self.__getattribute__(lang)
+            return model.debug(img_list)
 
     server = zerorpc.Server(OCRServer())
     try: