v02 update

README.md

@@ -4,7 +4,7 @@ emoji: 🚀
 colorFrom: gray
 colorTo: indigo
 sdk: gradio
-sdk_version: 3.0.18
+sdk_version: 3.0.26
 app_file: app.py
 pinned: false
 license: gpl-3.0

app.py

@@ -1,20 +1,21 @@
-# OCR Translate v0.1
+# OCR Translate v0.2
 # 创建人：曾逸夫
-# 创建时间：2022-06-14
-# email: zyfiy1314@163.com
-# 项目地址：https://gitee.com/CV_Lab/ocr-translate
+# 创建时间：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.1
+OCR_TR_DESCRIPTION = '''# OCR Translate v0.2
 <div id="content_align">基于Tesseract的OCR翻译系统</div>'''
 
 # 图片路径
@@ -26,6 +27,7 @@ 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}"  # 模型名称
 
@@ -55,23 +57,35 @@ def ocr_tesseract(img, languages):
     return ocr_str
 
 
-# 示例
-def set_example_image(example: list) -> dict:
-    return gr.Image.update(value=example[0])
-
-
 # 清除
 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("sudo apt-get install xclip")
+        print(e)
+
+
+# 清除剪贴板
+def cp_clear():
+    pyclip.clear()
+
+
 # 翻译
-def translate(input_text):
+def translate(input_text, inputs_transStyle):
     # 参考：https://huggingface.co/docs/transformers/model_doc/marian
     if input_text is None or input_text == "":
         return "系统提示：没有可翻译的内容！"
 
-    tokenizer, model = model_choice()
+    # 选择翻译模型
+    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")
@@ -90,6 +104,7 @@ def translate(input_text):
 
     return translate_text[2:]
 
+
 def main():
 
     with gr.Blocks(css='style.css') as ocr_tr:
@@ -106,16 +121,31 @@ def main():
                     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='语言')
+                        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():
-                    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.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():
@@ -124,29 +154,25 @@ def main():
                 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)
+                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])
 
-        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])
+        # ---------------------- 翻译 ----------------------
+        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)
 
 

packages.txt

@@ -1 +1,2 @@
-tesseract-ocr-all
+tesseract-ocr-all
+xclip