Spaces:
Sleeping
Sleeping
| # OCR Translate v0.2 | |
| # 创建人：曾逸夫 (Founder: Zeng Yifu) | |
| # Creation time：2022-07-19 | |
| import os | |
| #os.system("apt-get install xclip") | |
| import gradio as gr | |
| import nltk | |
| import pyclip | |
| import pytesseract | |
| import nltk | |
| from nltk.tokenize import sent_tokenize | |
| from transformers import MarianMTModel, MarianTokenizer | |
| from transformers import T5Tokenizer, T5ForConditionalGeneration | |
| nltk.download('punkt') | |
| nltk.download('punkt_tab') | |
| OCR_TR_DESCRIPTION = '''# OCR Translate v0.2 | |
| <div id="content_align">OCR translation system based on Tesseract</div>''' | |
| # Image Path | |
| img_dir = "./data" | |
| # Get tesseract language list | |
| choices = os.popen('tesseract --list-langs').read().split('\n')[1:-1] | |
| # Translation model selection | |
| def model_choice(src, trg): | |
| # https://huggingface.co/Helsinki-NLP/opus-mt-tc-big-pt-en | |
| # https://huggingface.co/Helsinki-NLP/opus-mt-tc-big-en-pt | |
| # https://huggingface.co/unicamp-dl/translation-pt-en-t5 | |
| # https://huggingface.co/unicamp-dl/translation-en-pt-t5 | |
| model_name = f"unicamp-dl/translation-{src}-{trg}-t5" # Model Name | |
| #tokenizer = MarianTokenizer.from_pretrained(model_name) # Tokenizer | |
| #model = MarianMTModel.from_pretrained(model_name) # Model | |
| tokenizer = T5Tokenizer.from_pretrained(model_name) | |
| model = T5ForConditionalGeneration.from_pretrained(model_name) | |
| return tokenizer, model | |
| # Convert tesseract language list to pytesseract language | |
| 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 | |
| # Clear | |
| def clear_content(): | |
| return None | |
| # Copy to Clipboard | |
| def cp_text(input_text): | |
| # sudo apt-get install xclip | |
| try: | |
| pyclip.copy(input_text) | |
| except Exception as e: | |
| print("sudo apt-get install xclip") | |
| print(e) | |
| # Clear Clipboard | |
| def cp_clear(): | |
| pyclip.clear() | |
| # Translate | |
| def translate(input_text, inputs_transStyle): | |
| # Reference：https://huggingface.co/docs/transformers/model_doc/marian | |
| if input_text is None or input_text == "": | |
| return "System prompt: There is no content to translate!" | |
| # Selecting a translation model | |
| trans_src, trans_trg = inputs_transStyle.split("-")[0], inputs_transStyle.split("-")[1] | |
| tokenizer, model = model_choice(trans_src, trans_trg) | |
| translate_text = "" | |
| input_text_list = input_text.split("\n\n") | |
| translate_text_list_tmp = [] | |
| for i in range(len(input_text_list)): | |
| if input_text_list[i] != "": | |
| translate_text_list_tmp.append(input_text_list[i]) | |
| for i in range(len(translate_text_list_tmp)): | |
| translated_sub = model.generate( | |
| **tokenizer(sent_tokenize(translate_text_list_tmp[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 Text Extraction -------------- | |
| with gr.Blocks(): | |
| with gr.Row(): | |
| gr.Markdown("### Step 01: Text Extraction") | |
| with gr.Row(): | |
| with gr.Column(): | |
| with gr.Row(): | |
| inputs_img = gr.Image(image_mode="RGB", sources="upload", type="pil", label="image") | |
| with gr.Row(): | |
| inputs_lang = gr.CheckboxGroup(choices=["por", "eng"], | |
| type="value", | |
| value=['eng'], | |
| label='language') | |
| with gr.Row(): | |
| clear_img_btn = gr.Button('Clear') | |
| ocr_btn = gr.Button(value='OCR Extraction', variant="primary") | |
| with gr.Column(): | |
| with gr.Row(): | |
| outputs_text = gr.Textbox(label="Extract content", lines=20) | |
| with gr.Row(): | |
| inputs_transStyle = gr.Radio(choices=["pt-en", "en-pt"], | |
| type="value", | |
| value="pt-en", | |
| label='translation mode') | |
| with gr.Row(): | |
| clear_text_btn = gr.Button('Clear') | |
| translate_btn = gr.Button(value='Translate', variant="primary") | |
| with gr.Row(): | |
| example_list = [["./data/test.png", ["eng"]], ["./data/test02.png", ["eng"]], | |
| ["./data/test03.png", ["por"]]] | |
| gr.Examples(example_list, [inputs_img, inputs_lang], outputs_text, ocr_tesseract, cache_examples=False) | |
| # -------------- Translate -------------- | |
| with gr.Blocks(): | |
| with gr.Row(): | |
| gr.Markdown("### Step 02: Translation") | |
| with gr.Row(): | |
| outputs_tr_text = gr.Textbox(label="Translate Content", lines=20) | |
| with gr.Row(): | |
| cp_clear_btn = gr.Button(value='Clear Clipboard') | |
| cp_btn = gr.Button(value='Copy to clipboard', variant="primary") | |
| # ---------------------- 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]) | |
| # ---------------------- Translate ---------------------- | |
| translate_btn.click(fn=translate, inputs=[outputs_text, inputs_transStyle], outputs=[outputs_tr_text]) | |
| clear_text_btn.click(fn=clear_content, inputs=[], outputs=[outputs_text]) | |
| # ---------------------- Copy to Clipboard ---------------------- | |
| cp_btn.click(fn=cp_text, inputs=[outputs_tr_text], outputs=[]) | |
| cp_clear_btn.click(fn=cp_clear, inputs=[], outputs=[]) | |
| ocr_tr.launch(inbrowser=True) | |
| if __name__ == '__main__': | |
| main() |