Merge pull request zyddnys#746 from moeflow-com/minor-fixes

Minor fixes

.dockerignore

@@ -8,3 +8,4 @@ ocrs
 models/*
 test/testdata/bboxes
 /venv
+.git

Dockerfile

@@ -22,7 +22,7 @@ RUN apt-get remove -y g++ && \
 COPY . /app
 
 # Prepare models
-RUN python -u docker_prepare.py
+RUN python -u docker_prepare.py --continue-on-error
 
 RUN rm -rf /tmp
 

docker_prepare.py

@@ -1,28 +1,55 @@
 import asyncio
-
+from argparse import ArgumentParser
 from manga_translator.utils import ModelWrapper
 from manga_translator.detection import DETECTORS
 from manga_translator.ocr import OCRS
 from manga_translator.inpainting import INPAINTERS
 
+
+arg_parser = ArgumentParser()
+arg_parser.add_argument("--models", default="")
+arg_parser.add_argument("--continue-on-error", action="store_true")
+
+
+cli_args = arg_parser.parse_args()
+
+
 async def download(dict):
-  for key, value in dict.items():
-    if issubclass(value, ModelWrapper):
-      print(' -- Downloading', key)
-      try:
-        inst = value()
-        await inst.download()
-      except Exception as e:
-        print('Failed to download', key, value)
-        print(e)
+    """ """
+    for key, value in dict.items():
+        if issubclass(value, ModelWrapper):
+            print(" -- Downloading", key)
+            try:
+                inst = value()
+                await inst.download()
+            except Exception as e:
+                print("Failed to download", key, value)
+                print(e)
+                if not cli_args.continue_on_error:
+                    raise
+
 
 async def main():
-  await download(DETECTORS)
-  await download(OCRS)
-  await download({
-    k: v for k, v in INPAINTERS.items() 
-      if k not in ['sd']
-  })
-
-if __name__ == '__main__':
-  asyncio.run(main())
+    models: set[str] = set(filter(None, cli_args.models.split(",")))
+
+    await download(
+        {
+            k: v
+            for k, v in DETECTORS.items()
+            if (not models) or (f"detector.{k}" in models)
+        }
+    )
+    await download(
+        {k: v for k, v in OCRS.items() if (not models) or (f"ocr.{k}" in models)}
+    )
+    await download(
+        {
+            k: v
+            for k, v in INPAINTERS.items()
+            if (not models) or (f"inpaint.{k}" in models) and (k not in ["sd"])
+        }
+    )
+
+
+if __name__ == "__main__":
+    asyncio.run(main())

manga_translator/utils/generic.py

@@ -1,5 +1,5 @@
 import os
-from typing import List, Callable, Tuple
+from typing import List, Callable, Tuple, Optional
 import numpy as np
 import cv2
 import functools
@@ -65,7 +65,7 @@ def _get_args(self):
 
 # TODO: Add TranslationContext for type linting
 
-def atoi(text):
+def atoi(text: str) -> int | str:
     return int(text) if text.isdigit() else text
 
 def natural_sort(l: List[str]):
@@ -246,7 +246,7 @@ def __call__(self, val = None):
         else:
             return 0
 
-def load_image(img: Image.Image):
+def load_image(img: Image.Image) -> Tuple[np.ndarray, Optional[Image.Image]]:
     if img.mode == 'RGBA':
         # from https://stackoverflow.com/questions/9166400/convert-rgba-png-to-rgb-with-pil
         img.load()  # needed for split()

manga_translator/utils/textblock.py

@@ -83,7 +83,7 @@ def __init__(self, lines: List[Tuple[int, int, int, int]],
         if self.text and len(texts) > 1:
             for txt in texts[1:]:
                 first_cjk = '\u3000' <= self.text[-1] <= '\u9fff'
-                second_cjk = '\u3000' <= txt[0] <= '\u9fff'
+                second_cjk = txt and ('\u3000' <= txt[0] <= '\u9fff')
                 if first_cjk or second_cjk :
                     self.text += txt
                 else :