diff --git a/config/deploy.template-AidLux-cn.yaml b/config/deploy.template-AidLux-cn.yaml index 8a57deb2a..da2be4a16 100644 --- a/config/deploy.template-AidLux-cn.yaml +++ b/config/deploy.template-AidLux-cn.yaml @@ -65,7 +65,7 @@ Deploy: InstallUiautomator2: true Ocr: - # Run Ocr as a service, can reduce memory usage by not import mxnet everytime you start an alas instance + # Run Ocr as a service, can reduce memory usage by not import OCR dependencies everytime you start an alas instance # Whether to use ocr server # [Default] false @@ -79,14 +79,13 @@ Deploy: # Address of ocr server for alas instance to connect # [Default] 127.0.0.1:22268 OcrClientAddress: 127.0.0.1:22268 - # Use higher-performance backend for OCR instead of MXNet - # Available backends: onnx + # Backend used for OCR + # Available backends: mxnet, onnx # [Default] onnx OcrBackend: onnx - # Number of threads used for intra-op parallelism in the ONNX session - # Leave as default if you are unsure of the impact + # Number of threads used for intra-op parallelism in ONNX Runtime # [Default] null - IntraOpThreads: null + OnnxIntraOpThreads: null Update: # Use auto update and builtin updater feature diff --git a/config/deploy.template-AidLux.yaml b/config/deploy.template-AidLux.yaml index 6423776e1..ac014f80c 100644 --- a/config/deploy.template-AidLux.yaml +++ b/config/deploy.template-AidLux.yaml @@ -65,7 +65,7 @@ Deploy: InstallUiautomator2: true Ocr: - # Run Ocr as a service, can reduce memory usage by not import mxnet everytime you start an alas instance + # Run Ocr as a service, can reduce memory usage by not import OCR dependencies everytime you start an alas instance # Whether to use ocr server # [Default] false @@ -79,14 +79,13 @@ Deploy: # Address of ocr server for alas instance to connect # [Default] 127.0.0.1:22268 OcrClientAddress: 127.0.0.1:22268 - # Use higher-performance backend for OCR instead of MXNet - # Available backends: onnx + # Backend used for OCR + # Available backends: mxnet, onnx # [Default] onnx OcrBackend: onnx - # Number of threads used for intra-op parallelism in the ONNX session - # Leave as default if you are unsure of the impact + # Number of threads used for intra-op parallelism in ONNX Runtime # [Default] null - IntraOpThreads: null + OnnxIntraOpThreads: null Update: # Use auto update and builtin updater feature diff --git a/config/deploy.template-cn.yaml b/config/deploy.template-cn.yaml index 6bcadc0a9..7de453a39 100644 --- a/config/deploy.template-cn.yaml +++ b/config/deploy.template-cn.yaml @@ -65,7 +65,7 @@ Deploy: InstallUiautomator2: true Ocr: - # Run Ocr as a service, can reduce memory usage by not import mxnet everytime you start an alas instance + # Run Ocr as a service, can reduce memory usage by not import OCR dependencies everytime you start an alas instance # Whether to use ocr server # [Default] false @@ -79,14 +79,13 @@ Deploy: # Address of ocr server for alas instance to connect # [Default] 127.0.0.1:22268 OcrClientAddress: 127.0.0.1:22268 - # Use higher-performance backend for OCR instead of MXNet - # Available backends: onnx + # Backend used for OCR + # Available backends: mxnet, onnx # [Default] onnx OcrBackend: onnx - # Number of threads used for intra-op parallelism in the ONNX session - # Leave as default if you are unsure of the impact + # Number of threads used for intra-op parallelism in ONNX Runtime # [Default] null - IntraOpThreads: null + OnnxIntraOpThreads: null Update: # Use auto update and builtin updater feature diff --git a/config/deploy.template-docker-cn.yaml b/config/deploy.template-docker-cn.yaml index 97df70b79..ab74bca96 100644 --- a/config/deploy.template-docker-cn.yaml +++ b/config/deploy.template-docker-cn.yaml @@ -65,7 +65,7 @@ Deploy: InstallUiautomator2: true Ocr: - # Run Ocr as a service, can reduce memory usage by not import mxnet everytime you start an alas instance + # Run Ocr as a service, can reduce memory usage by not import OCR dependencies everytime you start an alas instance # Whether to use ocr server # [Default] false @@ -79,14 +79,13 @@ Deploy: # Address of ocr server for alas instance to connect # [Default] 127.0.0.1:22268 OcrClientAddress: 127.0.0.1:22268 - # Use higher-performance backend for OCR instead of MXNet - # Available backends: onnx + # Backend used for OCR + # Available backends: mxnet, onnx # [Default] onnx OcrBackend: onnx - # Number of threads used for intra-op parallelism in the ONNX session - # Leave as default if you are unsure of the impact + # Number of threads used for intra-op parallelism in ONNX Runtime # [Default] null - IntraOpThreads: null + OnnxIntraOpThreads: null Update: # Use auto update and builtin updater feature diff --git a/config/deploy.template-docker.yaml b/config/deploy.template-docker.yaml index f9eadf8e2..325d466df 100644 --- a/config/deploy.template-docker.yaml +++ b/config/deploy.template-docker.yaml @@ -65,7 +65,7 @@ Deploy: InstallUiautomator2: true Ocr: - # Run Ocr as a service, can reduce memory usage by not import mxnet everytime you start an alas instance + # Run Ocr as a service, can reduce memory usage by not import OCR dependencies everytime you start an alas instance # Whether to use ocr server # [Default] false @@ -79,14 +79,13 @@ Deploy: # Address of ocr server for alas instance to connect # [Default] 127.0.0.1:22268 OcrClientAddress: 127.0.0.1:22268 - # Use higher-performance backend for OCR instead of MXNet - # Available backends: onnx + # Backend used for OCR + # Available backends: mxnet, onnx # [Default] onnx OcrBackend: onnx - # Number of threads used for intra-op parallelism in the ONNX session - # Leave as default if you are unsure of the impact + # Number of threads used for intra-op parallelism in ONNX Runtime # [Default] null - IntraOpThreads: null + OnnxIntraOpThreads: null Update: # Use auto update and builtin updater feature diff --git a/config/deploy.template-linux-cn.yaml b/config/deploy.template-linux-cn.yaml index e935ec259..e5d5f96bd 100644 --- a/config/deploy.template-linux-cn.yaml +++ b/config/deploy.template-linux-cn.yaml @@ -65,7 +65,7 @@ Deploy: InstallUiautomator2: true Ocr: - # Run Ocr as a service, can reduce memory usage by not import mxnet everytime you start an alas instance + # Run Ocr as a service, can reduce memory usage by not import OCR dependencies everytime you start an alas instance # Whether to use ocr server # [Default] false @@ -79,14 +79,13 @@ Deploy: # Address of ocr server for alas instance to connect # [Default] 127.0.0.1:22268 OcrClientAddress: 127.0.0.1:22268 - # Use higher-performance backend for OCR instead of MXNet - # Available backends: onnx + # Backend used for OCR + # Available backends: mxnet, onnx # [Default] onnx OcrBackend: onnx - # Number of threads used for intra-op parallelism in the ONNX session - # Leave as default if you are unsure of the impact + # Number of threads used for intra-op parallelism in ONNX Runtime # [Default] null - IntraOpThreads: null + OnnxIntraOpThreads: null Update: # Use auto update and builtin updater feature diff --git a/config/deploy.template-linux.yaml b/config/deploy.template-linux.yaml index 6a9662449..f94669bac 100644 --- a/config/deploy.template-linux.yaml +++ b/config/deploy.template-linux.yaml @@ -65,7 +65,7 @@ Deploy: InstallUiautomator2: true Ocr: - # Run Ocr as a service, can reduce memory usage by not import mxnet everytime you start an alas instance + # Run Ocr as a service, can reduce memory usage by not import OCR dependencies everytime you start an alas instance # Whether to use ocr server # [Default] false @@ -79,14 +79,13 @@ Deploy: # Address of ocr server for alas instance to connect # [Default] 127.0.0.1:22268 OcrClientAddress: 127.0.0.1:22268 - # Use higher-performance backend for OCR instead of MXNet - # Available backends: onnx + # Backend used for OCR + # Available backends: mxnet, onnx # [Default] onnx OcrBackend: onnx - # Number of threads used for intra-op parallelism in the ONNX session - # Leave as default if you are unsure of the impact + # Number of threads used for intra-op parallelism in ONNX Runtime # [Default] null - IntraOpThreads: null + OnnxIntraOpThreads: null Update: # Use auto update and builtin updater feature diff --git a/config/deploy.template.yaml b/config/deploy.template.yaml index 7d9d38320..fb34bb16f 100644 --- a/config/deploy.template.yaml +++ b/config/deploy.template.yaml @@ -65,7 +65,7 @@ Deploy: InstallUiautomator2: true Ocr: - # Run Ocr as a service, can reduce memory usage by not import mxnet everytime you start an alas instance + # Run Ocr as a service, can reduce memory usage by not import OCR dependencies everytime you start an alas instance # Whether to use ocr server # [Default] false @@ -79,14 +79,13 @@ Deploy: # Address of ocr server for alas instance to connect # [Default] 127.0.0.1:22268 OcrClientAddress: 127.0.0.1:22268 - # Use higher-performance backend for OCR instead of MXNet - # Available backends: onnx + # Backend used for OCR + # Available backends: mxnet, onnx # [Default] onnx OcrBackend: onnx - # Number of threads used for intra-op parallelism in the ONNX session - # Leave as default if you are unsure of the impact + # Number of threads used for intra-op parallelism in ONNX Runtime # [Default] null - IntraOpThreads: null + OnnxIntraOpThreads: null Update: # Use auto update and builtin updater feature diff --git a/deploy/Windows/config.py b/deploy/Windows/config.py index 3338808c1..2d62cc2d8 100644 --- a/deploy/Windows/config.py +++ b/deploy/Windows/config.py @@ -39,7 +39,7 @@ class ConfigModel: OcrServerPort: int = 22268 OcrClientAddress: str = "127.0.0.1:22268" OcrBackend: str = "onnx" - IntraOpThreads: Optional[int] = None + OnnxIntraOpThreads: Optional[int] = None # Update EnableReload: bool = True diff --git a/deploy/Windows/template.yaml b/deploy/Windows/template.yaml index f7983295c..4e9b66e63 100644 --- a/deploy/Windows/template.yaml +++ b/deploy/Windows/template.yaml @@ -65,7 +65,7 @@ Deploy: InstallUiautomator2: true Ocr: - # Run Ocr as a service, can reduce memory usage by not import mxnet everytime you start an alas instance + # Run Ocr as a service, can reduce memory usage by not import OCR dependencies everytime you start an alas instance # Whether to use ocr server # [Default] false @@ -79,14 +79,13 @@ Deploy: # Address of ocr server for alas instance to connect # [Default] 127.0.0.1:22268 OcrClientAddress: 127.0.0.1:22268 - # Use higher-performance backend for OCR instead of MXNet - # Available backends: onnx + # Backend used for OCR + # Available backends: mxnet, onnx # [Default] onnx OcrBackend: onnx - # Number of threads used for intra-op parallelism in the ONNX session - # Leave as default if you are unsure of the impact + # Number of threads used for intra-op parallelism in ONNX Runtime # [Default] null - IntraOpThreads: null + OnnxIntraOpThreads: null Update: # Use auto update and builtin updater feature diff --git a/deploy/config.py b/deploy/config.py index 121b51c15..38ba3408c 100644 --- a/deploy/config.py +++ b/deploy/config.py @@ -36,7 +36,7 @@ class ConfigModel: OcrServerPort: int = 22268 OcrClientAddress: str = "127.0.0.1:22268" OcrBackend: str = "onnx" - IntraOpThreads: Optional[int] = None + OnnxIntraOpThreads: Optional[int] = None # Update EnableReload: bool = True diff --git a/deploy/template b/deploy/template index 848762604..27e1e65ba 100644 --- a/deploy/template +++ b/deploy/template @@ -65,7 +65,7 @@ Deploy: InstallUiautomator2: true Ocr: - # Run Ocr as a service, can reduce memory usage by not import mxnet everytime you start an alas instance + # Run Ocr as a service, can reduce memory usage by not import OCR dependencies everytime you start an alas instance # Whether to use ocr server # [Default] false @@ -79,14 +79,13 @@ Deploy: # Address of ocr server for alas instance to connect # [Default] 127.0.0.1:22268 OcrClientAddress: 127.0.0.1:22268 - # Use higher-performance backend for OCR instead of MXNet - # Available backends: onnx + # Backend used for OCR + # Available backends: mxnet, onnx # [Default] onnx OcrBackend: onnx - # Number of threads used for intra-op parallelism in the ONNX session - # Leave as default if you are unsure of the impact + # Number of threads used for intra-op parallelism in ONNX Runtime # [Default] null - IntraOpThreads: null + OnnxIntraOpThreads: null Update: # Use auto update and builtin updater feature diff --git a/module/base/base.py b/module/base/base.py index e1e3d015e..9bea3898f 100644 --- a/module/base/base.py +++ b/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') diff --git a/module/ocr/al_ocr.py b/module/ocr/al_ocr.py index 9cf26fe09..bd907e768 100644 --- a/module/ocr/al_ocr.py +++ b/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) diff --git a/module/ocr/models.py b/module/ocr/models.py index ec7009312..45bb4616a 100644 --- a/module/ocr/models.py +++ b/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% diff --git a/module/ocr/onnx_ocr.py b/module/ocr/onnx_ocr.py new file mode 100644 index 000000000..f5925b5dc --- /dev/null +++ b/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('')] = ' ' + 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() diff --git a/module/ocr/rpc.py b/module/ocr/rpc.py index abaaa8ff2..1cd9ca347 100644 --- a/module/ocr/rpc.py +++ b/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: