Ankur Goyal commited on
Commit
87500f1
1 Parent(s): 15fad86

Support multi-page preview

Browse files
Files changed (2) hide show
  1. app.py +11 -8
  2. requirements.txt +1 -1
app.py CHANGED
@@ -45,8 +45,8 @@ def run_pipeline(model, question, document, top_k):
45
 
46
  # TODO: Move into docquery
47
  # TODO: Support words past the first page (or window?)
48
- def lift_word_boxes(document):
49
- return document.context["image"][0][1]
50
 
51
 
52
  def expand_bbox(word_boxes):
@@ -122,23 +122,26 @@ def process_question(question, document, model=list(CHECKPOINTS.keys())[0]):
122
  return None, None
123
 
124
  predictions = run_pipeline(model, question, document, 3)
125
- image = document.preview.copy()
126
- draw = ImageDraw.Draw(image, "RGBA")
127
  for i, p in enumerate(ensure_list(predictions)):
128
  if i > 0:
129
  # Keep the code around to produce multiple boxes, but only show the top
130
  # prediction for now
131
  break
132
 
133
- if "start" in p and "end" in p and p.get("page") == 0:
 
 
134
  x1, y1, x2, y2 = normalize_bbox(
135
- expand_bbox(lift_word_boxes(document)[p["start"] : p["end"] + 1]),
 
 
136
  image.width,
137
  image.height,
138
  )
139
  draw.rectangle(((x1, y1), (x2, y2)), fill=(0, 255, 0, int(0.4 * 255)))
140
 
141
- return gr.update(visible=True, value=image), gr.update(
142
  visible=True, value=predictions
143
  )
144
 
@@ -197,7 +200,7 @@ with gr.Blocks(css=CSS) as demo:
197
  submit_button = gr.Button("Submit", variant="primary", elem_id="submit-button")
198
 
199
  with gr.Row():
200
- image = gr.Image(visible=False)
201
  with gr.Column():
202
  output = gr.JSON(label="Output", visible=False)
203
 
 
45
 
46
  # TODO: Move into docquery
47
  # TODO: Support words past the first page (or window?)
48
+ def lift_word_boxes(document, page):
49
+ return document.context["image"][page][1]
50
 
51
 
52
  def expand_bbox(word_boxes):
 
122
  return None, None
123
 
124
  predictions = run_pipeline(model, question, document, 3)
125
+ pages = [x.copy() for x in document.preview]
 
126
  for i, p in enumerate(ensure_list(predictions)):
127
  if i > 0:
128
  # Keep the code around to produce multiple boxes, but only show the top
129
  # prediction for now
130
  break
131
 
132
+ if "start" in p and "end" in p:
133
+ image = pages[p["page"]]
134
+ draw = ImageDraw.Draw(image, "RGBA")
135
  x1, y1, x2, y2 = normalize_bbox(
136
+ expand_bbox(
137
+ lift_word_boxes(document, p["page"])[p["start"] : p["end"] + 1]
138
+ ),
139
  image.width,
140
  image.height,
141
  )
142
  draw.rectangle(((x1, y1), (x2, y2)), fill=(0, 255, 0, int(0.4 * 255)))
143
 
144
+ return gr.update(visible=True, value=pages), gr.update(
145
  visible=True, value=predictions
146
  )
147
 
 
200
  submit_button = gr.Button("Submit", variant="primary", elem_id="submit-button")
201
 
202
  with gr.Row():
203
+ image = gr.Gallery(visible=False)
204
  with gr.Column():
205
  output = gr.JSON(label="Output", visible=False)
206
 
requirements.txt CHANGED
@@ -1,4 +1,4 @@
1
  torch
2
  git+https://github.com/huggingface/transformers.git@21f6f58721dd9154357576be6de54eefef1f1818
3
- git+https://github.com/impira/docquery.git@5de147aa7df285829a8d064f3ca2c7ac74efe0f8
4
  sentencepiece
 
1
  torch
2
  git+https://github.com/huggingface/transformers.git@21f6f58721dd9154357576be6de54eefef1f1818
3
+ git+https://github.com/impira/docquery.git@3aa3cc6ca6624d6371db9cd4732cbbcb9c8b3ea0
4
  sentencepiece