Merge pull request #815 from k-aito/rewrite-save-text

Rewrite save text and add load text

manga_translator/args.py

@@ -117,6 +117,7 @@ def reparse(arr: list):
 parser_batch.add_argument('--use-mtpe', action='store_true', help='Turn on/off machine translation post editing (MTPE) on the command line (works only on linux right now)')
 g_batch = parser_batch.add_mutually_exclusive_group()
 g_batch.add_argument('--save-text', action='store_true', help='Save extracted text and translations into a text file.')
+g_batch.add_argument('--load-text', action='store_true', help='Load extracted text and translations from a text file.')
 g_batch.add_argument('--save-text-file', default='', type=str, help='Like --save-text but with a specified file path.')
 parser_batch.add_argument('--prep-manual', action='store_true', help='Prepare for manual typesetting by outputting blank, inpainted images, plus copies of the original for reference')
 parser_batch.add_argument('--save-quality', default=100, type=int, help='Quality of saved JPEG image, range from 0 to 100 with 100 being best')
@@ -136,4 +137,4 @@ def reparse(arr: list):
 parser_api.add_argument('--nonce', default=os.getenv('MT_WEB_NONCE', ''), type=str, help='Nonce for securing internal API server communication')
 parser_api.add_argument("--report", default=None,type=str, help='reports to server to register instance')
 
-subparsers.add_parser('config-help', help='Print help information for config file')
+subparsers.add_parser('config-help', help='Print help information for config file')

manga_translator/manga_translator.py

@@ -1,4 +1,5 @@
 import cv2
+import json
 import langcodes
 import langdetect
 import os
@@ -137,6 +138,12 @@ def parse_init_params(self, params: dict):
             ModelWrapper._MODEL_DIR = params.get('model_dir')
         #todo: fix why is kernel size loaded in the constructor
         self.kernel_size=int(params.get('kernel_size'))
+        # Set input files
+        self.input_files = params.get('input', [])
+        # Set save_text
+        self.save_text = params.get('save_text', False)
+        # Set load_text
+        self.load_text = params.get('load_text', False)
 
     @property
     def using_gpu(self):
@@ -173,7 +180,6 @@ async def translate(self, image: Image.Image, config: Config) -> Context:
         return await self._translate(config, ctx)
 
     async def _translate(self, config: Config, ctx: Context) -> Context:
-
         # -- Colorization
         if config.colorizer.colorizer != Colorizer.none:
             await self._report_progress('colorizing')
@@ -249,7 +255,6 @@ async def _translate(self, config: Config, ctx: Context) -> Context:
         ctx.text_regions = await self._run_text_translation(config, ctx)
         await self._report_progress('after-translating')
 
-
         if not ctx.text_regions:
             await self._report_progress('error-translating', True)
             ctx.result = ctx.upscaled
@@ -412,12 +417,25 @@ async def _run_textline_merge(self, config: Config, ctx: Context):
         return text_regions
 
     async def _run_text_translation(self, config: Config, ctx: Context):
-        translated_sentences = \
-            await dispatch_translation(config.translator.translator_gen,
-                                       [region.text for region in ctx.text_regions],
-                                       config.translator,
-                                       self.use_mtpe,
-                                       ctx, 'cpu' if self._gpu_limited_memory else self.device)
+        if self.load_text:
+            input_filename = os.path.splitext(os.path.basename(self.input_files[0]))[0]
+            with open(self._result_path(f"{input_filename}_translations.txt"), "r") as f:
+                    translated_sentences = json.load(f)
+        else:
+            translated_sentences = \
+                await dispatch_translation(config.translator.translator_gen,
+                                           [region.text for region in ctx.text_regions],
+                                           config.translator,
+                                           self.use_mtpe,
+                                           ctx, 'cpu' if self._gpu_limited_memory else self.device)
+
+            # Save translation if args.save_text is set and quit
+            if self.save_text:
+                input_filename = os.path.splitext(os.path.basename(self.input_files[0]))[0]
+                with open(self._result_path(f"{input_filename}_translations.txt"), "w") as f:
+                    json.dump(translated_sentences, f, indent=4)
+                print("Don't continue if --save-text is used")
+                exit(-1)
 
         for region, translation in zip(ctx.text_regions, translated_sentences):
             if config.render.uppercase: