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
240
module/ocr/onnx_ocr.py
Normal file
240
module/ocr/onnx_ocr.py
Normal 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()
|
||||
Reference in New Issue
Block a user