import gradio as gr
from paddleocr import PaddleOCR, draw_ocr
import json
import os
import cv2
import numpy as np
from PIL import Image

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='cyrillic')
ocr_ar = PaddleOCR(use_angle_cls=True, lang='arabic')
#ocr_ch = PaddleOCR(det_model_dir="models/det/ch/ch_PP-OCRv4_det_infer", rec_model_dir="models/rec/ch/ch_PP-OCRv4_rec_infer", cls_model_dir="models/cls/ch_ppocr_mobile_v2.0_cls_infer", rec_char_dict_path="models/dict/ppocr_keys_v1.txt", lang="ch")


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

    result = ocr.ocr(img, cls=True)
    final_result = ""
 
    image = Image.open(img).convert('RGB')
    # boxes = [line[0] for line in result]
    # txts = [line[1][0] for line in result]
    # scores = [line[1][1] for line in result]
    # im_show = draw_ocr(image, boxes, txts, scores, font_path='fonts/simfang.ttf')
    # im_show = Image.fromarray(im_show)
    return [image, result]

demo = gr.Blocks()

with demo:
    gr.Markdown("# Multilingual OCR")
    with gr.Row():
        with gr.Column():
            input_image = gr.Image(source="upload", type="filepath")
            input_radio = gr.Radio(["English", "Chinese (Simplified)", "Russian", "Ukrainian", "Arabic", "Persian"], label="Languages"),
            input_button = gr.Button("Run!")
        with gr.Column():
            output_image = gr.Image()
            output_text = gr.Textbox(label="Results")

    input_button.click(fn=perform_ocr, inputs=[input_image], outputs=[output_image, output_text])

demo.launch(server_name="0.0.0.0", server_port=7860)