mirror of
https://github.com/LmeSzinc/StarRailCopilot.git
synced 2024-11-25 10:01:10 +00:00
83 lines
1.9 KiB
Python
83 lines
1.9 KiB
Python
from pponnxcr import TextSystem as TextSystem_
|
|
|
|
from module.base.decorator import cached_property, del_cached_property
|
|
from module.exception import ScriptError
|
|
|
|
DIC_LANG_TO_MODEL = {
|
|
'cn': 'zhs',
|
|
'en': 'en',
|
|
'jp': 'ja',
|
|
'tw': 'zht',
|
|
}
|
|
|
|
|
|
def lang2model(lang: str) -> str:
|
|
"""
|
|
Args:
|
|
lang: In-game language name, defined in VALID_LANG
|
|
|
|
Returns:
|
|
str: Model name, defined in pponnxcr.utility
|
|
"""
|
|
return DIC_LANG_TO_MODEL.get(lang, lang)
|
|
|
|
|
|
def model2lang(model: str) -> str:
|
|
"""
|
|
Args:
|
|
model: Model name, defined in pponnxcr.utility
|
|
|
|
Returns:
|
|
str: In-game language name, defined in VALID_LANG
|
|
"""
|
|
for k, v in DIC_LANG_TO_MODEL.items():
|
|
if model == v:
|
|
return k
|
|
return model
|
|
|
|
|
|
class TextSystem(TextSystem_):
|
|
def __init__(self, *args, **kwargs):
|
|
super().__init__(*args, **kwargs)
|
|
self.text_recognizer.rec_batch_num = 1
|
|
|
|
|
|
class OcrModel:
|
|
def get_by_model(self, model: str) -> TextSystem:
|
|
try:
|
|
return self.__getattribute__(model)
|
|
except AttributeError:
|
|
raise ScriptError(f'OCR model "{model}" does not exists')
|
|
|
|
def get_by_lang(self, lang: str) -> TextSystem:
|
|
try:
|
|
model = lang2model(lang)
|
|
return self.__getattribute__(model)
|
|
except AttributeError:
|
|
raise ScriptError(f'OCR model under lang "{lang}" does not exists')
|
|
|
|
def resource_release(self):
|
|
del_cached_property(self, 'zhs')
|
|
del_cached_property(self, 'en')
|
|
del_cached_property(self, 'ja')
|
|
del_cached_property(self, 'zht')
|
|
|
|
@cached_property
|
|
def zhs(self):
|
|
return TextSystem('zhs')
|
|
|
|
@cached_property
|
|
def en(self):
|
|
return TextSystem('en')
|
|
|
|
@cached_property
|
|
def ja(self):
|
|
return TextSystem('ja')
|
|
|
|
@cached_property
|
|
def zht(self):
|
|
return TextSystem('zht')
|
|
|
|
|
|
OCR_MODEL = OcrModel()
|