import gradio as gr
import json

from ocr_text import main as ocr_main
from detect_layout import main as layout_main
from reading_order import main as reading_main

with open("languages.json", "r", encoding='utf-8') as file:
    language_map = json.load(file)

def model1(image_path, languages):
    langs = ""
    if languages == [] or not languages:
        langs = "English"
    else:
        for lang in languages:
            langs += f"{lang},"
        langs = langs[:-1]
    print(langs)
    annotated = ocr_main(image_path, langs=langs)
    return annotated

def model2(image_path):

    annotated = layout_main(image_path)
    return annotated

def model3(image_path):

    annotated = reading_main(image_path)
    return annotated

def process_image(image, model_func):
    if image is None:
        return None
    result = model_func(image)
    return result

with gr.Blocks() as demo:
    gr.Markdown("<center><h1>Surya - Image OCR/Layout/Reading Order</h1></center>")
    
    with gr.Row():
        with gr.Column():
            with gr.Row():
                input_image = gr.Image(type="filepath", label="Input Image", sources="upload")
            with gr.Row():
                dropdown = gr.Dropdown(label="Select Languages for OCR", choices=list(language_map.keys()), multiselect=True, value=["English"], interactive=True)
            with gr.Row():
                btn1 = gr.Button("OCR", variant="primary")
                btn2 = gr.Button("Layout", variant="primary")
                btn3 = gr.Button("Reading Order", variant="primary")
            with gr.Row():
                clear = gr.ClearButton()

        with gr.Column():
            with gr.Tabs():
                with gr.TabItem("OCR"):
                    output_image1 = gr.Image()
                with gr.TabItem("Layout"):
                    output_image2 = gr.Image()
                with gr.TabItem("Reading Order"):
                    output_image3 = gr.Image()

    btn1.click(fn=model1, inputs=[input_image, dropdown], outputs=output_image1)
    btn2.click(fn=model2, inputs=[input_image], outputs=output_image2)
    btn3.click(fn=model3, inputs=[input_image], outputs=output_image3)
    clear.add(components=[input_image, output_image1, output_image2, output_image3])

demo.launch()