File size: 7,537 Bytes
c6d26b8
9a6d49e
c6d26b8
 
 
 
 
 
 
 
 
 
 
56cf37c
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
c6d26b8
 
 
 
 
 
9a6d49e
c6d26b8
 
9a6d49e
c6d26b8
 
 
9a6d49e
c6d26b8
 
 
9a6d49e
c6d26b8
9a6d49e
 
c6d26b8
 
 
 
9a6d49e
c6d26b8
 
 
 
 
 
 
 
 
 
 
 
 
 
 
8c811eb
e4dd4df
c6d26b8
 
 
 
9a6d49e
c6d26b8
 
 
 
9a6d49e
c6d26b8
 
 
 
 
 
 
 
 
9a6d49e
c6d26b8
 
 
 
9a6d49e
c6d26b8
9a6d49e
c6d26b8
 
 
9a6d49e
c6d26b8
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
9a6d49e
c6d26b8
 
 
e59da9c
c6d26b8
 
 
 
 
 
 
 
 
 
 
 
 
e59da9c
c6d26b8
 
 
 
 
 
 
 
 
 
 
e59da9c
c6d26b8
 
 
 
 
 
9a6d49e
c6d26b8
 
 
 
 
 
e59da9c
c6d26b8
 
e59da9c
 
c6d26b8
 
 
 
 
 
9a6d49e
c6d26b8
 
 
9a6d49e
c6d26b8
 
 
 
 
 
 
 
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
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
# OCR Translate v0.2


import os

os.system("sudo 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
# Added below code
from fastapi import FastAPI, File, Request, UploadFile, Body, Depends, HTTPException
from fastapi.security.api_key import APIKeyHeader
from typing import Optional, Annotated
from fastapi.encoders import jsonable_encoder
from PIL import Image
from io import BytesIO

API_KEY = os.environ.get("API_KEY")

app = FastAPI()
api_key_header = APIKeyHeader(name="api_key", auto_error=False)

def get_api_key(api_key: Optional[str] = Depends(api_key_header)):
    if api_key is None or api_key != API_KEY:
        raise HTTPException(status_code=401, detail="Unauthorized access")
    return api_key

@app.post("/api/ocr", response_model=dict)
async def ocr(
    api_key: str = Depends(get_api_key),
    image: UploadFile = File(...),
    # languages: list = Body(["eng"])
):
    try:
        content = await image.read()
        image = Image.open(BytesIO(content))
        print("[image]",image)
        if hasattr(pytesseract, "image_to_string"):
            print("Image to string function is available")
            # print(pytesseract.image_to_string(image, lang = 'eng'))
            text = ocr_tesseract(image, ['eng'])
        else:
            print("Image to string function is not available")
        # text = pytesseract.image_to_string(image, lang="+".join(languages))
    except Exception as e:
        return {"error": str(e)}, 500

    return {"ImageText": "text"}

nltk.download('punkt')

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="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}"  # Model name

    tokenizer = MarianTokenizer.from_pretrained(model_name)  # tokenizer
    model = MarianMTModel.from_pretrained(model_name)  # Model

    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):
    print("[img]", img)
    print("[languages]", 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!"

    # Select 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.Box():

            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", source="upload", type="pil", label="image")
                    with gr.Row():
                        inputs_lang = gr.CheckboxGroup(choices=["chi_sim", "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=["zh-en", "en-zh"],
                                                     type="value",
                                                     value="zh-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", ["chi_sim"]]]
                gr.Examples(example_list, [inputs_img, inputs_lang], outputs_text, ocr_tesseract, cache_examples=False)

        # -------------- translate --------------
        with gr.Box():

            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()