Update app.py

app.py

@@ -9,26 +9,28 @@ import torch
 from PIL import Image
 
 from donut import DonutModel
-from transformers import DonutProcessor, VisionEncoderDecoderModel
-
 
 def demo_process(input_img):
     global pretrained_model, task_prompt, task_name
     # input_img = Image.fromarray(input_img)
-    output = model.inference(image=input_img, prompt=task_prompt)["predictions"][0]
+    output = pretrained_model.inference(image=input_img, prompt=task_prompt)["predictions"][0]
     return output
 
 task_prompt = f"<s_cord-v2>"
 
 image = Image.open("./cord_sample_receipt1.png")
 image.save("cord_sample_receipt1.png")
+image = Image.open("./cord_sample_receipt2.png")
+image.save("cord_sample_receipt2.png")
+image = Image.open("./receipt_00015.png")
+image.save("sroie_receipt_00015.png")
+image = Image.open("./receipt_00049.png")
+image.save("sroie_receipt_00049.png")
 image = Image.open("./000.jpg")
-image.save("sroie_sample_receipt1.jpg")
+image.save("sroie_receipt_00048.jpg")
 
-#philschmid
-#pretrained_model = DonutModel.from_pretrained("philschmid/donut-base-sroie")
-model = VisionEncoderDecoderModel.from_pretrained("philschmid/donut-base-sroie")
-model.eval()
+pretrained_model = DonutModel.from_pretrained("naver-clova-ix/donut-base-finetuned-cord-v2")
+pretrained_model.eval()
 
 demo = gr.Interface(
     fn=demo_process,
@@ -41,6 +43,8 @@ More CORD receipt images are available at https://huggingface.co/datasets/naver-
 More details are available at:
 - Paper: https://arxiv.org/abs/2111.15664
 - GitHub: https://github.com/clovaai/donut""",
-    examples=[["cord_sample_receipt1.png"],["sroie_sample_receipt1.jpg"]],
+    examples=[["cord_sample_receipt1.png"], ["cord_sample_receipt2.png"]],
     cache_examples=False,
-)
+)
+
+demo.launch()