From c5feafea5ead4b2426e801b87a70909351b0d650 Mon Sep 17 00:00:00 2001 From: frederik-uni <147479464+frederik-uni@users.noreply.github.com> Date: Sat, 16 Nov 2024 23:50:18 +0100 Subject: [PATCH] update to enum --- manga_translator/colorization/__init__.py | 3 +- manga_translator/config.py | 66 ++++++++++++++++++++--- manga_translator/detection/__init__.py | 11 ++-- manga_translator/inpainting/__init__.py | 13 ++--- manga_translator/ocr/__init__.py | 9 ++-- manga_translator/translators/__init__.py | 49 ++++++++--------- manga_translator/upscaling/__init__.py | 7 +-- 7 files changed, 109 insertions(+), 49 deletions(-) diff --git a/manga_translator/colorization/__init__.py b/manga_translator/colorization/__init__.py index aef0dcd0..d7423286 100644 --- a/manga_translator/colorization/__init__.py +++ b/manga_translator/colorization/__init__.py @@ -2,9 +2,10 @@ from .common import CommonColorizer, OfflineColorizer from .manga_colorization_v2 import MangaColorizationV2 +from ..config import Colorizer COLORIZERS = { - 'mc2': MangaColorizationV2, + Colorizer.mc2: MangaColorizationV2, } colorizer_cache = {} diff --git a/manga_translator/config.py b/manga_translator/config.py index 5340b2be..30c44193 100644 --- a/manga_translator/config.py +++ b/manga_translator/config.py @@ -23,6 +23,60 @@ class InpaintPrecision(IntEnum): fp16 = 1 bf16 = 2 +class Detector(IntEnum): + default = 0 + dbconvnext = 1 + ctd = 2 + craft = 3 + none = 4 + +class Inpainter(IntEnum): + default = 0 + lama_large = 1 + lama_mpe = 2 + sd = 3 + none = 4 + original = 5 + +class Colorizer(IntEnum): + none = 0 + mc2 = 1 + +class Ocr(IntEnum): + ocr32px = 0 + ocr48px = 1 + ocr48px_ctc = 2 + mocr = 3 +class Translator(IntEnum): + youdao = 0 + baidu = 1 + deepl = 2 + papago = 3 + caiyun = 4 + gpt3 = 5 + gpt3_5 = 6 + gpt4 = 7 + none = 8 + original = 9 + sakura = 10 + deepseek = 11 + groq = 12 + offline = 13 + nllb = 14 + nllb_big = 15 + sugoi = 16 + jparacrawl = 17 + jparacrawl_big = 18 + m2m100 = 19 + m2m100_big = 20 + mbart50 = 21 + qwen2 = 22 + qwen2_big = 23 +class Upscaler: + waifu2x = 0 + esrgan = 1 + upscler4xultrasharp = 2 + class RenderConfig(BaseModel): renderer: Renderer = Renderer.default """Render english text translated from manga with some additional typesetting. Ignores some other argument options""" @@ -53,7 +107,7 @@ class RenderConfig(BaseModel): class UpscaleConfig(BaseModel): - upscaler: str = 'esrgan' #todo: validate UPSCALERS #todo: convert to enum + upscaler: Upscaler = Upscaler.esrgan """Upscaler to use. --upscale-ratio has to be set for it to take effect""" revert_upscaling: bool = False """Downscales the previously upscaled image after translation back to original size (Use with --upscale-ratio).""" @@ -61,7 +115,7 @@ class UpscaleConfig(BaseModel): """Image upscale ratio applied before detection. Can improve text detection.""" class TranslatorConfig(BaseModel): - translator: str = "google" #todo: validate TRANSLATORS todo: convert to enum + translator: Translator = Translator.sugoi """Language translator to use""" target_lang: str = 'ENG' #todo: validate VALID_LANGUAGES #todo: convert to enum """Destination language""" @@ -78,7 +132,7 @@ class TranslatorConfig(BaseModel): class DetectorConfig(BaseModel): """""" - detector: str = 'default' #todo: validate DETECTORS #todo: convert to enum + detector: Detector =Detector.default """"Text detector used for creating a text mask from an image, DO NOT use craft for manga, it\'s not designed for it""" detection_size: int = 1536 """Size of image used for detection""" @@ -96,7 +150,7 @@ class DetectorConfig(BaseModel): """The threshold for ignoring text in non bubble areas, with valid values ranging from 1 to 50, does not ignore others. Recommendation 5 to 10. If it is too low, normal bubble areas may be ignored, and if it is too large, non bubble areas may be considered normal bubbles""" class InpainterConfig(BaseModel): - inpainter: str = 'lama_large' #todo: validate INPAINTERS #todo: convert to enum + inpainter: Inpainter = Inpainter.lama_large """Inpainting model to use""" inpainting_size: int = 2048 """Size of image used for inpainting (too large will result in OOM)""" @@ -109,7 +163,7 @@ class ColorizerConfig(BaseModel): """Size of image used for colorization. Set to -1 to use full image size""" denoise_sigma: int = 30 """Used by colorizer and affects color strength, range from 0 to 255 (default 30). -1 turns it off.""" - colorizer: Optional[str] = None # todo: validate COLORIZERS #todo: convert to enum + colorizer: Colorizer = Colorizer.none """Colorization model to use.""" @@ -117,7 +171,7 @@ class ColorizerConfig(BaseModel): class OcrConfig(BaseModel): use_mocr_merge: bool = False """Use bbox merge when Manga OCR inference.""" - ocr: str = '48px' #todo: validate OCRS #todo: convert to enum + ocr: Ocr = Ocr.ocr48px """Optical character recognition (OCR) model to use""" min_text_length: int = 0 """Minimum text length of a text region""" diff --git a/manga_translator/detection/__init__.py b/manga_translator/detection/__init__.py index 96e930f5..2f8bd24e 100644 --- a/manga_translator/detection/__init__.py +++ b/manga_translator/detection/__init__.py @@ -6,13 +6,14 @@ from .craft import CRAFTDetector from .none import NoneDetector from .common import CommonDetector, OfflineDetector +from ..config import Detector DETECTORS = { - 'default': DefaultDetector, - 'dbconvnext': DBConvNextDetector, - 'ctd': ComicTextDetector, - 'craft': CRAFTDetector, - 'none': NoneDetector, + Detector.default: DefaultDetector, + Detector.dbconvnext: DBConvNextDetector, + Detector.ctd: ComicTextDetector, + Detector.craft: CRAFTDetector, + Detector.none: NoneDetector, } detector_cache = {} diff --git a/manga_translator/inpainting/__init__.py b/manga_translator/inpainting/__init__.py index 5b7a9b16..16f24212 100644 --- a/manga_translator/inpainting/__init__.py +++ b/manga_translator/inpainting/__init__.py @@ -6,14 +6,15 @@ from .inpainting_sd import StableDiffusionInpainter from .none import NoneInpainter from .original import OriginalInpainter +from ..config import Inpainter INPAINTERS = { - 'default': AotInpainter, - 'lama_large': LamaLargeInpainter, - 'lama_mpe': LamaMPEInpainter, - 'sd': StableDiffusionInpainter, - 'none': NoneInpainter, - 'original': OriginalInpainter, + Inpainter.default: AotInpainter, + Inpainter.lama_large: LamaLargeInpainter, + Inpainter.lama_mpe: LamaMPEInpainter, + Inpainter.sd: StableDiffusionInpainter, + Inpainter.none: NoneInpainter, + Inpainter.original: OriginalInpainter, } inpainter_cache = {} diff --git a/manga_translator/ocr/__init__.py b/manga_translator/ocr/__init__.py index 15d7a018..75390d76 100644 --- a/manga_translator/ocr/__init__.py +++ b/manga_translator/ocr/__init__.py @@ -6,13 +6,14 @@ from .model_48px import Model48pxOCR from .model_48px_ctc import Model48pxCTCOCR from .model_manga_ocr import ModelMangaOCR +from ..config import Ocr from ..utils import Quadrilateral OCRS = { - '32px': Model32pxOCR, - '48px': Model48pxOCR, - '48px_ctc': Model48pxCTCOCR, - 'mocr': ModelMangaOCR, + Ocr.ocr32px: Model32pxOCR, + Ocr.ocr48px: Model48pxOCR, + Ocr.ocr48px_ctc: Model48pxCTCOCR, + Ocr.mocr: ModelMangaOCR, } ocr_cache = {} diff --git a/manga_translator/translators/__init__.py b/manga_translator/translators/__init__.py index 2ff45ea9..30bff773 100644 --- a/manga_translator/translators/__init__.py +++ b/manga_translator/translators/__init__.py @@ -19,36 +19,37 @@ from .sakura import SakuraTranslator from .qwen2 import Qwen2Translator, Qwen2BigTranslator from .groq import GroqTranslator +from ..config import Translator OFFLINE_TRANSLATORS = { - 'offline': SelectiveOfflineTranslator, - 'nllb': NLLBTranslator, - 'nllb_big': NLLBBigTranslator, - 'sugoi': SugoiTranslator, - 'jparacrawl': JparacrawlTranslator, - 'jparacrawl_big': JparacrawlBigTranslator, - 'm2m100': M2M100Translator, - 'm2m100_big': M2M100BigTranslator, - 'mbart50': MBart50Translator, - 'qwen2': Qwen2Translator, - 'qwen2_big': Qwen2BigTranslator, + Translator.offline: SelectiveOfflineTranslator, + Translator.nllb: NLLBTranslator, + Translator.nllb_big: NLLBBigTranslator, + Translator.sugoi: SugoiTranslator, + Translator.jparacrawl: JparacrawlTranslator, + Translator.jparacrawl_big: JparacrawlBigTranslator, + Translator.m2m100: M2M100Translator, + Translator.m2m100_big: M2M100BigTranslator, + Translator.mbart50: MBart50Translator, + Translator.qwen2: Qwen2Translator, + Translator.qwen2_big: Qwen2BigTranslator, } TRANSLATORS = { # 'google': GoogleTranslator, - 'youdao': YoudaoTranslator, - 'baidu': BaiduTranslator, - 'deepl': DeeplTranslator, - 'papago': PapagoTranslator, - 'caiyun': CaiyunTranslator, - 'gpt3': GPT3Translator, - 'gpt3.5': GPT35TurboTranslator, - 'gpt4': GPT4Translator, - 'none': NoneTranslator, - 'original': OriginalTranslator, - 'sakura': SakuraTranslator, - 'deepseek': DeepseekTranslator, - 'groq':GroqTranslator, + Translator.youdao: YoudaoTranslator, + Translator.baidu: BaiduTranslator, + Translator.deepl: DeeplTranslator, + Translator.papago: PapagoTranslator, + Translator.caiyun: CaiyunTranslator, + Translator.gpt3: GPT3Translator, + Translator.gpt3_5: GPT35TurboTranslator, + Translator.gpt4: GPT4Translator, + Translator.none: NoneTranslator, + Translator.original: OriginalTranslator, + Translator.sakura: SakuraTranslator, + Translator.deepseek: DeepseekTranslator, + Translator.groq:GroqTranslator, **OFFLINE_TRANSLATORS, } translator_cache = {} diff --git a/manga_translator/upscaling/__init__.py b/manga_translator/upscaling/__init__.py index b5a972dd..76852789 100644 --- a/manga_translator/upscaling/__init__.py +++ b/manga_translator/upscaling/__init__.py @@ -5,11 +5,12 @@ from .waifu2x import Waifu2xUpscaler from .esrgan import ESRGANUpscaler from .esrgan_pytorch import ESRGANUpscalerPytorch +from ..config import Upscaler UPSCALERS = { - 'waifu2x': Waifu2xUpscaler, - 'esrgan': ESRGANUpscaler, - '4xultrasharp': ESRGANUpscalerPytorch, + Upscaler.waifu2x: Waifu2xUpscaler, + Upscaler.esrgan: ESRGANUpscaler, + Upscaler.upscler4xultrasharp: ESRGANUpscalerPytorch, } upscaler_cache = {}