Initial commit
This commit is contained in:
@@ -0,0 +1,88 @@
|
||||
from PIL import Image, ImageDraw, ImageFont
|
||||
from dotenv import load_dotenv
|
||||
import os
|
||||
import io
|
||||
import numpy as np
|
||||
import ast
|
||||
from helpers.utils import romanize, intercepts
|
||||
load_dotenv()
|
||||
|
||||
MAX_TRANSLATE = int(os.getenv('MAX_TRANSLATION', 200))
|
||||
FONT_FILE = os.getenv('FONT_FILE')
|
||||
FONT_SIZE = int(os.getenv('FONT_SIZE', 16))
|
||||
LINE_SPACING = int(os.getenv('LINE_SPACING', 3))
|
||||
TEXT_COLOR = os.getenv('TEXT_COLOR', "#ff0000")
|
||||
LINE_HEIGHT = FONT_SIZE
|
||||
TO_ROMANIZE = ast.literal_eval(os.getenv('TO_ROMANIZE', 'True'))
|
||||
font = ImageFont.truetype(FONT_FILE, FONT_SIZE)
|
||||
|
||||
|
||||
|
||||
def modify_image_bytes(image_bytes: io.BytesIO, ocr_output, translation: list) -> bytes:
|
||||
# Load the image from bytes
|
||||
with io.BytesIO(image_bytes) as byte_stream:
|
||||
image = Image.open(byte_stream)
|
||||
draw = ImageDraw.Draw(image)
|
||||
translate_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 translate_image(draw: ImageDraw, translation: list, ocr_output: list, max_translate: int) -> ImageDraw:
|
||||
translation
|
||||
translated_number = 0
|
||||
bounding_boxes = []
|
||||
for i, (position, untranslated_phrase, confidence) in enumerate(ocr_output):
|
||||
if translated_number >= max_translate:
|
||||
break
|
||||
translate_one_phrase(draw, translation[i], position, bounding_boxes, untranslated_phrase)
|
||||
translated_number += 1
|
||||
return draw
|
||||
|
||||
def translate_one_phrase(draw: ImageDraw, translated_phrase: str, position: tuple, bounding_boxes: list, untranslated_phrase: str) -> ImageDraw:
|
||||
# Draw the bounding box
|
||||
top_left, _, _, _ = position
|
||||
position = (top_left[0], top_left[1] - 60)
|
||||
text_content = f"{translated_phrase}\n{romanize(untranslated_phrase, TO_ROMANIZE)}\n{untranslated_phrase}"
|
||||
lines = text_content.split('\n')
|
||||
x,y = position
|
||||
max_width = 0
|
||||
total_height = 0
|
||||
total_height = len(lines) * (LINE_HEIGHT + LINE_SPACING)
|
||||
for line in lines:
|
||||
bbox = draw.textbbox(position, line, font=font)
|
||||
line_width = bbox[2] - bbox[0]
|
||||
max_width = max(max_width, line_width)
|
||||
bounding_box = (x, y, x + max_width, y + total_height, untranslated_phrase)
|
||||
print(f"Bounding Box of Interest: {bounding_box}")
|
||||
|
||||
adjust_if_intersects(x, y, 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= TEXT_COLOR, font=font)
|
||||
adjusted_y += FONT_SIZE + LINE_SPACING
|
||||
position = (adjusted_x,adjusted_y)
|
||||
print(f"Adjusted_bounding_box: {bounding_box[-1]}.\n")
|
||||
|
||||
def adjust_if_intersects(x: int, y: int, bounding_box: tuple, bounding_boxes: list, untranslated_phrase: str, max_width: int, total_height: int) -> tuple:
|
||||
y = np.max([y,0])
|
||||
if len(bounding_boxes) > 0:
|
||||
for box in bounding_boxes:
|
||||
print(f'Investigating box: {box}')
|
||||
if intercepts((box[0],box[2]),(bounding_box[0],bounding_box[2])) and intercepts((box[1],box[3]),(y, y+total_height)):
|
||||
print(f'Overlapping change adjustment to {untranslated_phrase}')
|
||||
y = np.max([y,box[3]]) + LINE_SPACING
|
||||
print(y, box[3])
|
||||
print(f'Changed to {(x,y, x+max_width, y+total_height, untranslated_phrase)}')
|
||||
adjusted_bounding_box = (x, y, x + max_width, y + total_height, untranslated_phrase)
|
||||
bounding_boxes.append(adjusted_bounding_box)
|
||||
return adjusted_bounding_box
|
||||
|
||||
|
||||
|
||||
@@ -0,0 +1,87 @@
|
||||
from paddleocr import PaddleOCR
|
||||
import easyocr
|
||||
from rapidocr_onnxruntime import RapidOCR
|
||||
import langid
|
||||
from helpers.utils import contains_lang
|
||||
from concurrent.futures import ThreadPoolExecutor
|
||||
# PaddleOCR
|
||||
# Paddleocr supports Chinese, English, French, German, Korean and Japanese.
|
||||
# You can set the parameter `lang` as `ch`, `en`, `fr`, `german`, `korean`, `japan`
|
||||
# to switch the language model in order.
|
||||
# need to run only once to download and load model into memory
|
||||
|
||||
def _paddle_init(lang='ch', use_angle_cls=False, use_GPU=True):
|
||||
return PaddleOCR(use_angle_cls=use_angle_cls, lang=lang, use_GPU=use_GPU)
|
||||
|
||||
|
||||
def _paddle_ocr(ocr, image) -> list:
|
||||
### return a list containing the bounding box, text and confidence of the detected text
|
||||
result = ocr.ocr(image, cls=False)[0]
|
||||
if not isinstance(result, list):
|
||||
return []
|
||||
result = [ (pos, text[0], text[1]) for pos, text in result]
|
||||
return result
|
||||
|
||||
# EasyOCR has support for many languages
|
||||
|
||||
def _easy_init(ocr_languages: list, use_GPU=True):
|
||||
return easyocr.Reader(ocr_languages, gpu=use_GPU)
|
||||
|
||||
def _easy_ocr(ocr,image) -> list:
|
||||
return ocr.readtext(image)
|
||||
|
||||
# RapidOCR mostly for mandarin and some other asian languages
|
||||
|
||||
def _rapid_init(use_GPU=True):
|
||||
return RapidOCR(use_gpu=use_GPU)
|
||||
|
||||
def _rapid_ocr(ocr, image) -> list:
|
||||
return ocr(image)
|
||||
|
||||
### Initialize the OCR model
|
||||
def init_OCR(model='paddle', **kwargs):
|
||||
if model == 'paddle':
|
||||
return _paddle_init(**kwargs)
|
||||
elif model == 'easy':
|
||||
return _easy_init(**kwargs)
|
||||
elif model == 'rapid':
|
||||
return _rapid_init(**kwargs)
|
||||
|
||||
### Perform OCR on the image
|
||||
def identify(ocr, image) -> list:
|
||||
if isinstance(ocr, PaddleOCR):
|
||||
return _paddle_ocr(ocr, image)
|
||||
elif isinstance(ocr, easyocr.Reader):
|
||||
return _easy_ocr(ocr, image)
|
||||
elif isinstance(ocr, RapidOCR):
|
||||
return _rapid_ocr(ocr, image)
|
||||
else:
|
||||
raise ValueError("Invalid OCR model. Please initialise the OCR model first with init() and pass it as an argument to identify().")
|
||||
|
||||
|
||||
### Filter out the results that are not in the source language
|
||||
def id_filtered(ocr, image, lang) -> list:
|
||||
result = identify(ocr, image)
|
||||
|
||||
### Parallelise since langid is slow
|
||||
def classify_text(entry):
|
||||
return entry if langid.classify(entry[1])[0] == lang else None
|
||||
with ThreadPoolExecutor() as executor:
|
||||
results_no_eng = list(filter(None, executor.map(classify_text, result)))
|
||||
return results_no_eng
|
||||
|
||||
|
||||
# zh, ja, ko
|
||||
def id_lang(ocr, image, lang) -> list:
|
||||
result = identify(ocr, image)
|
||||
filtered = [entry for entry in result if contains_lang(entry[1], lang)]
|
||||
return filtered
|
||||
|
||||
def get_words(ocr_output) -> list:
|
||||
return [entry[1] for entry in ocr_output]
|
||||
|
||||
def get_positions(ocr_output) -> list:
|
||||
return [entry[0] for entry in ocr_output]
|
||||
|
||||
def get_confidences(ocr_output) -> list:
|
||||
return [entry[2] for entry in ocr_output]
|
||||
@@ -0,0 +1,76 @@
|
||||
from transformers import M2M100Tokenizer, M2M100ForConditionalGeneration, AutoTokenizer, AutoModelForSeq2SeqLM, GPTQConfig
|
||||
import torch, os
|
||||
from dotenv import load_dotenv
|
||||
load_dotenv()
|
||||
|
||||
if os.getenv('TRANSLATION_USE_GPU') in ['False', '0', 'false', 'no', 'No', 'NO', 'FALSE']:
|
||||
device = torch.device("cpu")
|
||||
else:
|
||||
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
|
||||
|
||||
### Batch translate a list of strings
|
||||
|
||||
|
||||
# M2M100 model
|
||||
|
||||
|
||||
def init_M2M():
|
||||
global tokenizer, model
|
||||
tokenizer = M2M100Tokenizer.from_pretrained("facebook/m2m100_418M", local_files_only=True)
|
||||
model = M2M100ForConditionalGeneration.from_pretrained("facebook/m2m100_418M", local_files_only=True, torch_dtype=torch.float16).to(device)
|
||||
model.eval()
|
||||
|
||||
|
||||
|
||||
def translate_M2M(text, from_lang = 'zh', target_lang = 'en'):
|
||||
if len(text) == 0:
|
||||
return []
|
||||
tokenizer.src_lang = from_lang
|
||||
with torch.no_grad():
|
||||
encoded = tokenizer(text, return_tensors="pt", padding=True).to(device)
|
||||
generated_tokens = model.generate(**encoded,
|
||||
forced_bos_token_id=tokenizer.get_lang_id(target_lang))
|
||||
translated = tokenizer.batch_decode(generated_tokens, skip_special_tokens=True)
|
||||
return translated
|
||||
|
||||
|
||||
# Helsinki-NLP model Opus MT
|
||||
|
||||
|
||||
def init_OPUS():
|
||||
global tokenizer, model
|
||||
tokenizer = AutoTokenizer.from_pretrained("Helsinki-NLP/opus-mt-tc-bible-big-zhx-en", local_files_only=True)
|
||||
model = AutoModelForSeq2SeqLM.from_pretrained("Helsinki-NLP/opus-mt-tc-bible-big-zhx-en", local_files_only=True, torch_dtype=torch.float16).to(device)
|
||||
model.eval()
|
||||
|
||||
def translate_OPUS(text: list[str]) -> list[str]:
|
||||
if len(text) == 0:
|
||||
return []
|
||||
with torch.no_grad():
|
||||
encoded = tokenizer(text, return_tensors="pt", padding=True).to(device)
|
||||
generated_tokens = model.generate(**encoded)
|
||||
translated = tokenizer.batch_decode(generated_tokens, skip_special_tokens=True)
|
||||
return translated
|
||||
|
||||
###
|
||||
def init_TRANSLATE(model): # model = 'opus' or 'm2m'
|
||||
if model == 'opus':
|
||||
init_OPUS()
|
||||
elif model == 'm2m':
|
||||
init_M2M()
|
||||
else:
|
||||
raise ValueError("Invalid model. Please use 'opus' or 'm2m'.")
|
||||
|
||||
|
||||
###
|
||||
def translate(text, model, **kwargs):
|
||||
if model == 'opus':
|
||||
return translate_OPUS(text)
|
||||
elif model == 'm2m':
|
||||
try:
|
||||
return translate_M2M(text, **kwargs)
|
||||
except:
|
||||
raise ValueError("Please provide the from_lang and target_lang variables if you are using the M2M model.")
|
||||
|
||||
else:
|
||||
raise ValueError("Invalid model. Please use 'opus' or 'm2m'.")
|
||||
@@ -0,0 +1,155 @@
|
||||
import re, uroman as ur
|
||||
from pypinyin import pinyin
|
||||
import pyscreenshot as ImageGrab # wayland tings not sure if it will work on other machines alternatively use mss
|
||||
import mss, io, os
|
||||
from PIL import Image
|
||||
import jaconv
|
||||
import MeCab
|
||||
import unidic
|
||||
mecab = MeCab.Tagger('-d "{}"'.format(unidic.DICDIR))
|
||||
uroman = ur.Uroman()
|
||||
|
||||
|
||||
# define a function for checking whether one axis of a shape intercepts with another
|
||||
def intercepts(x,y):
|
||||
# both x and y are two dimensional tuples representing the ends of a line on one dimension.
|
||||
x1, x2 = x
|
||||
y1, y2 = y
|
||||
return (x1 <= y1 <= x2) or (x1 <= y2 <= x2) or (y1 <= x1 <= y2) or (y1 <= x2 <= y2)
|
||||
|
||||
def is_wayland():
|
||||
return 'WAYLAND_DISPLAY' in os.environ
|
||||
|
||||
# path to save screenshot of monitor to
|
||||
def printsc_wayland(region, save: bool = False, path: str = None):
|
||||
if save:
|
||||
im = ImageGrab.grab(bbox=region)
|
||||
im.save(path)
|
||||
else:
|
||||
return ImageGrab.grab(bbox=region)
|
||||
|
||||
|
||||
def printsc_non_wayland(region, save: bool = False, path: str = None):
|
||||
# use mss to capture the screen
|
||||
with mss.mss() as sct:
|
||||
# grab the screen
|
||||
img = sct.grab(region)
|
||||
# convert the image to a PIL image
|
||||
image = Image.frombytes("RGB", img.size, img.bgra, "raw", "BGRX")
|
||||
# save the image if save is True
|
||||
if save:
|
||||
image.save(path)
|
||||
|
||||
|
||||
def printsc(region, save: bool = False, path: str = None):
|
||||
try:
|
||||
if is_wayland():
|
||||
return printsc_wayland(region, save, path)
|
||||
else:
|
||||
return printsc_non_wayland(region, save, path)
|
||||
except Exception as e:
|
||||
print(f'Error {e}')
|
||||
|
||||
def convert_image_to_bytes(img):
|
||||
with io.BytesIO() as byte_stream:
|
||||
img.save(byte_stream, format='PNG') # Save the image to the byte stream
|
||||
return byte_stream.getvalue() # Get the byte representation
|
||||
|
||||
|
||||
|
||||
def bytes_to_image(image_bytes):
|
||||
# Load the image from bytes
|
||||
byte_stream = io.BytesIO(image_bytes)
|
||||
# Open the image from the BytesIO stream
|
||||
image = Image.open(byte_stream)
|
||||
return image
|
||||
|
||||
# for japanese
|
||||
|
||||
def add_furigana(text):
|
||||
parsed = mecab.parse(text).split('\n')[:-2]
|
||||
furigana_string = ''
|
||||
for i in parsed:
|
||||
words = i.split('\t')[0]
|
||||
add = f'({jaconv.kata2hira(i.split(',')[6])})'
|
||||
to_add = add if contains_kanji(words) else ''
|
||||
furigana_string += i.split('\t')[0] + to_add
|
||||
return furigana_string
|
||||
|
||||
|
||||
def contains_kanji(text):
|
||||
return bool(re.search(r'[\u4E00-\u9FFF]', text))
|
||||
|
||||
def contains_hiragana(text):
|
||||
return bool(re.search(r'[\u3040-\u309F]', text))
|
||||
|
||||
def contains_katakana(text):
|
||||
return bool(re.search(r'[\u30A0-\u30FF]', text))
|
||||
|
||||
|
||||
|
||||
def romanize(text, piny=False):
|
||||
if piny:
|
||||
return ' '.join([ py[0] for py in pinyin(text, heteronym=True)])
|
||||
return uroman.romanize_string(text)
|
||||
|
||||
# check if a string contains words from a language
|
||||
def contains_lang(text, lang):
|
||||
# Matches any character in the Unicode range of the language
|
||||
if lang == 'zh':
|
||||
return bool(re.search(r'[\u4e00-\u9fff]', text))
|
||||
elif lang == 'ja':
|
||||
return bool(re.search(r'[\u3040-\u30ff]', text))
|
||||
elif lang == 'ko':
|
||||
return bool(re.search(r'[\uac00-\ud7af]', text))
|
||||
elif lang == 'en':
|
||||
return bool(re.search(r'[a-zA-Z]', text))
|
||||
else:
|
||||
raise ValueError("Invalid language. Please use one of 'en', 'zh', 'ja', or 'ko'.")
|
||||
|
||||
### en, ch_sim, ja, ko rapidocr only has chinese and en at the moment
|
||||
def standardize_lang(lang):
|
||||
if lang == 'ch_sim':
|
||||
easyocr_lang = 'ch_sim'
|
||||
paddleocr_lang = 'ch'
|
||||
rapidocr_lang = 'ch'
|
||||
translation_model_lang = 'zh'
|
||||
elif lang == 'ch_tra':
|
||||
easyocr_lang = 'ch_tra'
|
||||
paddleocr_lang = 'ch'
|
||||
rapidocr_lang = 'ch'
|
||||
translation_model_lang = 'zh'
|
||||
elif lang == 'ja':
|
||||
easyocr_lang = 'ja'
|
||||
paddleocr_lang = 'ja'
|
||||
rapidocr_lang = 'ja'
|
||||
translation_model_lang = 'ja'
|
||||
elif lang == 'ko':
|
||||
easyocr_lang = 'korean'
|
||||
paddleocr_lang = 'ko'
|
||||
rapidocr_lang = 'ko'
|
||||
translation_model_lang = 'ko'
|
||||
elif lang == 'en':
|
||||
easyocr_lang = 'en'
|
||||
paddleocr_lang = 'en'
|
||||
rapidocr_lang = 'en'
|
||||
translation_model_lang = 'en'
|
||||
else:
|
||||
raise ValueError("Invalid language. Please use one of 'en', 'ch_sim', 'ch_tra', 'ja', or 'ko'.")
|
||||
return {'easyocr_lang': easyocr_lang, 'paddleocr_lang': paddleocr_lang, 'rapidocr_lang': rapidocr_lang, 'translation_model_lang': translation_model_lang}
|
||||
|
||||
def which_ocr_lang(model):
|
||||
if model == 'easy':
|
||||
return 'easyocr_lang'
|
||||
elif model == 'paddle':
|
||||
return 'paddleocr_lang'
|
||||
elif model == 'rapid':
|
||||
return 'rapidocr_lang'
|
||||
else:
|
||||
raise ValueError("Invalid OCR model. Please use one of 'easy', 'paddle', or 'rapid'.")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
# Example usage
|
||||
japanesetext = "本が好きにちは"
|
||||
print(add_furigana(japanesetext))
|
||||
Reference in New Issue
Block a user