Merge pull request #9 from moeflow-com/gradio-add-translate-packager

gradio UI: export to moeflow

.gitignore

@@ -46,3 +46,4 @@ MANIFEST
 # Input and Output
 /input/
 /input-translated/
+/storage

Makefile

@@ -6,6 +6,9 @@ CONDA_YML ?= conda.yaml
 default:
 	@echo Please use other targets
 
+run-gradio:
+	venv/bin/gradio gradio-multi.py
+
 run-worker:
 	conda run -n mit-py311 --no-capture-output celery --app moeflow_worker worker --queues mit --loglevel=debug --concurrency=1
 

gradio-multi.py

@@ -1,25 +1,10 @@
 import logging
-from typing import List
 import gradio as gr
-import asyncio
-from pathlib import Path
-import json
-import uuid
-from PIL import Image
-import manga_translator.detection as mit_detection
-import manga_translator.ocr as mit_ocr
-import manga_translator.textline_merge as textline_merge
-import manga_translator.utils.generic as utils_generic
-from manga_translator.gradio import (
-    mit_detect_text_default_params,
-    mit_ocr_default_params,
-    storage_dir,
-    load_model_mutex,
-    MitJSONEncoder,
+from manga_translator.moeflow import (
+    process_files,
+    export_moeflow_project,
+    is_cuda_avaiable,
 )
-from manga_translator.utils.textblock import TextBlock
-
-STORAGE_DIR_RESOLVED = storage_dir.resolve()
 
 if gr.NO_RELOAD:
     logging.basicConfig(
@@ -35,99 +20,22 @@
 logger.setLevel(logging.INFO)
 
 
-async def copy_files(gradio_temp_files: list[str]) -> list[str]:
-    new_root: Path = storage_dir / uuid.uuid4().hex
-    new_root.mkdir(parents=True, exist_ok=True)
-
-    ret: list[str] = []
-    for f in gradio_temp_files:
-        new_file = new_root / f.split("/")[-1]
-        new_file.write_bytes(Path(f).read_bytes())
-        ret.append(str(new_file.relative_to(storage_dir)))
-        logger.debug("copied %s to %s", f, new_file)
-
-    return ret
-
-
-def log_file(basename: str, result: List[TextBlock]):
-    logger.info("file: %s", basename)
-    for i, b in enumerate(result):
-        logger.info("  block %d: %s", i, b.text)
-
-
-async def process_files(
-    filename_list: list[str], detector_key: str, ocr_key: str, device: str
-) -> str:
-    path_list: list[Path] = []
-    for f in filename_list:
-        assert f
-        # p = (storage_dir / f).resolve()
-        # assert p.is_file() and STORAGE_DIR_RESOLVED in p.parents, f"illegal path: {f}"
-        path_list.append(Path(f))
-
-    with load_model_mutex:
-        await mit_detection.prepare(detector_key)
-        await mit_ocr.prepare(ocr_key, device)
-
-    result = await asyncio.gather(
-        *[process_file(p, detector_key, ocr_key, device) for p in path_list]
-    )
-
-    for r in result:
-        log_file(r["filename"], r["text_blocks"])
-
-    return json.dumps(result, cls=MitJSONEncoder)
-
-
-async def process_file(
-    img_path: Path, detector: str, ocr_key: str, device: str
-) -> dict:
-    pil_img = Image.open(img_path)
-    img, mask = utils_generic.load_image(pil_img)
-    img_w, img_h = img.shape[:2]
-
-    try:
-        # detector
-        detector_args = {
-            **mit_detect_text_default_params,
-            "detector_key": detector,
-            "device": device,
-        }
-        regions, mask_raw, mask = await mit_detection.dispatch(
-            image=img, **detector_args
-        )
-        # ocr
-        ocr_args = {**mit_ocr_default_params, "ocr_key": ocr_key, "device": device}
-        textlines = await mit_ocr.dispatch(image=img, regions=regions, **ocr_args)
-        # textline merge
-        text_blocks = await textline_merge.dispatch(
-            textlines=textlines, width=img_w, height=img_h
-        )
-    except Exception as e:
-        logger.error("error processing %s: %s", img_path, e)
-        print(e)
-        text_blocks = []
-    else:
-        logger.debug("processed %s", img_path)
-
-    return {
-        "filename": img_path.name,
-        "text_blocks": text_blocks,
-    }
-
-
 with gr.Blocks() as demo:
     file_input = gr.File(
         label="upload file",
         file_count="multiple",
         type="filepath",
     )
 
-    ocr_output = gr.JSON(
-        label="OCR output",
+    target_language_input = gr.Radio(
+        ("ENG", "CHS", "CHT", None), label="translate into language", value="CHS"
     )
 
-    device_input = gr.Radio(choices=["cpu", "cuda"], label="device", value="cuda")
+    device_input = gr.Radio(
+        choices=["cpu", "cuda"],
+        label="device",
+        value="cuda" if is_cuda_avaiable() else "cpu",
+    )
     detector_key_input = gr.Radio(
         choices=[
             "default",
@@ -141,20 +49,56 @@ async def process_file(
         label="detector",
     )
 
+    export_moeflow_project_name = gr.Text(None, label="moeflow project name")
+
     ocr_key_input = gr.Radio(
         choices=["48px", "48px_ctc", "mocr"], label="ocr", value="48px"
     )
-    run_button = gr.Button("upload + text detection + OCR + textline_merge")
+    run_button = gr.Button("run")
+
+    file_output = gr.File(label="moeflow project zip", type="filepath")
+
+    ocr_output = gr.JSON(
+        label="process result",
+    )
 
     @run_button.click(
-        inputs=[file_input, detector_key_input, ocr_key_input, device_input],
-        outputs=[ocr_output],
+        inputs=[
+            file_input,
+            detector_key_input,
+            ocr_key_input,
+            device_input,
+            target_language_input,
+        ],
+        outputs=[ocr_output, file_output],
     )
     async def on_run_button(
-        gradio_temp_files: list[str], detector_key: str, ocr_key: str, device: str
-    ) -> str:
-        res = await process_files(gradio_temp_files, detector_key, ocr_key, device)
-        return res
+        gradio_temp_files: list[str],
+        detector_key: str,
+        ocr_key: str,
+        device: str,
+        target_language: str | None,
+        export_moeflow_project_name: str | None,
+    ) -> tuple[str, bytes]:
+        res = await process_files(
+            gradio_temp_files,
+            detector_key=detector_key,
+            ocr_key=ocr_key,
+            device=device,
+            # translator_key="gpt4",
+            target_language=target_language,
+        )
+        if res:
+            moeflow_zip = str(export_moeflow_project(res, export_moeflow_project_name))
+        else:
+            moeflow_zip = None
+
+        output_json = {
+            "project_name": "unnamed",
+            "files": [f.model_dump() for f in res.files],
+        }
+
+        return output_json, moeflow_zip
 
 
 if __name__ == "__main__":

manga_translator/moeflow/__init__.py

@@ -0,0 +1,12 @@
+from .process_file import process_files
+from .export_moeflow_project import export_moeflow_project
+from ._const import is_cuda_avaiable
+
+# from .export_moeflow_project import ex
+
+__all__ = [
+    # "process_file",
+    "process_files",
+    "export_moeflow_project",
+    "is_cuda_avaiable",
+]

manga_translator/moeflow/_const.py

@@ -0,0 +1,27 @@
+import uuid
+import datetime
+from pathlib import Path
+import functools
+
+_storage_dir = Path(__file__).parent.parent.parent / "storage"
+storage_dir = _storage_dir.resolve()
+
+
+def create_unique_dir(suffix: str | None = None) -> Path:
+    if suffix is None:
+        suffix = ""
+    parts = [
+        datetime.datetime.now().strftime("%Y%m%d-%H%M%S"),
+        uuid.uuid4().hex[0:8],
+    ]
+    if suffix:
+        parts.append(suffix)
+
+    return _storage_dir / "-".join(parts)
+
+
+@functools.lru_cache(maxsize=1)
+def is_cuda_avaiable():
+    import torch
+
+    return torch.cuda.is_available()

manga_translator/gradio/detection.py → manga_translator/moeflow/detection.py

@@ -1,8 +1,7 @@
 from dataclasses import dataclass
 from typing import List, Optional
 
-from gradio_client import file
-from .json_encoder import to_json
+from .model import to_json
 import manga_translator.utils.generic as utils_generic
 import numpy as np
 

manga_translator/moeflow/export_moeflow_project.py

@@ -0,0 +1,69 @@
+from pydantic import BaseModel
+from .model import FileBatchProcessResult, FileProcessResult
+from ._const import create_unique_dir
+import zipfile
+import logging
+from pathlib import Path
+
+logger = logging.getLogger(__name__)
+logger.setLevel(logging.DEBUG)
+
+
+def export_moeflow_project(
+    process_result: FileBatchProcessResult, project_name: str | None
+) -> Path:
+    if not project_name:
+        project_name = process_result.files[0].local_path.name.rsplit(".", 1)[0]
+
+    meta_json = MoeflowProjectMeta(name=project_name, intro="").model_dump_json()
+
+    output_dir = create_unique_dir("moeflow-export")
+    output_dir.mkdir(parents=True, exist_ok=True)
+    output_file = output_dir / f"{project_name}.zip"
+    output_file.parent.mkdir(parents=True, exist_ok=True)
+
+    with zipfile.ZipFile(output_file, "w") as zf:
+        zf.writestr("project.json", meta_json)
+        zf.writestr("translations.txt", _build_file(process_result.files))
+        for f in process_result.files:
+            zf.write(f.local_path, f"images/{f.local_path.name}")
+    return output_file
+
+
+class MoeflowProjectMeta(BaseModel):
+    name: str
+    intro: str
+    output_language: str = "en"
+    default_role: str = "supporter"
+    allow_apply_type: int = 3
+    application_check_type: int = 1
+    is_need_check_application: bool = False
+    source_language: str = "ja"
+
+
+def _build_file(files: list[FileProcessResult]) -> str:
+    result: list[str] = []
+    for file in files:
+        result.append(f">>>>[{file.local_path.name}]<<<<")
+        if file.translated:
+            translated_texts: list[str] = next(iter(file.translated.values()))
+        else:
+            translated_texts = None
+        logging.debug(
+            "file: %s %s x %s", file.local_path.name, file.image_w, file.image_h
+        )
+
+        for idx, block in enumerate(file.text_blocks):
+            t = translated_texts[idx] if translated_texts else ""
+            x = (
+                block.center_x / file.image_h
+            )  # this works but IDK why. Is mit using a swapped coordinate system?
+            y = block.center_y / file.image_w
+            logging.debug(
+                "block: %s,%s / %s", block.center_x, block.center_y, block.text
+            )
+            position_type = 1
+            result.append(f"----[{idx}]----[{x},{y},{position_type}]")
+            logging.debug("serialized block: %s,%s / %s", x, y, result[-1])
+            result.append(t)
+    return "\n".join(result + [""])