sav7669 commited on
Commit
a92be85
1 Parent(s): 2838f66

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +4 -8
app.py CHANGED
@@ -10,16 +10,12 @@ from PIL import Image
10
 
11
  from donut import DonutModel
12
  from transformers import DonutProcessor, VisionEncoderDecoderModel
13
- from huggingface_hub import notebook_login
14
-
15
- HF_TOKEN = os.environ.get("HF_TOKEN", None)
16
- huggingface_hub.notebook_login(token=HF_TOKEN)
17
 
18
 
19
  def demo_process(input_img):
20
  global pretrained_model, task_prompt, task_name
21
  # input_img = Image.fromarray(input_img)
22
- output = pretrained_model.inference(image=input_img, prompt=task_prompt)["predictions"][0]
23
  return output
24
 
25
  task_prompt = f"<s_cord-v2>"
@@ -34,9 +30,9 @@ image = Image.open("./001.jpg")
34
  image.save("sroie_test2.jpg")
35
 
36
  #philschmid
37
- pretrained_model = DonutModel.from_pretrained("philschmid/donut-base-sroie")
38
- #model = VisionEncoderDecoderModel.from_pretrained("philschmid/donut-base-sroie")
39
- pretrained_model.eval()
40
 
41
  demo = gr.Interface(
42
  fn=demo_process,
 
10
 
11
  from donut import DonutModel
12
  from transformers import DonutProcessor, VisionEncoderDecoderModel
 
 
 
 
13
 
14
 
15
  def demo_process(input_img):
16
  global pretrained_model, task_prompt, task_name
17
  # input_img = Image.fromarray(input_img)
18
+ output = model.inference(image=input_img, prompt=task_prompt)["predictions"][0]
19
  return output
20
 
21
  task_prompt = f"<s_cord-v2>"
 
30
  image.save("sroie_test2.jpg")
31
 
32
  #philschmid
33
+ #pretrained_model = DonutModel.from_pretrained("philschmid/donut-base-sroie")
34
+ model = VisionEncoderDecoderModel.from_pretrained("philschmid/donut-base-sroie")
35
+ model.eval()
36
 
37
  demo = gr.Interface(
38
  fn=demo_process,