From ba17fd3ee5b8aa4be7203d0a8d4a9ca67efb7947 Mon Sep 17 00:00:00 2001 From: Wang Guan Date: Tue, 23 Apr 2024 00:59:25 +0900 Subject: [PATCH 1/5] allow set of CELERY_BROKER_URL CELERY_BACKEND_URL this enables other Celery backend like Redis --- app/config.py | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) diff --git a/app/config.py b/app/config.py index 32c0627..54056f3 100644 --- a/app/config.py +++ b/app/config.py @@ -112,8 +112,11 @@ # ----------- # Celery # ----------- -CELERY_BROKER_URL = f"amqp://{env['RABBITMQ_USER']}:{env['RABBITMQ_PASS']}@moeflow-rabbitmq:5672/{env['RABBITMQ_VHOST_NAME']}" -CELERY_BACKEND_URL = DB_URI +CELERY_BROKER_URL = env.get( + "CELERY_BROKER_URL", + f"amqp://{env['RABBITMQ_USER']}:{env['RABBITMQ_PASS']}@moeflow-rabbitmq:5672/{env['RABBITMQ_VHOST_NAME']}", +) +CELERY_BACKEND_URL = env.get("CELERY_BACKEND_URL", DB_URI) CELERY_MONGODB_BACKEND_SETTINGS = { "database": env["MONGODB_DB_NAME"], "taskmeta_collection": "celery_taskmeta", From 846fec89f47fadfefa791dc8c7138b6e3c7b4c43 Mon Sep 17 00:00:00 2001 From: Wang Guan Date: Tue, 23 Apr 2024 01:54:36 +0900 Subject: [PATCH 2/5] provide mit-worker -based apis --- app/__init__.py | 1 + app/apis/manga_image_translator.py | 76 +++++++++++++ app/apis/urls.py | 32 +++++- app/models/target.py | 2 +- app/tasks/__init__.py | 35 ++++++ app/tasks/mit.py | 166 +++++++++++++++++++++++++++++ manage.py | 33 ++++++ 7 files changed, 343 insertions(+), 2 deletions(-) create mode 100644 app/apis/manga_image_translator.py create mode 100644 app/tasks/mit.py diff --git a/app/__init__.py b/app/__init__.py index 49abe67..b0acfd9 100644 --- a/app/__init__.py +++ b/app/__init__.py @@ -139,6 +139,7 @@ def create_celery() -> Celery: "app.tasks.ocr", "app.tasks.import_from_labelplus", "app.tasks.thumbnail", + "app.tasks.mit", # only included for completeness's sake. its impl is in other repo. ], related_name=None, ) diff --git a/app/apis/manga_image_translator.py b/app/apis/manga_image_translator.py new file mode 100644 index 0000000..b130798 --- /dev/null +++ b/app/apis/manga_image_translator.py @@ -0,0 +1,76 @@ +# Translation preprocess API backed by manga-image-translator worker +from app.core.views import MoeAPIView +from flask import request + +from app.exceptions.base import UploadFileNotFoundError +from app.tasks.mit import preprocess_mit, mit_ocr, mit_detect_text, mit_translate +from app.tasks import queue_task, wait_result_sync +from app import app_config +from werkzeug.datastructures import FileStorage +from tempfile import NamedTemporaryFile +import os + +MIT_STORAGE_ROOT = app_config.get("MIT_STORAGE_ROOT", "/MIT_STORAGE_ROOT_UNDEFINED") + + +def _wait_task_result(task_id: str): + try: + result = wait_result_sync(task_id, timeout=1) + return {"id": task_id, "status": "success", "result": result} + except TimeoutError: + return { + "id": task_id, + "status": "pending", + } + except Exception as e: + return {"id": task_id, "status": "fail", "message": str(e)} + + +class MitImageApi(MoeAPIView): + # upload image file for other APIs + def post(self): + blob: None | FileStorage = request.files.get("file", None) + if not (blob and blob.name.endswith((".jpg", ".jpeg", ".png", ".gif"))): + raise UploadFileNotFoundError("Please select an image file") + tmpfile = NamedTemporaryFile(dir=MIT_STORAGE_ROOT, delete=False) + tmpfile.write(blob.read()) + tmpfile.close() + return {"filename": tmpfile.name} + + +class MitImageTaskApi(MoeAPIView): + _MIT_IMAGE_TASKS = { + 'mit_detect_text': mit_detect_text, + 'mit_ocr': mit_ocr, + } + + def post(self): + task_params = self.get_json() + task_name = task_params.get("task_name", None) + if task_name not in self._MIT_IMAGE_TASKS: + raise ValueError("Invalid task name") + if 'filename' not in task_params: + raise ValueError("Filename required") + tmpfile_name = task_params['filename'] + tmpfile_path = os.path.join(MIT_STORAGE_ROOT, tmpfile_name) + if os.path.commonprefix([tmpfile_path, MIT_STORAGE_ROOT]) != MIT_STORAGE_ROOT: + raise ValueError("Invalid filename") + if not os.path.isfile(tmpfile_path): + raise ValueError("File not found") + task_id = queue_task(self._MIT_IMAGE_TASKS[task_name], tmpfile_path, **task_params) + return {"id": task_id} + + def get(self, task_id: str): + return _wait_task_result(task_id) + + +class MitTranslateTaskApi(MoeAPIView): + def post(self): + task_params = self.get_json() + text: str = task_params.get('text') + target_lang = task_params.get('target_lang', 'CHT') + task_id = queue_task(mit_translate, text, target_lang) + return {'task_id': task_id} + + def get(self, task_id: str): + return _wait_task_result(task_id) diff --git a/app/apis/urls.py b/app/apis/urls.py index 0202b8a..9523b1a 100644 --- a/app/apis/urls.py +++ b/app/apis/urls.py @@ -68,6 +68,8 @@ ) from app.apis.language import LanguageListAPI from app.apis.target import TargetAPI +from app.apis.manga_image_translator import MitImageApi, MitImageTaskApi, MitTranslateTaskApi +from app import app_config v1_prefix = "/v1" # api主页 @@ -246,7 +248,7 @@ methods=["GET", "PUT", "DELETE", "OPTIONS"], view_func=ProjectSetAPI.as_view("team_project_set"), ) -# 术语库模块 +# 术语库模块 TODO: not used (yet?) in moeflow v1.1.0 term_bank = Blueprint("term_bank", __name__, url_prefix=v1_prefix + "/term-banks") term_bank.add_url_rule( "/", @@ -457,3 +459,31 @@ methods=["GET", "OPTIONS"], view_func=AdminVCodeListAPI.as_view("admin_v_code_list"), ) + +if "MIT_STORAGE_ROOT" in app_config: + mit = Blueprint("manga-image-translator", __name__, url_prefix=v1_prefix + "/mit") + mit.add_url_rule( + '/images', + methods=['POST', 'OPTIONS'], + view_func=MitImageApi.as_view('mit_images') + ) + mit.add_url_rule( + "/image-tasks", + methods=["POST", "OPTIONS"], + view_func=MitImageTaskApi.as_view("mit_image_tasks"), + ) + mit.add_url_rule( + "/image-tasks/", + methods=["GET", "OPTIONS"], + view_func=MitImageTaskApi.as_view("mit_image_tasks"), + ) + mit.add_url_rule( + "/translate-tasks", + methods=["POST", "OPTIONS"], + view_func=MitTranslateTaskApi.as_view("mit_translate_tasks"), + ) + mit.add_url_rule( + "/translate-tasks/", + methods=["GET", "OPTIONS"], + view_func=MitTranslateTaskApi.as_view("mit_translate_tasks"), + ) diff --git a/app/models/target.py b/app/models/target.py index a625fea..5cb6355 100644 --- a/app/models/target.py +++ b/app/models/target.py @@ -59,7 +59,7 @@ def create( file.create_target_cache(target) return target - def outputs(self): + def outputs(self) -> list[Output]: """所有导出""" outputs = Output.objects(project=self.project, target=self).order_by( "-create_time" diff --git a/app/tasks/__init__.py b/app/tasks/__init__.py index 789ba54..22577de 100644 --- a/app/tasks/__init__.py +++ b/app/tasks/__init__.py @@ -7,8 +7,43 @@ flower --port=5555 --broker=redis://localhost:6379/1 """ +import asyncio +import datetime +from typing import Any +from celery import Task +from celery.result import AsyncResult +from celery.exceptions import TimeoutError as CeleryTimeoutError +from app import celery as celery_app +from asgiref.sync import async_to_sync + class SyncResult: """和celery的delay异步返回类似的结果，用于同步、异步切换""" task_id = "sync" + + +def queue_task(task: Task, *args, **kwargs) -> str: + result = task.delay(*args, **kwargs) + result.forget() + return result.id + + +def wait_result_sync(task_id: str, timeout: int = 10) -> Any: + result = AsyncResult(id=task_id, app=celery_app) + try: + return result.get(timeout=timeout) + except CeleryTimeoutError: + raise TimeoutError + + +@async_to_sync +async def wait_result(task_id: str, timeout: int = 10) -> Any: + start = datetime.datetime.now().timestamp() + result = AsyncResult(id=task_id, app=celery_app) + while not result.ready(): + if (datetime.datetime.now().timestamp() - start) > timeout: + result.forget() + raise TimeoutError + await asyncio.sleep(0.5e3) + return result.get() # type: ignore diff --git a/app/tasks/mit.py b/app/tasks/mit.py new file mode 100644 index 0000000..3e1ca2d --- /dev/null +++ b/app/tasks/mit.py @@ -0,0 +1,166 @@ +""" +Text segmentation + OCR using manga-image-translator +""" + +from dataclasses import dataclass +from typing import Any +from app import celery as celery_app +from celery import Task +from celery.result import AsyncResult +import logging + +logger = logging.getLogger(__name__) + + +gpu_options = {"device": "cuda"} + + +@celery_app.task(name="tasks.mit.detect_text") +def _mit_detect_text(path_or_url: str, **kwargs): + pass # Real implementation is in manga_translator/moeflow_worker.py + + +@celery_app.task(name="tasks.mit.ocr") +def _mit_ocr(path_or_url: str, **kwargs): + pass + + +@celery_app.task(name="tasks.mit.translate") +def _mit_translate(**kwargs): + pass + + +@celery_app.task(name="tasks.mit.inpaint") +def _mit_inpaint(path_or_url: str, **kwargs): + pass + + +def _run_mit_detect_text(image_path: str) -> dict: + detect_text: AsyncResult = mit_detect_text.delay( + image_path, + detector_key="default", + # mostly defaults from manga-image-translator/args.py + detect_size=2560, + text_threshold=0.5, + box_threshold=0.7, + unclip_ratio=2.3, + invert=False, + gamma_correct=False, + rotate=False, + verbose=True, + **gpu_options, + ) + # XXX unrecommended but should not cause dead lock + result: dict = detect_text.get(disable_sync_subtasks=False) # type: ignore + logger.info("detect_text finished: %s", result) + return result + + +def _run_mit_ocr(image_path: str, regions: list[dict]) -> list[dict]: + ocr: AsyncResult = mit_ocr.delay( + image_path, + ocr_key="48px", # recommended by rowland + # ocr_key="48px_ctc", + # ocr_key="mocr", # XXX: mocr may have different output format + # use_mocr_merge=True, + regions=regions, + verbose=True, + **gpu_options, + ) + # XXX unrecommended but should not cause dead lock + ocred: list[dict] = ocr.get(disable_sync_subtasks=False) + logger.info("ocr finished: %s", ocred) + for t in ocred: + logger.info("ocr extracted text: %s", t) + return ocred + + +def _run_mit_translate( + text: str, + translator: str = "gpt3.5", + target_lang: str = "CHT", +) -> str: + t: AsyncResult = mit_translate.delay( + query=text, + translator=translator, + target_lang=target_lang, + ) + # XXX unrecommended but should not cause dead lock + + result: str = t.get(disable_sync_subtasks=False) + logger.info("translated %s %s", text, result) + return result[0] + + +@celery_app.task(name="tasks.preprocess_mit") +def _preprocess_mit(image_path: str, target_lang: str): + detected = _run_mit_detect_text(image_path) + ocred = _run_mit_ocr(image_path, detected["textlines"]) + translated_texts = [ + _run_mit_translate(t["text"], target_lang=target_lang) for t in ocred + ] + quads = [ + { + "pts": t["pts"], + "raw_text": t["text"], + "translated": translated_texts[i], + } + for i, t in enumerate(ocred) + ] + return { + "image_path": image_path, + "target_lang": target_lang, + "text_quads": quads, + } + + +@dataclass(frozen=True) +class MitTextQuad: + pts: list[tuple[int, int]] + raw_text: str + translated: str + + def to_dict(self) -> dict[str, Any]: + return { + "pts": self.pts, + "raw_text": self.raw_text, + "translated": self.translated, + } + + @classmethod + def from_dict(cls, d: dict[str, Any]) -> "MitTextQuad": + return cls( + pts=d["pts"], + raw_text=d["raw_text"], + translated=d["translated"], + ) + + +@dataclass(frozen=True) +class MitPreprocessedImage: + image_path: str + target_lang: str + text_quads: list[MitTextQuad] + + def to_dict(self) -> dict[str, Any]: + return { + "image_path": self.image_path, + "target_lang": self.target_lang, + "text_quads": [t.to_dict() for t in self.text_quads], + } + + @classmethod + def from_dict(cls, d: dict[str, Any]) -> "MitPreprocessedImage": + return cls( + image_path=d["image_path"], + target_lang=d["target_lang"], + text_quads=[MitTextQuad.from_dict(t) for t in d["text_quads"]], + ) + + +# export tasks with a better type +mit_detect_text: Task = _mit_detect_text # type: ignore +mit_ocr: Task = _mit_ocr # type: ignore +mit_translate: Task = _mit_translate # type: ignore +mit_inpaint: Task = _mit_inpaint # type: ignore +preprocess_mit: Task = _preprocess_mit # type: ignore diff --git a/manage.py b/manage.py index 0f6b6e3..d9b2c53 100644 --- a/manage.py +++ b/manage.py @@ -91,9 +91,42 @@ def local(action): ) +@click.command("mit_file") +@click.option("--file", help="path to image file") +def mit_preprocess_file(file: str): + from app.tasks.mit import preprocess_mit, MitPreprocessedImage + + proprocessed = preprocess_mit.delay(file, "CHT") + proprocessed_result: dict = proprocessed.get() + + print("proprocessed", proprocessed_result) + print("proprocessed", MitPreprocessedImage.from_dict(proprocessed_result)) + + +@click.command("mit_dir") +@click.option("--dir", help="absolute path to a dir containing image files") +def mit_preprocess_dir(dir: str): + from app.tasks.mit import preprocess_mit, MitPreprocessedImage + + for file in os.listdir(dir): + if not re.match(r".*\.(jpg|png|jpeg)$", file): + continue + full_path = os.path.join(dir, file) + proprocessed = preprocess_mit.delay(full_path, "CHT") + proprocessed_result = MitPreprocessedImage.from_dict(proprocessed.get()) + + print("proprocessed", proprocessed_result) + for q in proprocessed_result.text_quads: + print("text block", q.pts) + print(" ", q.raw_text) + print(" ", q.translated) + + main.add_command(local) main.add_command(docs) main.add_command(migrate) +main.add_command(mit_preprocess_file) +main.add_command(mit_preprocess_dir) if __name__ == "__main__": main() From a74cadec98f34949c50b976bf9a62024649adfae Mon Sep 17 00:00:00 2001 From: Wang Guan Date: Tue, 23 Apr 2024 02:24:02 +0900 Subject: [PATCH 3/5] make ruff happy --- app/apis/manga_image_translator.py | 2 +- manage.py | 3 +-- 2 files changed, 2 insertions(+), 3 deletions(-) diff --git a/app/apis/manga_image_translator.py b/app/apis/manga_image_translator.py index b130798..358d0d4 100644 --- a/app/apis/manga_image_translator.py +++ b/app/apis/manga_image_translator.py @@ -3,7 +3,7 @@ from flask import request from app.exceptions.base import UploadFileNotFoundError -from app.tasks.mit import preprocess_mit, mit_ocr, mit_detect_text, mit_translate +from app.tasks.mit import mit_ocr, mit_detect_text, mit_translate from app.tasks import queue_task, wait_result_sync from app import app_config from werkzeug.datastructures import FileStorage diff --git a/manage.py b/manage.py index d9b2c53..9d48102 100644 --- a/manage.py +++ b/manage.py @@ -1,7 +1,6 @@ import os - +import re import click - import logging from app import create_app, init_db From 1d921973a204798db5e9ab29d970cefd32687786 Mon Sep 17 00:00:00 2001 From: Wang Guan Date: Mon, 29 Apr 2024 18:05:13 +0900 Subject: [PATCH 4/5] make flask happy --- app/apis/urls.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/app/apis/urls.py b/app/apis/urls.py index 9523b1a..569a144 100644 --- a/app/apis/urls.py +++ b/app/apis/urls.py @@ -470,20 +470,20 @@ mit.add_url_rule( "/image-tasks", methods=["POST", "OPTIONS"], - view_func=MitImageTaskApi.as_view("mit_image_tasks"), + view_func=MitImageTaskApi.as_view("mit_image_tasks_create"), ) mit.add_url_rule( "/image-tasks/", methods=["GET", "OPTIONS"], - view_func=MitImageTaskApi.as_view("mit_image_tasks"), + view_func=MitImageTaskApi.as_view("mit_image_tasks_query"), ) mit.add_url_rule( "/translate-tasks", methods=["POST", "OPTIONS"], - view_func=MitTranslateTaskApi.as_view("mit_translate_tasks"), + view_func=MitTranslateTaskApi.as_view("mit_translate_tasks_create"), ) mit.add_url_rule( "/translate-tasks/", methods=["GET", "OPTIONS"], - view_func=MitTranslateTaskApi.as_view("mit_translate_tasks"), + view_func=MitTranslateTaskApi.as_view("mit_translate_tasks_query"), ) From 96d0d3da85343197da85671527865856d13c4653 Mon Sep 17 00:00:00 2001 From: Wang Guan Date: Sun, 19 May 2024 16:54:46 +0900 Subject: [PATCH 5/5] get new api set working --- app/apis/manga_image_translator.py | 57 ++++++++++++++++++------------ app/apis/urls.py | 14 +++++--- app/tasks/mit.py | 44 +++++++++++++++-------- 3 files changed, 72 insertions(+), 43 deletions(-) diff --git a/app/apis/manga_image_translator.py b/app/apis/manga_image_translator.py index 358d0d4..f1fbc18 100644 --- a/app/apis/manga_image_translator.py +++ b/app/apis/manga_image_translator.py @@ -3,12 +3,19 @@ from flask import request from app.exceptions.base import UploadFileNotFoundError -from app.tasks.mit import mit_ocr, mit_detect_text, mit_translate +from app.tasks.mit import ( + mit_ocr, + mit_detect_text, + mit_translate, + mit_detect_text_default_params, + mit_ocr_default_params, +) from app.tasks import queue_task, wait_result_sync from app import app_config from werkzeug.datastructures import FileStorage from tempfile import NamedTemporaryFile import os +from app.utils.logging import logger MIT_STORAGE_ROOT = app_config.get("MIT_STORAGE_ROOT", "/MIT_STORAGE_ROOT_UNDEFINED") @@ -16,21 +23,23 @@ def _wait_task_result(task_id: str): try: result = wait_result_sync(task_id, timeout=1) - return {"id": task_id, "status": "success", "result": result} + return {"task_id": task_id, "status": "success", "result": result} except TimeoutError: return { - "id": task_id, + "task_id": task_id, "status": "pending", } except Exception as e: - return {"id": task_id, "status": "fail", "message": str(e)} + return {"task_id": task_id, "status": "fail", "message": str(e)} class MitImageApi(MoeAPIView): # upload image file for other APIs def post(self): - blob: None | FileStorage = request.files.get("file", None) - if not (blob and blob.name.endswith((".jpg", ".jpeg", ".png", ".gif"))): + logger.info("files: %s", request.files) + blob: None | FileStorage = request.files.get("file") + logger.info("blob: %s", blob) + if not (blob and blob.filename.endswith((".jpg", ".jpeg", ".png", ".gif"))): raise UploadFileNotFoundError("Please select an image file") tmpfile = NamedTemporaryFile(dir=MIT_STORAGE_ROOT, delete=False) tmpfile.write(blob.read()) @@ -39,26 +48,30 @@ def post(self): class MitImageTaskApi(MoeAPIView): - _MIT_IMAGE_TASKS = { - 'mit_detect_text': mit_detect_text, - 'mit_ocr': mit_ocr, - } - def post(self): - task_params = self.get_json() - task_name = task_params.get("task_name", None) - if task_name not in self._MIT_IMAGE_TASKS: - raise ValueError("Invalid task name") - if 'filename' not in task_params: + task_params: dict[str, str] = self.get_json() + logger.info("task_params: %s", task_params) + tmpfile_name = task_params.pop("filename", None) + if not tmpfile_name: raise ValueError("Filename required") - tmpfile_name = task_params['filename'] tmpfile_path = os.path.join(MIT_STORAGE_ROOT, tmpfile_name) if os.path.commonprefix([tmpfile_path, MIT_STORAGE_ROOT]) != MIT_STORAGE_ROOT: raise ValueError("Invalid filename") if not os.path.isfile(tmpfile_path): raise ValueError("File not found") - task_id = queue_task(self._MIT_IMAGE_TASKS[task_name], tmpfile_path, **task_params) - return {"id": task_id} + task_name = task_params.pop("task_name", None) + if task_name == "mit_detect_text": + merged_params = mit_detect_text_default_params.copy() + merged_params.update(task_params) + task_id = queue_task(mit_detect_text, tmpfile_path, **merged_params) + return {"task_id": task_id} + elif task_name == "mit_ocr": + merged_params = mit_ocr_default_params.copy() + merged_params.update(task_params) + task_id = queue_task(mit_ocr, tmpfile_path, **merged_params) + return {"task_id": task_id} + else: + raise ValueError("Invalid task name") def get(self, task_id: str): return _wait_task_result(task_id) @@ -67,10 +80,8 @@ def get(self, task_id: str): class MitTranslateTaskApi(MoeAPIView): def post(self): task_params = self.get_json() - text: str = task_params.get('text') - target_lang = task_params.get('target_lang', 'CHT') - task_id = queue_task(mit_translate, text, target_lang) - return {'task_id': task_id} + task_id = queue_task(mit_translate, **task_params) + return {"task_id": task_id} def get(self, task_id: str): return _wait_task_result(task_id) diff --git a/app/apis/urls.py b/app/apis/urls.py index 569a144..249200b 100644 --- a/app/apis/urls.py +++ b/app/apis/urls.py @@ -68,7 +68,11 @@ ) from app.apis.language import LanguageListAPI from app.apis.target import TargetAPI -from app.apis.manga_image_translator import MitImageApi, MitImageTaskApi, MitTranslateTaskApi +from app.apis.manga_image_translator import ( + MitImageApi, + MitImageTaskApi, + MitTranslateTaskApi, +) from app import app_config v1_prefix = "/v1" @@ -460,12 +464,12 @@ view_func=AdminVCodeListAPI.as_view("admin_v_code_list"), ) -if "MIT_STORAGE_ROOT" in app_config: +if app_config["MIT_STORAGE_ROOT"]: mit = Blueprint("manga-image-translator", __name__, url_prefix=v1_prefix + "/mit") mit.add_url_rule( - '/images', - methods=['POST', 'OPTIONS'], - view_func=MitImageApi.as_view('mit_images') + "/images", + methods=["POST", "OPTIONS"], + view_func=MitImageApi.as_view("mit_image_upload"), ) mit.add_url_rule( "/image-tasks", diff --git a/app/tasks/mit.py b/app/tasks/mit.py index 3e1ca2d..1d3801b 100644 --- a/app/tasks/mit.py +++ b/app/tasks/mit.py @@ -35,19 +35,28 @@ def _mit_inpaint(path_or_url: str, **kwargs): pass +def to_dict(**kwargs) -> dict[str, Any]: + return kwargs + + +mit_detect_text_default_params = to_dict( + detector_key="default", + # mostly defaults from manga-image-translator/args.py + detect_size=2560, + text_threshold=0.5, + box_threshold=0.7, + unclip_ratio=2.3, + invert=False, + gamma_correct=False, + rotate=False, + verbose=True, +) + + def _run_mit_detect_text(image_path: str) -> dict: detect_text: AsyncResult = mit_detect_text.delay( image_path, - detector_key="default", - # mostly defaults from manga-image-translator/args.py - detect_size=2560, - text_threshold=0.5, - box_threshold=0.7, - unclip_ratio=2.3, - invert=False, - gamma_correct=False, - rotate=False, - verbose=True, + **mit_detect_text_default_params, **gpu_options, ) # XXX unrecommended but should not cause dead lock @@ -56,15 +65,20 @@ def _run_mit_detect_text(image_path: str) -> dict: return result +mit_ocr_default_params = to_dict( + ocr_key="48px", # recommended by rowland + # ocr_key="48px_ctc", + # ocr_key="mocr", # XXX: mocr may have different output format + # use_mocr_merge=True, + verbose=True, +) + + def _run_mit_ocr(image_path: str, regions: list[dict]) -> list[dict]: ocr: AsyncResult = mit_ocr.delay( image_path, - ocr_key="48px", # recommended by rowland - # ocr_key="48px_ctc", - # ocr_key="mocr", # XXX: mocr may have different output format - # use_mocr_merge=True, regions=regions, - verbose=True, + **mit_ocr_default_params, **gpu_options, ) # XXX unrecommended but should not cause dead lock