# OCR Translate v0.1
# 创建人：曾逸夫
# 创建时间：2022-06-14
# email: zyfiy1314@163.com
# 项目地址：https://gitee.com/CV_Lab/ocr-translate

import os

import gradio as gr
import nltk
import pytesseract
from nltk.tokenize import sent_tokenize
from transformers import MarianMTModel, MarianTokenizer

nltk.download('punkt')

OCR_TR_DESCRIPTION = '''# OCR Translate v0.1
<div id="content_align">基于Tesseract的OCR翻译系统</div>'''

# 图片路径
img_dir = "./data"

# 获取tesseract语言列表
choices = os.popen('tesseract --list-langs').read().split('\n')[1:-1]


# 翻译模型选择
def model_choice(src="en", trg="zh"):
    # https://huggingface.co/Helsinki-NLP/opus-mt-en-zh
    model_name = f"Helsinki-NLP/opus-mt-{src}-{trg}"  # 模型名称

    tokenizer = MarianTokenizer.from_pretrained(model_name)  # 分词器
    model = MarianMTModel.from_pretrained(model_name)  # 模型

    return tokenizer, model


# tesseract语言列表转pytesseract语言
def ocr_lang(lang_list):
    lang_str = ""
    lang_len = len(lang_list)
    if lang_len == 1:
        return lang_list[0]
    else:
        for i in range(lang_len):
            lang_list.insert(lang_len - i, "+")

        lang_str = "".join(lang_list[:-1])
        return lang_str


# ocr tesseract
def ocr_tesseract(img, languages):
    ocr_str = pytesseract.image_to_string(img, lang=ocr_lang(languages))
    return ocr_str


# 示例
def set_example_image(example: list) -> dict:
    return gr.Image.update(value=example[0])


# 清除
def clear_content():
    return None


# 翻译
def translate(input_text):
    # 参考：https://huggingface.co/docs/transformers/model_doc/marian
    if input_text is None or input_text == "":
        return "系统提示：没有可翻译的内容！"

    tokenizer, model = model_choice()

    translate_text = ""
    input_text_list = input_text.split("\n\n")

    for i in range(len(input_text_list)):
        translated_sub = model.generate(
            **tokenizer(sent_tokenize(input_text_list[i]), return_tensors="pt", truncation=True, padding=True))
        tgt_text_sub = [tokenizer.decode(t, skip_special_tokens=True) for t in translated_sub]
        translate_text_sub = "".join(tgt_text_sub)
        translate_text = translate_text + "\n\n" + translate_text_sub

    return translate_text[2:]


def main():

    with gr.Blocks(css='style.css') as ocr_tr:
        gr.Markdown(OCR_TR_DESCRIPTION)

        # -------------- OCR 文字提取 --------------
        with gr.Box():

            with gr.Row():
                gr.Markdown("### Step 01: 文字提取")

            with gr.Row():
                with gr.Column():
                    with gr.Row():
                        inputs_img = gr.Image(image_mode="RGB", source="upload", type="pil", label="图片")
                    with gr.Row():
                        inputs_lang = gr.CheckboxGroup(choices=choices, type="value", value=['eng'], label='语言')

                    with gr.Row():
                        clear_img_btn = gr.Button('Clear')
                        ocr_btn = gr.Button(value='OCR 提取', variant="primary")

                with gr.Column():
                    imgs_path = sorted(os.listdir(img_dir))
                    example_images = gr.Dataset(components=[inputs_img],
                                                samples=[[f"{img_dir}/{i}"] for i in imgs_path])

        # -------------- 翻译 --------------
        with gr.Box():

            with gr.Row():
                gr.Markdown("### Step 02: 翻译")

            with gr.Row():
                with gr.Column():
                    with gr.Row():
                        outputs_text = gr.Textbox(label="提取内容", lines=20)
                    with gr.Row():
                        clear_text_btn = gr.Button('Clear')
                        translate_btn = gr.Button(value='翻译', variant="primary")
                with gr.Column():
                    outputs_tr_text = gr.Textbox(label="翻译内容", lines=20)

        # ---------------------- OCR Tesseract ----------------------
        ocr_btn.click(fn=ocr_tesseract, inputs=[inputs_img, inputs_lang], outputs=[
            outputs_text,])

        clear_img_btn.click(fn=clear_content, inputs=[], outputs=[inputs_img])

        example_images.click(fn=set_example_image, inputs=[
            example_images,], outputs=[
                inputs_img,])

        # ---------------------- OCR Tesseract ----------------------
        translate_btn.click(fn=translate, inputs=[outputs_text], outputs=[outputs_tr_text])
        clear_text_btn.click(fn=clear_content, inputs=[], outputs=[outputs_text])

    ocr_tr.launch(inbrowser=True)


if __name__ == '__main__':
    main()