Upload code

module/ocr/models.py

@@ -0,0 +1,76 @@
+from pponnxcr import TextSystem as TextSystem_
+
+from module.base.decorator import 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')
+
+    @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()