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

lang_id = {
    "Afrikaans": "af",
    "Amharic": "am",
    "Arabic": "ar",
    "Asturian": "ast",
    "Azerbaijani": "az",
    "Bashkir": "ba",
    "Belarusian": "be",
    "Bulgarian": "bg",
    "Bengali": "bn",
    "Breton": "br",
    "Bosnian": "bs",
    "Catalan": "ca",
    "Cebuano": "ceb",
    "Czech": "cs",
    "Welsh": "cy",
    "Danish": "da",
    "German": "de",
    "Greeek": "el",
    "English": "en",
    "Spanish": "es",
    "Estonian": "et",
    "Persian": "fa",
    "Fulah": "ff",
    "Finnish": "fi",
    "French": "fr",
    "Western Frisian": "fy",
    "Irish": "ga",
    "Gaelic": "gd",
    "Galician": "gl",
    "Gujarati": "gu",
    "Hausa": "ha",
    "Hebrew": "he",
    "Hindi": "hi",
    "Croatian": "hr",
    "Haitian": "ht",
    "Hungarian": "hu",
    "Armenian": "hy",
    "Indonesian": "id",
    "Igbo": "ig",
    "Iloko": "ilo",
    "Icelandic": "is",
    "Italian": "it",
    "Japanese": "ja",
    "Javanese": "jv",
    "Georgian": "ka",
    "Kazakh": "kk",
    "Central Khmer": "km",
    "Kannada": "kn",
    "Korean": "ko",
    "Luxembourgish": "lb",
    "Ganda": "lg",
    "Lingala": "ln",
    "Lao": "lo",
    "Lithuanian": "lt",
    "Latvian": "lv",
    "Malagasy": "mg",
    "Macedonian": "mk",
    "Malayalam": "ml",
    "Mongolian": "mn",
    "Marathi": "mr",
    "Malay": "ms",
    "Burmese": "my",
    "Nepali": "ne",
    "Dutch": "nl",
    "Norwegian": "no",
    "Northern Sotho": "ns",
    "Occitan": "oc",
    "Oriya": "or",
    "Panjabi": "pa",
    "Polish": "pl",
    "Pushto": "ps",
    "Portuguese": "pt",
    "Romanian": "ro",
    "Russian": "ru",
    "Sindhi": "sd",
    "Sinhala": "si",
    "Slovak": "sk",
    "Slovenian": "sl",
    "Somali": "so",
    "Albanian": "sq",
    "Serbian": "sr",
    "Swati": "ss",
    "Sundanese": "su",
    "Swedish": "sv",
    "Swahili": "sw",
    "Tamil": "ta",
    "Thai": "th",
    "Tagalog": "tl",
    "Tswana": "tn",
    "Turkish": "tr",
    "Ukrainian": "uk",
    "Urdu": "ur",
    "Uzbek": "uz",
    "Vietnamese": "vi",
    "Wolof": "wo",
    "Xhosa": "xh",
    "Yiddish": "yi",
    "Yoruba": "yo",
    "Chinese": "zh",
    "Zulu": "zu",
}

ocr_lang=[
'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',
    
]




def draw_boxes(image, bounds, color='blue', width=1):
    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 detect(img, target_lang):
    lang=[f"{lang_id[target_lang]}"]
    reader = easyocr.Reader(lang)
    bounds = reader.readtext(img)
    bounds2=reader.readtext(img)
    im = PIL.Image.open(img)
    im_out=draw_boxes(im, bounds)
    return im_out,pd.DataFrame(bounds),pd.DataFrame(bounds2)
with gr.Blocks() as robot:
    im=gr.Image(type="filepath")
    target_lang = gr.Dropdown(label="Detect language", choices=list(lang_id.keys()),value="English")
    
    go_btn=gr.Button()
    out_im=gr.Image()
    with gr.Row():
        out_txt=gr.Textbox(lines=8)
        data_f=gr.Dataframe()
    go_btn.click(detect,[im,target_lang],[out_im,out_txt,data_f])
robot.queue(concurrency_count=10).launch()