import pandas as pd
import PIL
from PIL import Image
from PIL import ImageDraw
import gradio as gr
import torch
import easyocr

# Download examples
urls = {
    "english.png": "https://github.com/JaidedAI/EasyOCR/raw/master/examples/english.png",
    "thai.jpg": "https://github.com/JaidedAI/EasyOCR/raw/master/examples/thai.jpg",
    "french.jpg": "https://github.com/JaidedAI/EasyOCR/raw/master/examples/french.jpg",
    "chinese.jpg": "https://github.com/JaidedAI/EasyOCR/raw/master/examples/chinese.jpg",
    "japanese.jpg": "https://github.com/JaidedAI/EasyOCR/raw/master/examples/japanese.jpg",
    "korean.png": "https://github.com/JaidedAI/EasyOCR/raw/master/examples/korean.png",
    "Hindi.jpeg": "https://i.imgur.com/mwQFd7G.jpeg"
}

for filename, url in urls.items():
    torch.hub.download_url_to_file(url, filename)

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

def inference(img_path, lang):
    reader = easyocr.Reader(lang)
    bounds = reader.readtext(img_path)
    
    im = PIL.Image.open(img_path)
    draw_boxes(im, bounds)
    im.save("result.jpg")

    df = pd.DataFrame(bounds)[[1, 2]]
    df.columns = ["text", "confidence"]
    
    return "result.jpg", df

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"
]

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"]]
]

with gr.Blocks(css=".output_image, .input_image {height: 40rem; width: 100%;}") as demo:

    gr.Markdown("# **EasyOCR**")
    gr.Markdown("Upload an image and select languages to extract text.")

    with gr.Row():
        image_input = gr.Image(type="filepath", label="Input Image")
        lang_input = gr.CheckboxGroup(choices=choices, value=["en"], label="Language")

    run_btn = gr.Button("Run OCR")

    with gr.Row():
        output_img = gr.Image(label="Output")
        output_df = gr.Dataframe(headers=["text", "confidence"], label="Detected Text")

    gr.Examples(examples, [image_input, lang_input], [output_img, output_df], fn=inference)

    run_btn.click(
        inference,
        inputs=[image_input, lang_input],
        outputs=[output_img, output_df]
    )

demo.launch()