Ankur Goyal commited on
Commit
cf2cdaf
1 Parent(s): 2f6c963

Include invoices

Browse files
Files changed (1) hide show
  1. app.py +4 -4
app.py CHANGED
@@ -21,6 +21,7 @@ def ensure_list(x):
21
 
22
  CHECKPOINTS = {
23
  "LayoutLMv1 🦉": "impira/layoutlm-document-qa",
 
24
  "Donut 🍩": "naver-clova-ix/donut-base-finetuned-docvqa",
25
  }
26
 
@@ -148,13 +149,12 @@ def process_question(question, document, model=list(CHECKPOINTS.keys())[0]):
148
  # prediction for now
149
  break
150
 
151
- if "start" in p and "end" in p:
152
  image = pages[p["page"]]
153
  draw = ImageDraw.Draw(image, "RGBA")
 
154
  x1, y1, x2, y2 = normalize_bbox(
155
- expand_bbox(
156
- lift_word_boxes(document, p["page"])[p["start"] : p["end"] + 1]
157
- ),
158
  image.width,
159
  image.height,
160
  )
 
21
 
22
  CHECKPOINTS = {
23
  "LayoutLMv1 🦉": "impira/layoutlm-document-qa",
24
+ "LayoutLMv1 for Invoices 🧾": "impira/layoutlm-invoices",
25
  "Donut 🍩": "naver-clova-ix/donut-base-finetuned-docvqa",
26
  }
27
 
 
149
  # prediction for now
150
  break
151
 
152
+ if "word_ids" in p:
153
  image = pages[p["page"]]
154
  draw = ImageDraw.Draw(image, "RGBA")
155
+ word_boxes = lift_word_boxes(document, p["page"])
156
  x1, y1, x2, y2 = normalize_bbox(
157
+ expand_bbox([word_boxes[i] for i in p["word_ids"]]),
 
 
158
  image.width,
159
  image.height,
160
  )