Spaces:
Sleeping
Sleeping
File size: 6,050 Bytes
940a520 5a0f532 940a520 5a0f532 31d35a6 940a520 5a0f532 940a520 5a0f532 940a520 5a0f532 e16eaa7 940a520 e16eaa7 5a0f532 940a520 31d35a6 940a520 5a0f532 940a520 5a0f532 940a520 e16eaa7 755ce29 e16eaa7 755ce29 e16eaa7 755ce29 e16eaa7 5a0f532 e16eaa7 5a0f532 940a520 5a0f532 940a520 5a0f532 940a520 5a0f532 940a520 5a0f532 940a520 5a0f532 940a520 5a0f532 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 |
# OCR Translate v0.2
# 创建人:曾逸夫
# 创建时间:2022-07-19
import os
#os.system("apt-get install xclip")
import gradio as gr
import nltk
import pyclip
import pytesseract
from nltk.tokenize import sent_tokenize
from transformers import MarianMTModel, MarianTokenizer
nltk.download('punkt')
OCR_TR_DESCRIPTION = '''# OCR Translate v0.2
<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-zh-en
# 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 clear_content():
return None
# 复制到剪贴板
def cp_text(input_text):
# sudo apt-get install xclip
try:
pyclip.copy(input_text)
except Exception as e:
print("install package xclip")
print(e)
# 清除剪贴板
def cp_clear():
pyclip.clear()
# 翻译
def translate(input_text, inputs_transStyle):
# 参考:https://huggingface.co/docs/transformers/model_doc/marian
if input_text is None or input_text == "":
return "系统提示:没有可翻译的内容!"
# 选择翻译模型
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 文字提取 --------------
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=["chi_sim", "eng"],
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():
with gr.Row():
outputs_text = gr.Textbox(label="提取内容", lines=20)
with gr.Row():
inputs_transStyle = gr.Radio(choices=["zh-en", "en-zh"],
type="value",
value="zh-en",
label='翻译模式')
with gr.Row():
clear_text_btn = gr.Button('Clear')
translate_btn = gr.Button(value='翻译', variant="primary")
with gr.Row():
example_list = [["./data/test.png", ["eng"]], ["./data/test02.png", ["eng"]],
["./data/test03.png", ["chi_sim"]]]
gr.Examples(example_list, [inputs_img, inputs_lang], outputs_text, ocr_tesseract, cache_examples=False)
# -------------- 翻译 --------------
with gr.Box():
with gr.Row():
gr.Markdown("### Step 02: 翻译")
with gr.Row():
outputs_tr_text = gr.Textbox(label="翻译内容", lines=20)
with gr.Row():
cp_clear_btn = gr.Button(value='清除剪贴板')
cp_btn = gr.Button(value='复制到剪贴板', 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_btn.click(fn=translate, inputs=[outputs_text, inputs_transStyle], outputs=[outputs_tr_text])
clear_text_btn.click(fn=clear_content, inputs=[], outputs=[outputs_text])
# ---------------------- 复制到剪贴板 ----------------------
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()
|