import os
from matplotlib.pyplot import switch_backend, text
os.system('pip install paddlepaddle')
os.system('pip install paddleocr')
from paddleocr import PaddleOCR, draw_ocr
from PIL import Image
import gradio as gr
import torch

from transformers import M2M100ForConditionalGeneration, M2M100Tokenizer

model = M2M100ForConditionalGeneration.from_pretrained("facebook/m2m100_1.2B")
tokenizer = M2M100Tokenizer.from_pretrained("facebook/m2m100_1.2B")

title = 'OCR Translator'
description = 'This is a gradio demo for OCR and translating using the PaddleOCR and m2m100_418M model. It takes in an input of an image, the language to be read using OCR, and the language the result will be translated to. The PaddleOCR implementation is limited to English, Chinese, Japanese, German, and French, while the results can be translated to 100 languages.'
article = '<p>This is only a demo. The official repository can be found <a href="https://github.com/PaddlePaddle/PaddleOCR">here</a></p>'
examples = [['japan_to_en.png', 'Japanese', 'English'], ['en_to_fr.png', 'English', 'French'], ['german_to_en.jpg', 'German', 'English']]

def inference(img, src_lang, tgt_lang):
    if src_lang == 'Chinese':
        img_src_lang = 'ch'
        tokenizer.src_lang = 'zh'
    elif src_lang == 'English':
        img_src_lang = 'en'
        tokenizer.src_lang = 'en'
    elif src_lang == 'French':
        img_src_lang = 'fr'
        tokenizer.src_lang = 'fr'
    elif src_lang == 'Japanese':
        img_src_lang = 'japan'
        tokenizer.src_lang = "ja"
    elif src_lang == 'German':
        img_src_lang = 'german'
        tokenizer.src_lang = "de"

    if tgt_lang == 'Chinese':
        tgt_lang = 'zh'
    elif tgt_lang == 'English':
        tgt_lang = 'en'
    elif tgt_lang == 'French':
        tgt_lang = 'fr'
    elif src_lang == 'Japanese':
        tgt_lang = 'ja'
    elif src_lang == 'German':
        tgt_lang = 'de'

    # Use OCR Model
    ocr = PaddleOCR(use_angle_cls = True, lang = img_src_lang, use_gpu = False)
    img_path = img.name
    result = ocr.ocr(img_path, cls = True)
    image = Image.open(img_path).convert('RGB')
    boxes = [line[0] for line in result]
    txts = [line[1][0] for line in result]
    im_show = draw_ocr(image, boxes, font_path = 'chinese.simfang.ttf')
    im_show = Image.fromarray(im_show)
    im_show.save('result.jpg')

    # Parse OCR Text
    input_text = ''
    for txt in txts:
        input_text = input_text + " " + txt

    # Translate to Target Language    
    encoded_src = tokenizer(input_text, return_tensors="pt")
    generated_tokens = model.generate(**encoded_src, forced_bos_token_id=tokenizer.get_lang_id(tgt_lang), use_cache=True)
    result = tokenizer.batch_decode(generated_tokens, skip_special_tokens=True)[0]
    return ['result.jpg', input_text, result]

gr.Interface(
    inference,
    [gr.inputs.Image(type='file', label='Input'),
    gr.inputs.Dropdown(choices=['Chinese', 'English', 'French', 'German', 'Japanese'], type="value", default='en', label='Source Language'), 
    gr.inputs.Dropdown(choices=['Chinese', 'English', 'French', 'German', 'Japanese'], type="value", default='en', label='Translate to')],
    [gr.outputs.Image(type='file', label='Output'), gr.outputs.Textbox(label = 'Output Text'), gr.outputs.Textbox(label = 'Translated Text')],
    title=title,
    examples=examples,
    description=description,
    article=article,
    enable_queue=True
    ).launch(debug=True)