| import gradio as gr |
| import torch |
| from PIL import Image |
|
|
| from donut import DonutModel |
|
|
| def demo_process(input_img): |
| global pretrained_model, task_prompt, task_name |
| |
| try: |
| output = pretrained_model.inference(image=input_img, prompt=task_prompt)["predictions"][0] |
| except Exception as e: |
| return e |
| return output |
|
|
| task_prompt = f"<s_cord-v2>" |
|
|
| image = Image.open("./demo.png") |
| image.save("cord_sample_receipt1.png") |
|
|
| pretrained_model = DonutModel.from_pretrained("donut-base") |
| pretrained_model.eval() |
|
|
| demo = gr.Interface( |
| fn=demo_process, |
| inputs= gr.Image(type="pil"), |
| outputs="json", |
| title=f"Donut 🍩 demonstration for `cord-v2` task", |
| examples=[["cord_sample_receipt1.png"]], |
| cache_examples=False, |
| ) |
|
|
| demo.launch() |