import gradio as gr
import warnings
import os
import pix2struct, layoutlm, donut
warnings.filterwarnings('ignore')


def process_image_and_generate_output(image, model_selection, question):
    result = ''
    if image is None:
        return "Please select an image", None

    if model_selection == "LayoutLM":
        result = layoutlm.get_result(image, question)
        return result
    if model_selection == 'Pix2Struct':
        result = pix2struct.get_result(image, question)
        return result
    if model_selection == 'Donut':
        result = donut.get_result(image, question)
        return result

    return result


sample_images = [
    [os.path.join(os.path.dirname(__file__), "images/1.png"), "LayoutLM", "What is the NIC Code?"],
    [os.path.join(os.path.dirname(__file__), "images/1.png"), "Pix2Struct", "What is the NIC Code?"],
    [os.path.join(os.path.dirname(__file__), "images/1.png"), "Donut", "What is the NIC Code?"]
]

# Create a dropdown to select sample image
image_input = gr.Image(label="Upload Image", type='filepath')

# Create a dropdown to choose the model
model_selection_input = gr.Radio(["LayoutLM", "Pix2Struct", "Donut"],
                                 label="Choose Model")
question_input = gr.Text(label="Question")

iface = gr.Interface(fn=process_image_and_generate_output,
                     inputs=[image_input, model_selection_input, question_input],
                     outputs=gr.Text(label="Result"),
                     allow_flagging='never',
                     examples=sample_images,
                     title="DocVQA Sanctum")

iface.launch()