Update app.py

app.py

@@ -4,8 +4,6 @@ import torch
 from PIL import Image
 import re
 from transformers import DonutProcessor, VisionEncoderDecoderModel
-import json
-
 
 def demo_process(input_img, question=None):
     global processor, model
@@ -22,17 +20,17 @@ def demo_process(input_img, question=None):
 
     with torch.no_grad():
         outputs = model.generate(
-          pixel_values,
-          decoder_input_ids=decoder_input_ids,
-          max_length=1024,  # เปลี่ยนตามความต้องการ
-          early_stopping=True,
-          pad_token_id=processor.tokenizer.pad_token_id,
-          eos_token_id=processor.tokenizer.eos_token_id,
-          use_cache=True,
-          num_beams=1,
-          bad_words_ids=[[processor.tokenizer.unk_token_id]],
-          return_dict_in_generate=True,
-      )
+            pixel_values,
+            decoder_input_ids=decoder_input_ids,
+            max_length=1024,  
+            early_stopping=True,
+            pad_token_id=processor.tokenizer.pad_token_id,
+            eos_token_id=processor.tokenizer.eos_token_id,
+            use_cache=True,
+            num_beams=1,
+            bad_words_ids=[[processor.tokenizer.unk_token_id]],
+            return_dict_in_generate=True,
+        )
 
     seq = processor.batch_decode(outputs.sequences)[0]
     seq = seq.replace(processor.tokenizer.eos_token, "").replace(processor.tokenizer.pad_token, "")
@@ -40,9 +38,8 @@ def demo_process(input_img, question=None):
     seq = processor.token2json(seq)
     return seq
 
-
 parser = argparse.ArgumentParser()
-parser.add_argument("--task", type=str, default="cord-v2")  # Add argument for task
+parser.add_argument("--task", type=str, default="cord-v2")
 parser.add_argument("--pretrained_path", type=str, default="suthawadee/donut-demo_new")
 args, left_argv = parser.parse_known_args()
 
@@ -52,14 +49,24 @@ device = "cpu" if not torch.cuda.is_available() else "cuda"
 model.to(device)
 model.eval()
 
-inputs = ["image", "text"] if args.task == "docvqa" else "image"
-outputs = "json"
+# เพิ่มตัวอย่างรูปภาพที่มีอยู่เพื่อทดสอบ
+image1 = "8.jpg"
+image2 = "15.jpg"
+
+examples = [
+    [Image.open(image1)],
+    [Image.open(image2)]
+]
+
+def main(pretrained_path, examples):
+    demo = gr.Interface(
+        fn=demo_process,
+        inputs=["image", "text"] if args.task == "docvqa" else "image",
+        outputs="json",
+        title="🇹🇭🧾ThaiReceipt",
+        description="Upload image.",
+        examples=examples
+    )
+    demo.launch(debug=True)
 
-demo = gr.Interface(
-    fn=demo_process,
-    inputs=inputs,
-    outputs=outputs,
-    title="🇹🇭🧾ThaiReceipt",
-    description="Upload an image."
-)
-demo.launch(debug=True)
+main(args.pretrained_path, examples)