139 lines
6.3 KiB
Python
139 lines
6.3 KiB
Python
from PIL import Image, ImageDraw, ImageFont
|
|
import os, io, sys, numpy as np
|
|
|
|
sys.path.insert(0, os.path.join(os.path.dirname(__file__), 'helpers'))
|
|
from utils import romanize, intercepts, add_furigana
|
|
from logging_config import logger
|
|
from config import ADD_OVERLAY, SOURCE_LANG, MAX_TRANSLATE, FONT_FILE, FONT_SIZE, LINE_SPACING, FONT_COLOUR, LINE_HEIGHT, TO_ROMANIZE, FILL_COLOUR, REGION
|
|
|
|
from PySide6.QtGui import QFont
|
|
font = ImageFont.truetype(FONT_FILE, FONT_SIZE)
|
|
|
|
|
|
def modify_image_bytes(image_bytes: io.BytesIO, ocr_output, translation: list) -> bytes:
|
|
"""Modify the image bytes with the translated text and return the modified image bytes"""
|
|
|
|
with io.BytesIO(image_bytes) as byte_stream:
|
|
image = Image.open(byte_stream)
|
|
draw = ImageDraw.Draw(image)
|
|
draw_on_image(draw, translation, ocr_output, MAX_TRANSLATE)
|
|
|
|
# Save the modified image back to bytes without changing the format
|
|
with io.BytesIO() as byte_stream:
|
|
image.save(byte_stream, format=image.format) # Save in original format
|
|
modified_image_bytes = byte_stream.getvalue()
|
|
return modified_image_bytes
|
|
|
|
def draw_on_image(draw: ImageDraw, translation: list, ocr_output: list, max_translate: int, replace = False) -> ImageDraw:
|
|
"""Draw the original, translated and optionally the romanisation of the texts on the image"""
|
|
translated_number = 0
|
|
bounding_boxes = []
|
|
logger.debug(f"Translations: {len(translation)} {translation}")
|
|
logger.debug(f"OCR output: {len(ocr_output)} {ocr_output}")
|
|
for i, (position, untranslated_phrase, confidence) in enumerate(ocr_output):
|
|
logger.debug(f"Untranslated phrase: {untranslated_phrase}")
|
|
if translated_number >= max_translate - 1:
|
|
break
|
|
if replace:
|
|
draw = draw_one_phrase_replace(draw, translation[i], position, bounding_boxes, untranslated_phrase)
|
|
else:
|
|
draw_one_phrase_add(draw, translation[i], position, bounding_boxes, untranslated_phrase)
|
|
translated_number += 1
|
|
return draw
|
|
|
|
def draw_one_phrase_add(draw: ImageDraw,
|
|
translated_phrase: str,
|
|
position: tuple, bounding_boxes: list,
|
|
untranslated_phrase: str) -> ImageDraw:
|
|
"""Draw the bounding box rectangle and text on the image above the original text"""
|
|
|
|
if SOURCE_LANG == 'ja':
|
|
untranslated_phrase = add_furigana(untranslated_phrase)
|
|
romanized_phrase = romanize(untranslated_phrase, 'ja')
|
|
else:
|
|
romanized_phrase = romanize(untranslated_phrase, SOURCE_LANG)
|
|
if TO_ROMANIZE:
|
|
text_content = f"{translated_phrase}\n{romanized_phrase}\n{untranslated_phrase}"
|
|
else:
|
|
text_content = f"{translated_phrase}\n{untranslated_phrase}"
|
|
|
|
lines = text_content.split('\n')
|
|
|
|
# Draw the bounding box
|
|
top_left, _, _, _ = position
|
|
max_width = get_max_width(lines, FONT_FILE, FONT_SIZE)
|
|
total_height = get_max_height(lines, FONT_SIZE, LINE_SPACING)
|
|
right_edge = REGION[2]
|
|
|
|
# Ensure the text is within the screen. P.S. Text on the edge may still be squished together if there are too many to translate
|
|
x_onscreen = top_left[0] if top_left[0] + max_width <= right_edge else right_edge - max_width
|
|
y_onscreen = max(top_left[1] - total_height, 0)
|
|
bounding_box = (x_onscreen, y_onscreen, x_onscreen + max_width, y_onscreen + total_height, untranslated_phrase)
|
|
|
|
adjust_if_intersects(x_onscreen, y_onscreen, bounding_box, bounding_boxes, untranslated_phrase, max_width, total_height)
|
|
|
|
adjusted_x, adjusted_y, adjusted_max_x, adjusted_max_y, _ = bounding_boxes[-1]
|
|
draw.rectangle([(adjusted_x,adjusted_y), (adjusted_max_x, adjusted_max_y)], outline="black", width=1)
|
|
position = (adjusted_x,adjusted_y)
|
|
for line in lines:
|
|
draw.text(position, line, fill= FONT_COLOUR, font=font)
|
|
if ADD_OVERLAY:
|
|
overlay.add_next_text_at_position_no_update(position[0], position[1], line, text_color=FONT_COLOUR)
|
|
adjusted_y += FONT_SIZE + LINE_SPACING
|
|
position = (adjusted_x,adjusted_y)
|
|
|
|
|
|
|
|
|
|
### Only support for horizontal text atm, vertical text is on the todo list
|
|
def draw_one_phrase_replace(draw: ImageDraw,
|
|
translated_phrase: str,
|
|
position: tuple, bounding_boxes: list,
|
|
untranslated_phrase: str) -> ImageDraw:
|
|
"""Cover up old text and add translation directly on top"""
|
|
# Draw the bounding box
|
|
top_left, _, _, bottom_right = position
|
|
max_width = bottom_right[0] - top_left[0]
|
|
font_size = bottom_right[1] - top_left[1]
|
|
draw.rectangle([top_left, bottom_right], fill=FILL_COLOUR)
|
|
while True:
|
|
font = ImageFont.truetype(FONT_FILE, font_size)
|
|
if font.get_max_width < max_width:
|
|
draw.text(top_left, translated_phrase, fill= FONT_COLOUR, font=font)
|
|
break
|
|
elif font_size <= 1:
|
|
break
|
|
else:
|
|
font_size -= 1
|
|
|
|
def get_max_width(lines: list, font_path, font_size) -> int:
|
|
"""Get the maximum width of the text lines"""
|
|
font = ImageFont.truetype(font_path, font_size)
|
|
max_width = 0
|
|
dummy_image = Image.new("RGB", (1, 1))
|
|
draw = ImageDraw.Draw(dummy_image)
|
|
for line in lines:
|
|
bbox = draw.textbbox((0,0), line, font=font)
|
|
line_width = bbox[2] - bbox[0]
|
|
max_width = max(max_width, line_width)
|
|
return max_width
|
|
|
|
def get_max_height(lines: list, font_size, line_spacing) -> int:
|
|
"""Get the maximum height of the text lines"""
|
|
return len(lines) * (font_size + line_spacing)
|
|
|
|
def adjust_if_intersects(x: int, y: int,
|
|
bounding_box: tuple, bounding_boxes: list,
|
|
untranslated_phrase: str,
|
|
max_width: int, total_height: int) -> tuple:
|
|
"""Adjust the y coordinate if the bounding box intersects with any other bounding box"""
|
|
y = np.max([y,0])
|
|
if len(bounding_boxes) > 0:
|
|
for box in bounding_boxes:
|
|
if intercepts((box[0],box[2]),(bounding_box[0],bounding_box[2])) and intercepts((box[1],box[3]),(y, y+total_height)):
|
|
y = np.max([y,box[3]]) + LINE_SPACING
|
|
adjusted_bounding_box = (x, y, x + max_width, y + total_height, untranslated_phrase)
|
|
bounding_boxes.append(adjusted_bounding_box)
|
|
return adjusted_bounding_box
|
|
|