import gradio as gr
import argparse
import torch
from PIL import Image
from donut import DonutModel
def demo_process(input_img):
    global model, task_prompt, task_name
    input_img = Image.fromarray(input_img)
    output = model.inference(image=input_img, prompt=task_prompt)["predictions"][0]
    return output
parser = argparse.ArgumentParser()
parser.add_argument("--task", type=str, default="Booking")
parser.add_argument("--pretrained_path", type=str, default="uartimcs/donut-booking-extract")
args, left_argv = parser.parse_known_args()
task_name = args.task
task_prompt = f"<s_{task_name}>"

image = Image.open("./sample-booking/CMA_150.jpg")
image.save("CMA_sample.jpg")
image = Image.open("./sample-booking/COSCO_150.jpg")
image.save("COSCO_sample.jpg")
image = Image.open("./sample-booking/ONEY_150.jpg")
image.save("ONEY_sample.jpg")


model = DonutModel.from_pretrained("uartimcs/donut-booking-extract")
model.eval()
demo = gr.Interface(fn=demo_process,inputs="image",outputs="json", title=f"Donut 🍩 demonstration for `{task_name}` task", examples=[["CMA_sample.jpg"], ["COSCO_sample.jpg"], ["ONEY_sample.jpg"]],)
demo.launch()