Merge pull request #25 from TareHimself/dev

Dev

Author: TareHimself

.gitignore

@@ -10,7 +10,7 @@ datasets/*
 # See https://help.github.com/articles/ignoring-files/ for more about ignoring files.
 
 # dependencies
-/node_modules
+*/node_modules
 /.pnp
 .pnp.js
 

README.md

@@ -1,6 +1,6 @@
 # Manga Translator
 
-<img src="assets/examples/ui.png"/>
+<img src="assets/examples/ui_09_11_23.png"/>
 
 ## Why ?
 

assets/examples/ui.png

@@ -39,6 +39,21 @@ def run_task():
 
     return wrapper
 
+def run_in_thread(func):
+    async def wrapper(*args, **kwargs):
+        loop = asyncio.get_event_loop()
+        pending_task = asyncio.Future()
+
+        def run_task():
+            nonlocal loop
+            loop.call_soon_threadsafe(pending_task.set_result, func(*args, **kwargs))
+
+        Thread(target=run_task, group=None, daemon=True).start()
+        result = await pending_task
+        return result
+
+    return wrapper
+
 
 
 def cv2_image_from_url(url: str):
@@ -88,16 +103,22 @@ def set_default_headers(self):
     def options(self):
         self.set_status(200)
 
-    @run_in_thread
-    def post(self):
+    async def post(self):
         try:
-            data = json.loads(self.request.body)
 
-            image_url = data.get('image')
+            image = self.request.files.get('file')
+
+            if image is None:
+                raise BaseException("No Image Sent")
+
+            data = json.loads(self.get_argument('data'))
+
+
             cleaner_id, cleaner_params = data.get('cleaner', 0), data.get('cleanerArgs', {})
-            image_cv2 = cv2_image_from_url(image_url)
-            result = FullConversion(ocr=CleanOcr(),cleaner=get_cleaners()[cleaner_id](**cleaner_params))([image_cv2])[0]
-            converted_pil = cv2_to_pil(result)
+            image_cv2 = pil_to_cv2(Image.open(io.BytesIO(image[0]['body'])))
+            converter = FullConversion(ocr=CleanOcr(),cleaner=get_cleaners()[cleaner_id](**cleaner_params))
+            results = await converter([image_cv2])
+            converted_pil = cv2_to_pil(results[0])
             img_byte_arr = io.BytesIO()
             converted_pil.save(img_byte_arr, format="PNG")
             # Create response given the bytes
@@ -123,13 +144,15 @@ def set_default_headers(self):
     def options(self):
         self.set_status(200)
 
-    @run_in_thread
-    def post(self):
+    async def post(self):
         try:
 
-            data = json.loads(self.request.body)
+            image = self.request.files.get('file')
 
-            image_url = data.get('image')
+            if image is None:
+                raise BaseException("No Image Sent")
+
+            data = json.loads(self.get_argument('data'))
 
             translator_id, translator_params = data.get('translator', 0), data.get('translatorArgs', {})
 
@@ -139,14 +162,14 @@ def post(self):
 
             cleaner_id, cleaner_params = data.get('cleaner', 0), data.get('cleanerArgs', {})
 
-            image_cv2 = cv2_image_from_url(image_url)
+            image_cv2 = pil_to_cv2(Image.open(io.BytesIO(image[0]['body'])))
 
             converter = FullConversion(translator=get_translators()[translator_id](**translator_params),
                                        ocr=get_ocr()[ocr_id](**ocr_params), drawer=get_drawers()[drawer_id](**drawer_params),cleaner=get_cleaners()[cleaner_id](**cleaner_params),color_detect_model=None)
 
-            result = converter([image_cv2])[0]
+            results = await converter([image_cv2])
 
-            converted_pil = cv2_to_pil(result)
+            converted_pil = cv2_to_pil(results[0])
             img_byte_arr = io.BytesIO()
             converted_pil.save(img_byte_arr, format="PNG")
             # Create response given the bytes
@@ -257,8 +280,7 @@ def set_default_headers(self):
     def options(self):
         self.set_status(200)
 
-    @run_in_thread
-    def post(self):
+    async def post(self):
         try:
 
             image = self.request.files.get('file')
@@ -270,10 +292,10 @@ def post(self):
 
             converter = FullConversion(color_detect_model=None,translator=DeepLTranslator(auth_token=os.environ['DEEPL_AUTH']),ocr=MangaOcr(),translate_free_text=True)
 
-            translated = converter([to_convert])[0]
+            translated = await converter([to_convert])
 
             # display_image(translated,"Translated")
-            converted_pil = cv2_to_pil(translated)
+            converted_pil = cv2_to_pil(translated[0])
 
             img_byte_arr = io.BytesIO()
 
@@ -316,9 +338,9 @@ async def main():
         (r"/info", BaseHandler),
         (r"/clean", CleanFromWebHandler),
         (r"/translate", TranslateFromWebHandler),
-        (r"/images/.*", ImageHandler),
+        # (r"/images/.*", ImageHandler),
         (r"/mira/translate", MiraTranslateWebHandler),
-        (r"/(.*)", UiFilesHandler, dict(build_path=build_path)),
+        # (r"/(.*)", UiFilesHandler, dict(build_path=build_path)),
     ], **settings)
     app.listen(app_port)
     webbrowser.open(f'http://localhost:{app_port}')

assets/examples/ui_09_11_23.png

@@ -0,0 +1,7 @@
+from translator.utils import resize_and_pad
+import numpy as np
+
+
+img = np.zeros((300,180,3),dtype=np.uint8)
+
+print(resize_and_pad(img,target_size=(1000,500)).shape)

server.py

@@ -1,5 +1,5 @@
 import queue
-from typing import  Union
+from typing import Union
 import numpy as np
 from PIL import Image
 from numpy import ndarray
@@ -16,9 +16,11 @@
 import sys
 import atexit
 
-class DeepFillV2Cleaner(Cleaner):
 
-    IN_PAINT_MODEL_DEVICE = torch.device("cuda" if torch.cuda.is_available() and torch.cuda.device_count() > 0 else "cpu")
+class DeepFillV2Cleaner(Cleaner):
+    IN_PAINT_MODEL_DEVICE = torch.device(
+        "cuda" if torch.cuda.is_available() and torch.cuda.device_count() > 0 else "cpu"
+    )
 
     DEFAULT_MODEL_PATH = os.path.join("models", "inpainting.pth")
 
@@ -30,22 +32,24 @@ class DeepFillV2Cleaner(Cleaner):
 
     _pending_tasks = []
 
-    _thread: Union[None,threading.Thread] = None
+    _thread: Union[None, threading.Thread] = None
 
     @staticmethod
     def get_model(path: str):
-
         if path == DeepFillV2Cleaner._model_path:
             return DeepFillV2Cleaner._model
         else:
-            DeepFillV2Cleaner._model = load_model(path,DeepFillV2Cleaner.IN_PAINT_MODEL_DEVICE)
+            DeepFillV2Cleaner._model = load_model(
+                path, DeepFillV2Cleaner.IN_PAINT_MODEL_DEVICE
+            )
             DeepFillV2Cleaner._model_path = path
             return DeepFillV2Cleaner._model
+
     @staticmethod
     def in_paint(
-            image: Image,
-            mask: Image,
-            model_path: str = DEFAULT_MODEL_PATH,
+        image: Image,
+        mask: Image,
+        model_path: str = DEFAULT_MODEL_PATH,
     ) -> Image:
         generator = DeepFillV2Cleaner.get_model(model_path)
 
@@ -60,7 +64,9 @@ def in_paint(
         pad_height = grid - h % grid if h % grid > 0 else 0
         pad_width = grid - w % grid if w % grid > 0 else 0
 
-        image = torch.nn.functional.pad(image, (0, pad_width, 0, pad_height)).unsqueeze(0)
+        image = torch.nn.functional.pad(image, (0, pad_width, 0, pad_height)).unsqueeze(
+            0
+        )
         mask = torch.nn.functional.pad(mask, (0, pad_width, 0, pad_height)).unsqueeze(0)
 
         image = (image * 2 - 1.0).to(
@@ -73,7 +79,9 @@ def in_paint(
         image_masked = image * (1.0 - mask)  # mask image
 
         ones_x = torch.ones_like(image_masked)[:, 0:1, :, :]
-        x = torch.cat([image_masked, ones_x, ones_x * mask], dim=1)  # concatenate channels
+        x = torch.cat(
+            [image_masked, ones_x, ones_x * mask], dim=1
+        )  # concatenate channels
 
         with torch.inference_mode():
             _, x_stage2 = generator(x, mask)
@@ -90,14 +98,12 @@ def in_paint(
         # img_out = img_out.crop((0, 0, w, h))
 
         return img_out
-    
+
     @staticmethod
     def _in_paint_thread():
-
         payload = DeepFillV2Cleaner._job_queue.get()
 
         while payload is not None:
-
             data = payload
             pil_image, pil_mask, model_path, callback = data
 
@@ -106,31 +112,25 @@ def _in_paint_thread():
             payload = DeepFillV2Cleaner._job_queue.get()
 
     @staticmethod
-    async def in_paint_async(image: Image,
-                            mask: Image,
-                            model_path: str = DEFAULT_MODEL_PATH):
+    async def add_in_paint_task(
+        image: np.ndarray, mask: np.ndarray, model_path: str = DEFAULT_MODEL_PATH
+    ):
         loop = asyncio.get_event_loop()
 
         pending_task = asyncio.Future()
 
         def callback(in_paint_result):
             nonlocal loop
 
-            loop.call_soon_threadsafe(pending_task.set_result, in_paint_result)
+            loop.call_soon_threadsafe(pending_task.set_result, pil_to_cv2(in_paint_result))
 
-        DeepFillV2Cleaner._job_queue.put((image, mask, model_path, callback))
+        DeepFillV2Cleaner._job_queue.put((cv2_to_pil(image), cv2_to_pil(mask), model_path, callback))
         DeepFillV2Cleaner._pending_tasks.append(pending_task)
 
         result = await pending_task
         DeepFillV2Cleaner._pending_tasks.remove(pending_task)
         return result
 
-    @staticmethod
-    def in_paint_threadsafe(image: Image,
-                            mask: Image,
-                            model_path: str = DEFAULT_MODEL_PATH):
-        return asyncio.run(DeepFillV2Cleaner.in_paint_async(image, mask, model_path))
-
     @staticmethod
     def _stop_in_paint_thread(signum, frame):
         sys.exit(signum)
@@ -147,28 +147,34 @@ def _exit_thread():
         except:
             pass
 
-
-    
-
     def __init__(self) -> None:
         super().__init__()
 
         if DeepFillV2Cleaner._thread is None:
-            DeepFillV2Cleaner._thread = threading.Thread(target=lambda : DeepFillV2Cleaner._in_paint_thread(), group=None, daemon=True)
+            DeepFillV2Cleaner._thread = threading.Thread(
+                target=lambda: DeepFillV2Cleaner._in_paint_thread(),
+                group=None,
+                daemon=True,
+            )
             DeepFillV2Cleaner._thread.start()
             atexit.register(DeepFillV2Cleaner._exit_thread)
 
-
     @staticmethod
     def get_name() -> str:
         return "Deep Fill V2"
 
-    def clean(self, frame: ndarray, mask: ndarray, detection_results: list[tuple[tuple[int,int,int,int],str,float]] = []) -> tuple[ndarray, ndarray]:
-        return in_paint_optimized(
+    async def clean(
+        self,
+        frame: ndarray,
+        mask: ndarray,
+        detection_results: list[tuple[tuple[int, int, int, int], str, float]] = [],
+    ) -> tuple[ndarray, ndarray]:
+        loop = asyncio.get_event_loop()
+        return await in_paint_optimized(
             frame,
             mask=mask,
             filtered=detection_results,  # segmentation_results.boxes.xyxy.cpu().numpy()
-            inpaint_fun= lambda frame,mask : pil_to_cv2(
-                DeepFillV2Cleaner.in_paint_threadsafe(cv2_to_pil(frame), cv2_to_pil(mask))
-            )
-        )
+            inpaint_fun=lambda frame, mask: loop.create_task(DeepFillV2Cleaner.add_in_paint_task(
+                    frame, mask
+                )),
+        )