import gradio as gr
from paddleocr import PaddleOCR,draw_ocr
import json
import os

ocr_en = PaddleOCR(use_angle_cls=True, lang='en')
ocr_ch = PaddleOCR(use_angle_cls=True, lang='ch')
ocr_ru = PaddleOCR(use_angle_cls=True, lang='ru')
ocr_ru = PaddleOCR(use_angle_cls=True, lang='ar')


def perform_ocr(img, lang):
    ocr = ""
    if lang == "English":
        ocr = ocr_en
    elif lang == "Chinese (Simplified)":
        ocr = ocr_ch
    elif lang == "Russian":
        ocr = ocr_ru
    elif lang == "Arabic":
        ocr = ocr_ar

    result =  ocr.ocr(img.name, cls=True)
    final_result = ""
    for idx in range(len(result)):
        res = result[idx]
        for line in res:
            final_result += str(line)
            
    return [img.name, final_result]


demo = gr.Blocks()

with demo:
    gr.HTML("<h1>PaddleOCR v4")
    with gr.Row():
        with gr.Column():
            input_image = gr.Image()
            input_checkbox = gr.CheckboxGroup(["English", "Chinese (Simplified)", "Russian", "Arabic"], label="Languages"),
            input_button = gr.Button()
        with gr.Column():
            output_image = gr.Image()
            output_text = gr.Text()

    btn.click(fn=perform_ocr, inputs=[input_image, input_checkbox], outputs=[output_image, output_text])
    
demo.launch()