Skip to content

Commit

Permalink
wip
Browse files Browse the repository at this point in the history
  • Loading branch information
jokester committed Mar 4, 2024
1 parent d1987be commit 1dbd75e
Showing 3 changed files with 75 additions and 19 deletions.
37 changes: 25 additions & 12 deletions manga_translator/manga_translator.py
Original file line number Diff line number Diff line change
@@ -180,6 +180,7 @@ async def translate_path(self, path: str, dest: str = None, params: dict[str, Un
else:
logger.info(f'Done. Translated {translated_count} image{"" if translated_count == 1 else "s"}')

# translate a single file, with error handling
async def translate_file(self, path: str, dest: str, params: dict):
if not params.get('overwrite') and os.path.exists(dest):
logger.info(
@@ -217,7 +218,7 @@ async def translate_file(self, path: str, dest: str, params: dict):
attempts += 1
return False

# translate single file
# the hard lifting
async def _translate_file(self, path: str, dest: str, ctx: Context) -> bool:
if path.endswith('.txt'):
with open(path, 'r') as f:
@@ -1264,26 +1265,38 @@ def format_translate(self, ctx: Context, return_image: bool):
return web.json_response({'details': results, 'img': img})

class PostSchema(Schema):
target_language = fields.Str(required=False, validate=lambda a: a.upper() in VALID_LANGUAGES)
# common?
# verbose: boolean

# detector stage
detector = fields.Str(required=False, validate=lambda a: a.lower() in DETECTORS)
detection_size = fields.Integer(required=False)
text_threshold = fields.Float(required=False)
box_threshold = fields.Float(required=False)
unclip_ratio = fields.Float(required=False)
det_gamma_correct = fields.Bool(required=False)
det_invert = fields.Bool(required=False)
det_rotate = fields.Bool(required=False)
det_auto_rotate = fields.Bool(required=False)

# translation
target_language = fields.Str(required=False, validate=lambda a: a.upper() in VALID_LANGUAGES)

# OCR
ocr = fields.Str(required=False, validate=lambda a: a.lower() in OCRS)

translator = fields.Str(required=False, validate=lambda a: a.lower() in TRANSLATORS)
# inpaint
inpainter = fields.Str(required=False, validate=lambda a: a.lower() in INPAINTERS)

# upscale (optional)
upscaler = fields.Str(required=False, validate=lambda a: a.lower() in UPSCALERS)
translator = fields.Str(required=False, validate=lambda a: a.lower() in TRANSLATORS)
direction = fields.Str(required=False, validate=lambda a: a.lower() in {'auto', 'h', 'v'})
upscale_ratio = fields.Integer(required=False)
direction = fields.Str(required=False, validate=lambda a: a.lower() in {'auto', 'h', 'v'})
translator_chain = fields.Str(required=False)
selective_translation = fields.Str(required=False)
attempts = fields.Integer(required=False)
detection_size = fields.Integer(required=False)
text_threshold = fields.Float(required=False)
box_threshold = fields.Float(required=False)
unclip_ratio = fields.Float(required=False)
inpainting_size = fields.Integer(required=False)
det_rotate = fields.Bool(required=False)
det_auto_rotate = fields.Bool(required=False)
det_invert = fields.Bool(required=False)
det_gamma_correct = fields.Bool(required=False)
min_text_length = fields.Integer(required=False)
colorization_size = fields.Integer(required=False)
denoise_sigma = fields.Integer(required=False)
50 changes: 46 additions & 4 deletions manga_translator/moeflow_worker.py
Original file line number Diff line number Diff line change
@@ -1,9 +1,11 @@
# consume tasks from moeflow job worker with manga-image-translator code
from celery import Celery
from manga_translator import MangaTranslator
from asgiref.sync import async_to_sync

env = {
"RABBITMQ_USER": 'moeflow',
"RABBITMQ_PASS": 'PLEASE_CHANGE_THIS',
"RABBITMQ_USER": 'moeflow',
"RABBITMQ_PASS": 'PLEASE_CHANGE_THIS',
'RABBITMQ_VHOST_NAME': 'moeflow',
'MONGODB_USER': 'moeflow',
'MONGODB_PASS': 'PLEASE_CHANGE_THIS',
@@ -13,8 +15,8 @@

celery_app = Celery(
"manga-image-translator-moeflow-worker",
broker =f"amqp://{env['RABBITMQ_USER']}:{env['RABBITMQ_PASS']}@moeflow-rabbitmq:5672/{env['RABBITMQ_VHOST_NAME']}",
backend= f"mongodb://{env['MONGODB_USER']}:{env['MONGODB_PASS']}@moeflow-mongodb:27017/{env['MONGODB_DB_NAME']}?authSource=admin",
broker=f"amqp://{env['RABBITMQ_USER']}:{env['RABBITMQ_PASS']}@moeflow-rabbitmq:5672/{env['RABBITMQ_VHOST_NAME']}",
backend=f"mongodb://{env['MONGODB_USER']}:{env['MONGODB_PASS']}@moeflow-mongodb:27017/{env['MONGODB_DB_NAME']}?authSource=admin",
mongodb_backend_settings={
"database": env["MONGODB_DB_NAME"],
"taskmeta_collection": "celery_taskmeta",
@@ -24,3 +26,43 @@
@celery_app.task
def add(x, y):
return x + y


@celery_app.task
def translate(image_path: str, dest: str, args_dict: dict):
detector_args = {
'detector': 'default', # see detecton/__init__.py
# 'img_rgb': from image
# 'detect_size': 48,
}

ocr_dict = {
'ocr': '48px', # reportedly to work best
# textlines: from detector
# TODO more
}

translate_dict = {

}

inpaint_dict = {

}

ctx_dict = {
# **args_dict,
# upscale_ratio
**detector_args,
**ocr_dict,
**translate_dict,
**ocr_dict
}
return do_translate(image_path, dest, ctx_dict)


@async_to_sync
async def do_translate(image_path: str, dest_image_path: str, args_dict: dict):
translator = MangaTranslator()

await translator.translate_file(image_path, dest_image_path, args_dict)
7 changes: 4 additions & 3 deletions manga_translator/ocr/common.py
Original file line number Diff line number Diff line change
@@ -7,6 +7,7 @@

from ..utils import InfererModule, TextBlock, ModelWrapper, Quadrilateral


class CommonOCR(InfererModule):
def _generate_text_direction(self, bboxes: List[Union[Quadrilateral, TextBlock]]):
if len(bboxes) > 0:
@@ -19,7 +20,7 @@ def _generate_text_direction(self, bboxes: List[Union[Quadrilateral, TextBlock]]

G = nx.Graph()
for i, box in enumerate(bboxes):
G.add_node(i, box = box)
G.add_node(i, box=box)
for ((u, ubox), (v, vbox)) in itertools.combinations(enumerate(bboxes), 2):
if quadrilateral_can_merge_region(ubox, vbox, aspect_ratio_tol=1):
G.add_edge(u, v)
@@ -30,9 +31,9 @@ def _generate_text_direction(self, bboxes: List[Union[Quadrilateral, TextBlock]]
majority_dir = Counter(dirs).most_common(1)[0][0]
# sort
if majority_dir == 'h':
nodes = sorted(nodes, key = lambda x: bboxes[x].aabb.y + bboxes[x].aabb.h // 2)
nodes = sorted(nodes, key=lambda x: bboxes[x].aabb.y + bboxes[x].aabb.h // 2)
elif majority_dir == 'v':
nodes = sorted(nodes, key = lambda x: -(bboxes[x].aabb.x + bboxes[x].aabb.w))
nodes = sorted(nodes, key=lambda x: -(bboxes[x].aabb.x + bboxes[x].aabb.w))
# yield overall bbox and sorted indices
for node in nodes:
yield bboxes[node], majority_dir

0 comments on commit 1dbd75e

Please sign in to comment.