edit app.py

app.py

@@ -1,132 +1,81 @@
-import pandas as pd
-import PIL
+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
-from PIL import ImageDraw
 import gradio as gr
 import torch
-import easyocr
 
-torch.hub.download_url_to_file('https://github.com/JaidedAI/EasyOCR/raw/master/examples/english.png', 'english.png')
-torch.hub.download_url_to_file('https://github.com/JaidedAI/EasyOCR/raw/master/examples/thai.jpg', 'thai.jpg')
-torch.hub.download_url_to_file('https://github.com/JaidedAI/EasyOCR/raw/master/examples/french.jpg', 'french.jpg')
-torch.hub.download_url_to_file('https://github.com/JaidedAI/EasyOCR/raw/master/examples/chinese.jpg', 'chinese.jpg')
-torch.hub.download_url_to_file('https://github.com/JaidedAI/EasyOCR/raw/master/examples/japanese.jpg', 'japanese.jpg')
-torch.hub.download_url_to_file('https://github.com/JaidedAI/EasyOCR/raw/master/examples/korean.png', 'korean.png')
-torch.hub.download_url_to_file('https://i.imgur.com/mwQFd7G.jpeg', 'Hindi.jpeg')
+from transformers import M2M100ForConditionalGeneration, M2M100Tokenizer
 
-def draw_boxes(image, bounds, color='yellow', width=2):
-    draw = ImageDraw.Draw(image)
-    for bound in bounds:
-        p0, p1, p2, p3 = bound[0]
-        draw.line([*p0, *p1, *p2, *p3, *p0], fill=color, width=width)
-    return image
+model = M2M100ForConditionalGeneration.from_pretrained("facebook/m2m100_1.2B")
+tokenizer = M2M100Tokenizer.from_pretrained("facebook/m2m100_1.2B")
 
-def inference(img, lang):
-    reader = easyocr.Reader(lang)
-    bounds = reader.readtext(img.name)
-    im = PIL.Image.open(img.name)
-    draw_boxes(im, bounds)
-    im.save('result.jpg')
-    return ['result.jpg', pd.DataFrame(bounds).iloc[: , 1:]]
+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]
 
-title = 'EasyOCR'
-description = 'Gradio demo for EasyOCR. EasyOCR demo supports 80+ languages.To use it, simply upload your image and choose a language from the dropdown menu, or click one of the examples to load them. Read more at the links below.'
-article = "<p style='text-align: center'><a href='https://www.jaided.ai/easyocr/'>Ready-to-use OCR with 80+ supported languages and all popular writing scripts including Latin, Chinese, Arabic, Devanagari, Cyrillic and etc.</a> | <a href='https://github.com/JaidedAI/EasyOCR'>Github Repo</a></p>"
-examples = [['english.png',['en']],['thai.jpg',['th']],['french.jpg',['fr', 'en']],['chinese.jpg',['ch_sim', 'en']],['japanese.jpg',['ja', 'en']],['korean.png',['ko', 'en']],['Hindi.jpeg',['hi', 'en']]]
-css = ".output_image, .input_image {height: 40rem !important; width: 100% !important;}"
-choices = [
-    "abq",
-    "ady",
-    "af",
-    "ang",
-    "ar",
-    "as",
-    "ava",
-    "az",
-    "be",
-    "bg",
-    "bh",
-    "bho",
-    "bn",
-    "bs",
-    "ch_sim",
-    "ch_tra",
-    "che",
-    "cs",
-    "cy",
-    "da",
-    "dar",
-    "de",
-    "en",
-    "es",
-    "et",
-    "fa",
-    "fr",
-    "ga",
-    "gom",
-    "hi",
-    "hr",
-    "hu",
-    "id",
-    "inh",
-    "is",
-    "it",
-    "ja",
-    "kbd",
-    "kn",
-    "ko",
-    "ku",
-    "la",
-    "lbe",
-    "lez",
-    "lt",
-    "lv",
-    "mah",
-    "mai",
-    "mi",
-    "mn",
-    "mr",
-    "ms",
-    "mt",
-    "ne",
-    "new",
-    "nl",
-    "no",
-    "oc",
-    "pi",
-    "pl",
-    "pt",
-    "ro",
-    "ru",
-    "rs_cyrillic",
-    "rs_latin",
-    "sck",
-    "sk",
-    "sl",
-    "sq",
-    "sv",
-    "sw",
-    "ta",
-    "tab",
-    "te",
-    "th",
-    "tjk",
-    "tl",
-    "tr",
-    "ug",
-    "uk",
-    "ur",
-    "uz",
-    "vi"
-]
 gr.Interface(
     inference,
-    [gr.inputs.Image(type='file', label='Input'),gr.inputs.CheckboxGroup(choices, type="value", default=['en'], label='language')],
-    [gr.outputs.Image(type='file', label='Output'), gr.outputs.Dataframe(headers=['text', 'confidence'])],
+    [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,
-    examples=examples,
-    css=css,
     enable_queue=True
     ).launch(debug=True)

requirements.txt

@@ -1,4 +1,6 @@
 Pillow
-gradio
+Gradio
 torch
-easyocr
+transformers
+numpy 
+sentencepiece