import torch
from sconf import Config
from PIL import Image, ImageOps
from donut import DonutConfig, DonutModel
import warnings

warnings.filterwarnings("ignore")

from transformers import logging

logging.set_verbosity_warning()

config = Config(default="./config.yaml")

model = DonutModel.from_pretrained(
                config.pretrained_model_name_or_path,
                input_size=config.input_size,
                max_length=config.max_position_embeddings, #self.config.max_length,
                align_long_axis=config.align_long_axis,
                ignore_mismatched_sizes=True,
            )

task_name = "matricula"
task_prompt = f"<s_{task_name}>"
    
def predict_matricula(model, task_name, image):
    image = ImageOps.exif_transpose(image)
    device = "cuda" if torch.cuda.is_available() else "cpu"

    model.eval()
    model.to(device)

    result = model.inference(image=image, prompt=f"<s_{task_name}>")["predictions"][0]
    return result


import gradio as gr

with gr.Blocks() as demo:
    fn_predict=lambda x:predict_matricula(model, task_name="matricula", image=x)
    gr.Markdown(
    """
    # Demo: Donut 🍩 for DR Matriculas
    Dominican Vehicle **Matriculas OCR** Infering
    """)
    with gr.Row():
        input_image=gr.Image(label="Matricula", sources="upload", type="pil", show_label=True)
    with gr.Row():
        output_json=gr.JSON(label="Matricula JSON", show_label=True, value={})
    submit_btn = gr.Button("Submit")
    submit_btn.click(fn=fn_predict, inputs=[input_image], outputs=[output_json])

demo.launch(share=True)