1
0
mirror of https://github.com/sui-feng-cb/AzurLaneAutoScript1.git synced 2026-08-18 20:45:58 +08:00

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>
This commit is contained in:
positnuec
2026-08-16 21:56:43 +08:00
parent b1100369f7
commit 239f252881
17 changed files with 338 additions and 166 deletions

View File

@@ -65,7 +65,7 @@ Deploy:
InstallUiautomator2: true InstallUiautomator2: true
Ocr: 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 # Whether to use ocr server
# [Default] false # [Default] false
@@ -79,14 +79,13 @@ Deploy:
# Address of ocr server for alas instance to connect # Address of ocr server for alas instance to connect
# [Default] 127.0.0.1:22268 # [Default] 127.0.0.1:22268
OcrClientAddress: 127.0.0.1:22268 OcrClientAddress: 127.0.0.1:22268
# Use higher-performance backend for OCR instead of MXNet # Backend used for OCR
# Available backends: onnx # Available backends: mxnet, onnx
# [Default] onnx # [Default] onnx
OcrBackend: onnx OcrBackend: onnx
# Number of threads used for intra-op parallelism in the ONNX session # Number of threads used for intra-op parallelism in ONNX Runtime
# Leave as default if you are unsure of the impact
# [Default] null # [Default] null
IntraOpThreads: null OnnxIntraOpThreads: null
Update: Update:
# Use auto update and builtin updater feature # Use auto update and builtin updater feature

View File

@@ -65,7 +65,7 @@ Deploy:
InstallUiautomator2: true InstallUiautomator2: true
Ocr: 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 # Whether to use ocr server
# [Default] false # [Default] false
@@ -79,14 +79,13 @@ Deploy:
# Address of ocr server for alas instance to connect # Address of ocr server for alas instance to connect
# [Default] 127.0.0.1:22268 # [Default] 127.0.0.1:22268
OcrClientAddress: 127.0.0.1:22268 OcrClientAddress: 127.0.0.1:22268
# Use higher-performance backend for OCR instead of MXNet # Backend used for OCR
# Available backends: onnx # Available backends: mxnet, onnx
# [Default] onnx # [Default] onnx
OcrBackend: onnx OcrBackend: onnx
# Number of threads used for intra-op parallelism in the ONNX session # Number of threads used for intra-op parallelism in ONNX Runtime
# Leave as default if you are unsure of the impact
# [Default] null # [Default] null
IntraOpThreads: null OnnxIntraOpThreads: null
Update: Update:
# Use auto update and builtin updater feature # Use auto update and builtin updater feature

View File

@@ -65,7 +65,7 @@ Deploy:
InstallUiautomator2: true InstallUiautomator2: true
Ocr: 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 # Whether to use ocr server
# [Default] false # [Default] false
@@ -79,14 +79,13 @@ Deploy:
# Address of ocr server for alas instance to connect # Address of ocr server for alas instance to connect
# [Default] 127.0.0.1:22268 # [Default] 127.0.0.1:22268
OcrClientAddress: 127.0.0.1:22268 OcrClientAddress: 127.0.0.1:22268
# Use higher-performance backend for OCR instead of MXNet # Backend used for OCR
# Available backends: onnx # Available backends: mxnet, onnx
# [Default] onnx # [Default] onnx
OcrBackend: onnx OcrBackend: onnx
# Number of threads used for intra-op parallelism in the ONNX session # Number of threads used for intra-op parallelism in ONNX Runtime
# Leave as default if you are unsure of the impact
# [Default] null # [Default] null
IntraOpThreads: null OnnxIntraOpThreads: null
Update: Update:
# Use auto update and builtin updater feature # Use auto update and builtin updater feature

View File

@@ -65,7 +65,7 @@ Deploy:
InstallUiautomator2: true InstallUiautomator2: true
Ocr: 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 # Whether to use ocr server
# [Default] false # [Default] false
@@ -79,14 +79,13 @@ Deploy:
# Address of ocr server for alas instance to connect # Address of ocr server for alas instance to connect
# [Default] 127.0.0.1:22268 # [Default] 127.0.0.1:22268
OcrClientAddress: 127.0.0.1:22268 OcrClientAddress: 127.0.0.1:22268
# Use higher-performance backend for OCR instead of MXNet # Backend used for OCR
# Available backends: onnx # Available backends: mxnet, onnx
# [Default] onnx # [Default] onnx
OcrBackend: onnx OcrBackend: onnx
# Number of threads used for intra-op parallelism in the ONNX session # Number of threads used for intra-op parallelism in ONNX Runtime
# Leave as default if you are unsure of the impact
# [Default] null # [Default] null
IntraOpThreads: null OnnxIntraOpThreads: null
Update: Update:
# Use auto update and builtin updater feature # Use auto update and builtin updater feature

View File

@@ -65,7 +65,7 @@ Deploy:
InstallUiautomator2: true InstallUiautomator2: true
Ocr: 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 # Whether to use ocr server
# [Default] false # [Default] false
@@ -79,14 +79,13 @@ Deploy:
# Address of ocr server for alas instance to connect # Address of ocr server for alas instance to connect
# [Default] 127.0.0.1:22268 # [Default] 127.0.0.1:22268
OcrClientAddress: 127.0.0.1:22268 OcrClientAddress: 127.0.0.1:22268
# Use higher-performance backend for OCR instead of MXNet # Backend used for OCR
# Available backends: onnx # Available backends: mxnet, onnx
# [Default] onnx # [Default] onnx
OcrBackend: onnx OcrBackend: onnx
# Number of threads used for intra-op parallelism in the ONNX session # Number of threads used for intra-op parallelism in ONNX Runtime
# Leave as default if you are unsure of the impact
# [Default] null # [Default] null
IntraOpThreads: null OnnxIntraOpThreads: null
Update: Update:
# Use auto update and builtin updater feature # Use auto update and builtin updater feature

View File

@@ -65,7 +65,7 @@ Deploy:
InstallUiautomator2: true InstallUiautomator2: true
Ocr: 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 # Whether to use ocr server
# [Default] false # [Default] false
@@ -79,14 +79,13 @@ Deploy:
# Address of ocr server for alas instance to connect # Address of ocr server for alas instance to connect
# [Default] 127.0.0.1:22268 # [Default] 127.0.0.1:22268
OcrClientAddress: 127.0.0.1:22268 OcrClientAddress: 127.0.0.1:22268
# Use higher-performance backend for OCR instead of MXNet # Backend used for OCR
# Available backends: onnx # Available backends: mxnet, onnx
# [Default] onnx # [Default] onnx
OcrBackend: onnx OcrBackend: onnx
# Number of threads used for intra-op parallelism in the ONNX session # Number of threads used for intra-op parallelism in ONNX Runtime
# Leave as default if you are unsure of the impact
# [Default] null # [Default] null
IntraOpThreads: null OnnxIntraOpThreads: null
Update: Update:
# Use auto update and builtin updater feature # Use auto update and builtin updater feature

View File

@@ -65,7 +65,7 @@ Deploy:
InstallUiautomator2: true InstallUiautomator2: true
Ocr: 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 # Whether to use ocr server
# [Default] false # [Default] false
@@ -79,14 +79,13 @@ Deploy:
# Address of ocr server for alas instance to connect # Address of ocr server for alas instance to connect
# [Default] 127.0.0.1:22268 # [Default] 127.0.0.1:22268
OcrClientAddress: 127.0.0.1:22268 OcrClientAddress: 127.0.0.1:22268
# Use higher-performance backend for OCR instead of MXNet # Backend used for OCR
# Available backends: onnx # Available backends: mxnet, onnx
# [Default] onnx # [Default] onnx
OcrBackend: onnx OcrBackend: onnx
# Number of threads used for intra-op parallelism in the ONNX session # Number of threads used for intra-op parallelism in ONNX Runtime
# Leave as default if you are unsure of the impact
# [Default] null # [Default] null
IntraOpThreads: null OnnxIntraOpThreads: null
Update: Update:
# Use auto update and builtin updater feature # Use auto update and builtin updater feature

View File

@@ -65,7 +65,7 @@ Deploy:
InstallUiautomator2: true InstallUiautomator2: true
Ocr: 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 # Whether to use ocr server
# [Default] false # [Default] false
@@ -79,14 +79,13 @@ Deploy:
# Address of ocr server for alas instance to connect # Address of ocr server for alas instance to connect
# [Default] 127.0.0.1:22268 # [Default] 127.0.0.1:22268
OcrClientAddress: 127.0.0.1:22268 OcrClientAddress: 127.0.0.1:22268
# Use higher-performance backend for OCR instead of MXNet # Backend used for OCR
# Available backends: onnx # Available backends: mxnet, onnx
# [Default] onnx # [Default] onnx
OcrBackend: onnx OcrBackend: onnx
# Number of threads used for intra-op parallelism in the ONNX session # Number of threads used for intra-op parallelism in ONNX Runtime
# Leave as default if you are unsure of the impact
# [Default] null # [Default] null
IntraOpThreads: null OnnxIntraOpThreads: null
Update: Update:
# Use auto update and builtin updater feature # Use auto update and builtin updater feature

View File

@@ -39,7 +39,7 @@ class ConfigModel:
OcrServerPort: int = 22268 OcrServerPort: int = 22268
OcrClientAddress: str = "127.0.0.1:22268" OcrClientAddress: str = "127.0.0.1:22268"
OcrBackend: str = "onnx" OcrBackend: str = "onnx"
IntraOpThreads: Optional[int] = None OnnxIntraOpThreads: Optional[int] = None
# Update # Update
EnableReload: bool = True EnableReload: bool = True

View File

@@ -65,7 +65,7 @@ Deploy:
InstallUiautomator2: true InstallUiautomator2: true
Ocr: 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 # Whether to use ocr server
# [Default] false # [Default] false
@@ -79,14 +79,13 @@ Deploy:
# Address of ocr server for alas instance to connect # Address of ocr server for alas instance to connect
# [Default] 127.0.0.1:22268 # [Default] 127.0.0.1:22268
OcrClientAddress: 127.0.0.1:22268 OcrClientAddress: 127.0.0.1:22268
# Use higher-performance backend for OCR instead of MXNet # Backend used for OCR
# Available backends: onnx # Available backends: mxnet, onnx
# [Default] onnx # [Default] onnx
OcrBackend: onnx OcrBackend: onnx
# Number of threads used for intra-op parallelism in the ONNX session # Number of threads used for intra-op parallelism in ONNX Runtime
# Leave as default if you are unsure of the impact
# [Default] null # [Default] null
IntraOpThreads: null OnnxIntraOpThreads: null
Update: Update:
# Use auto update and builtin updater feature # Use auto update and builtin updater feature

View File

@@ -36,7 +36,7 @@ class ConfigModel:
OcrServerPort: int = 22268 OcrServerPort: int = 22268
OcrClientAddress: str = "127.0.0.1:22268" OcrClientAddress: str = "127.0.0.1:22268"
OcrBackend: str = "onnx" OcrBackend: str = "onnx"
IntraOpThreads: Optional[int] = None OnnxIntraOpThreads: Optional[int] = None
# Update # Update
EnableReload: bool = True EnableReload: bool = True

View File

@@ -65,7 +65,7 @@ Deploy:
InstallUiautomator2: true InstallUiautomator2: true
Ocr: 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 # Whether to use ocr server
# [Default] false # [Default] false
@@ -79,14 +79,13 @@ Deploy:
# Address of ocr server for alas instance to connect # Address of ocr server for alas instance to connect
# [Default] 127.0.0.1:22268 # [Default] 127.0.0.1:22268
OcrClientAddress: 127.0.0.1:22268 OcrClientAddress: 127.0.0.1:22268
# Use higher-performance backend for OCR instead of MXNet # Backend used for OCR
# Available backends: onnx # Available backends: mxnet, onnx
# [Default] onnx # [Default] onnx
OcrBackend: onnx OcrBackend: onnx
# Number of threads used for intra-op parallelism in the ONNX session # Number of threads used for intra-op parallelism in ONNX Runtime
# Leave as default if you are unsure of the impact
# [Default] null # [Default] null
IntraOpThreads: null OnnxIntraOpThreads: null
Update: Update:
# Use auto update and builtin updater feature # Use auto update and builtin updater feature

View File

@@ -66,7 +66,7 @@ class ModuleBase:
def early_ocr_import(self): 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, 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. 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') logger.info('No ocr in daemon task, skip early_ocr_import')
return 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(): def do_ocr_import():
# Wait first image # Wait first image
import time import time
@@ -88,8 +93,8 @@ class ModuleBase:
time.sleep(0.01) time.sleep(0.01)
logger.info('early_ocr_import start') logger.info('early_ocr_import start')
from module.ocr.al_ocr import AlOcr from module.ocr.models import OCR_MODEL
_ = AlOcr OCR_MODEL.load()
logger.info('early_ocr_import finish') logger.info('early_ocr_import finish')
logger.info('early_ocr_import call') logger.info('early_ocr_import call')

View File

@@ -235,79 +235,3 @@ class AlOcr(CnOcr):
img_list, img_widths = self._pad_arrays(img_list) img_list, img_widths = self._pad_arrays(img_list)
image = cv2.hconcat(img_list)[0, :, :] image = cv2.hconcat(img_list)[0, :, :]
Image.fromarray(image).show() 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)

View File

@@ -4,18 +4,31 @@ from module.base.decorator import cached_property
class OcrModel: class OcrModel:
@cached_property @cached_property
def _OCR(self): def _OCR(self):
from module.exception import ScriptError
from module.webui.setting import State from module.webui.setting import State
backend = State.deploy_config.OcrBackend backend = State.deploy_config.OcrBackend
if backend == 'onnx': if backend == 'mxnet':
from module.ocr.al_ocr import AlOcrOnnx as _OCR
else:
from module.ocr.al_ocr import AlOcr as _OCR 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 return _OCR
def load(self):
_ = self._OCR
@cached_property @cached_property
def azur_lane(self): def azur_lane(self):
# Folder: ./bin/cnocr_models/azur_lane # Folder: ./bin/cnocr_models/azur_lane
# Size: 3.25MB (MXNet) / 3.34MB (ONNX) # Size: 3.25MB (MXNet) / 3.26MB (ONNX)
# Model: densenet-lite-gru # Model: densenet-lite-gru
# Epoch: 15 # Epoch: 15
# Validation accuracy: 99.43% # Validation accuracy: 99.43%
@@ -28,7 +41,7 @@ class OcrModel:
@cached_property @cached_property
def azur_lane_jp(self): def azur_lane_jp(self):
# Folder: ./bin/cnocr_models/azur_lane_jp # 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 # Model: densenet-lite-gru
# Epoch: 20 # Epoch: 20
# Validation accuracy: 99.01% # Validation accuracy: 99.01%
@@ -41,7 +54,7 @@ class OcrModel:
@cached_property @cached_property
def cnocr(self): def cnocr(self):
# Folder: ./bin/cnocr_models/cnocr # Folder: ./bin/cnocr_models/cnocr
# Size: 9.51MB (MXNet) / 9.75MB (ONNX) # Size: 9.52MB (MXNet) / 9.52MB (ONNX)
# Model: densenet-lite-gru # Model: densenet-lite-gru
# Epoch: 39 # Epoch: 39
# Validation accuracy: 99.04% # Validation accuracy: 99.04%
@@ -54,16 +67,17 @@ class OcrModel:
@cached_property @cached_property
def jp(self): def jp(self):
# Folder: ./bin/cnocr_models/jp # Folder: ./bin/cnocr_models/jp
# Size: 6.36MB (ONNX) # Size: 6.21MB (MXNet) / 6.22MB (ONNX)
# Model: densenet-lite-gru # Model: densenet-lite-gru
# Epoch: 125 # Epoch: 125
# _num_classes: 3052
return self._OCR(model_name='densenet-lite-gru', model_epoch=125, return self._OCR(model_name='densenet-lite-gru', model_epoch=125,
root='./bin/cnocr_models/jp', name='jp') root='./bin/cnocr_models/jp', name='jp')
@cached_property @cached_property
def tw(self): def tw(self):
# Folder: ./bin/cnocr_models/tw # Folder: ./bin/cnocr_models/tw
# Size: 8.43MB (MXNet) / 8.64MB (ONNX) # Size: 8.43MB (MXNet) / 8.44MB (ONNX)
# Model: densenet-lite-gru # Model: densenet-lite-gru
# Epoch: 63 # Epoch: 63
# Validation accuracy: 99.24% # Validation accuracy: 99.24%

240
module/ocr/onnx_ocr.py Normal file
View File

@@ -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()

View File

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