laverdes commited on
Commit
df73b43
·
1 Parent(s): f7fe7ff

feat: select default image and parse receipt with U🍩, raw output

Browse files
Files changed (1) hide show
  1. app.py +12 -3
app.py CHANGED
@@ -6,7 +6,7 @@ from transformers import VisionEncoderDecoderModel, VisionEncoderDecoderConfig
6
 
7
 
8
  def demo_process(input_img):
9
- global pretrained_model, task_prompt, task_name
10
  # input_img = Image.fromarray(input_img)
11
  output = pretrained_model.inference(image=input_img, prompt=task_prompt)["predictions"][0]
12
  return output
@@ -22,9 +22,18 @@ with st.sidebar:
22
  information = st.radio(
23
  "What information inside the are you interested in?",
24
  ('Receipt Summary', 'Receipt Menu Details', 'Extract all!'))
25
- receipt = st.selectbox('Pick one receipt', ['1', '2', '3', '4', '5', '6'])
26
 
27
  st.text(f'{information} mode is ON!\nTarget receipt: {receipt}\n(opening image @:./img/receipt-{receipt}.png)')
28
 
29
  image = Image.open(f"./img/receipt-{receipt}.jpg")
30
- st.image(image, caption='Your target receipt')
 
 
 
 
 
 
 
 
 
 
6
 
7
 
8
  def demo_process(input_img):
9
+ global pretrained_model, task_prompt # , task_name
10
  # input_img = Image.fromarray(input_img)
11
  output = pretrained_model.inference(image=input_img, prompt=task_prompt)["predictions"][0]
12
  return output
 
22
  information = st.radio(
23
  "What information inside the are you interested in?",
24
  ('Receipt Summary', 'Receipt Menu Details', 'Extract all!'))
25
+ receipt = st.selectbox('Pick one receipt', ['1', '2', '3', '4', '5', '6'], index='6')
26
 
27
  st.text(f'{information} mode is ON!\nTarget receipt: {receipt}\n(opening image @:./img/receipt-{receipt}.png)')
28
 
29
  image = Image.open(f"./img/receipt-{receipt}.jpg")
30
+ st.image(image, caption='Your target receipt')
31
+
32
+ st.text(f'baking the 🍩...')
33
+ pretrained_model = VisionEncoderDecoderModel.from_pretrained("unstructured/donut-base-sroie")
34
+ pretrained_model.encoder.to(torch.bfloat16)
35
+ pretrained_model.eval()
36
+
37
+ st.text(f'parsing receipt..')
38
+ parsed_receipt_info = demo_process(image)
39
+ st.text(f'\nRaw output:\n{parsed_receipt_info}')