Ankur Goyal commited on
Commit
194858a
1 Parent(s): 67ef815

Update title

Browse files
Files changed (1) hide show
  1. app.py +45 -14
app.py CHANGED
@@ -93,6 +93,7 @@ def process_path(path):
93
  gr.update(visible=True, value=document.preview),
94
  gr.update(visible=True),
95
  gr.update(visible=False, value=None),
 
96
  )
97
  except Exception:
98
  pass
@@ -101,6 +102,7 @@ def process_path(path):
101
  gr.update(visible=False, value=None),
102
  gr.update(visible=False),
103
  gr.update(visible=False, value=None),
 
104
  )
105
 
106
 
@@ -113,6 +115,7 @@ def process_upload(file):
113
  gr.update(visible=False, value=None),
114
  gr.update(visible=False),
115
  gr.update(visible=False, value=None),
 
116
  )
117
 
118
 
@@ -121,7 +124,7 @@ colors = ["#64A087", "green", "black"]
121
 
122
  def process_question(question, document, model=list(CHECKPOINTS.keys())[0]):
123
  if document is None:
124
- return None, None
125
 
126
  predictions = run_pipeline(model, question, document, 3)
127
  pages = [x.copy().convert("RGB") for x in document.preview]
@@ -143,16 +146,21 @@ def process_question(question, document, model=list(CHECKPOINTS.keys())[0]):
143
  )
144
  draw.rectangle(((x1, y1), (x2, y2)), fill=(0, 255, 0, int(0.4 * 255)))
145
 
146
- return gr.update(visible=True, value=pages), gr.update(
147
- visible=True, value=predictions
 
 
 
 
 
148
  )
149
 
150
 
151
  def load_example_document(img, question, model):
152
  if img is not None:
153
  document = ImageDocument(Image.fromarray(img))
154
- preview, answer = process_question(question, document, model)
155
- return document, question, preview, gr.update(visible=True), answer
156
  else:
157
  return None, None, None, gr.update(visible=False), None
158
 
@@ -202,7 +210,7 @@ CSS = """
202
  """
203
 
204
  with gr.Blocks(css=CSS) as demo:
205
- gr.Markdown("# DocQuery: Query Documents w/ NLP")
206
  gr.Markdown(
207
  "DocQuery uses LayoutLMv1 fine-tuned on DocVQA, a document visual question"
208
  " answering dataset, as well as SQuAD, which boosts its English-language comprehension."
@@ -261,6 +269,7 @@ with gr.Blocks(css=CSS) as demo:
261
  "Submit", variant="primary", elem_id="submit-button"
262
  )
263
  with gr.Column():
 
264
  output = gr.JSON(label="Output", visible=False)
265
 
266
  img_clear_button.click(
@@ -268,13 +277,23 @@ with gr.Blocks(css=CSS) as demo:
268
  gr.update(visible=False, value=None),
269
  None,
270
  gr.update(visible=False, value=None),
 
271
  gr.update(visible=False),
272
  None,
273
  None,
274
  None,
275
  ),
276
  inputs=img_clear_button,
277
- outputs=[image, document, output, img_clear_button, example_image, upload, url],
 
 
 
 
 
 
 
 
 
278
  )
279
  clear_button.click(
280
  lambda _: (
@@ -282,45 +301,57 @@ with gr.Blocks(css=CSS) as demo:
282
  None,
283
  None,
284
  gr.update(visible=False, value=None),
 
285
  None,
286
  None,
287
  None,
288
  ),
289
  inputs=clear_button,
290
- outputs=[image, document, question, output, example_image, upload, url],
 
 
 
 
 
 
 
 
 
291
  )
292
 
293
  upload.change(
294
  fn=process_upload,
295
  inputs=[upload],
296
- outputs=[document, image, img_clear_button, output],
297
  )
298
  url.change(
299
  fn=process_path,
300
  inputs=[url],
301
- outputs=[document, image, img_clear_button, output],
302
  )
303
 
304
  question.submit(
305
  fn=process_question,
306
  inputs=[question, document, model],
307
- outputs=[image, output],
308
  )
309
 
310
  submit_button.click(
311
  process_question,
312
  inputs=[question, document, model],
313
- outputs=[image, output],
314
  )
315
 
316
  model.change(
317
- process_question, inputs=[question, document, model], outputs=[image, output]
 
 
318
  )
319
 
320
  example_image.change(
321
  fn=load_example_document,
322
  inputs=[example_image, example_question, model],
323
- outputs=[document, question, image, img_clear_button, output],
324
  )
325
 
326
  if __name__ == "__main__":
93
  gr.update(visible=True, value=document.preview),
94
  gr.update(visible=True),
95
  gr.update(visible=False, value=None),
96
+ gr.update(visible=False, value=None),
97
  )
98
  except Exception:
99
  pass
102
  gr.update(visible=False, value=None),
103
  gr.update(visible=False),
104
  gr.update(visible=False, value=None),
105
+ gr.update(visible=False, value=None),
106
  )
107
 
108
 
115
  gr.update(visible=False, value=None),
116
  gr.update(visible=False),
117
  gr.update(visible=False, value=None),
118
+ gr.update(visible=False, value=None),
119
  )
120
 
121
 
124
 
125
  def process_question(question, document, model=list(CHECKPOINTS.keys())[0]):
126
  if document is None:
127
+ return None, None, None
128
 
129
  predictions = run_pipeline(model, question, document, 3)
130
  pages = [x.copy().convert("RGB") for x in document.preview]
146
  )
147
  draw.rectangle(((x1, y1), (x2, y2)), fill=(0, 255, 0, int(0.4 * 255)))
148
 
149
+ return (
150
+ gr.update(visible=True, value=pages),
151
+ gr.update(visible=True, value=predictions),
152
+ gr.update(
153
+ visible=True,
154
+ value=predictions[0]["answer"] if ensure_list(predictions) else None,
155
+ ),
156
  )
157
 
158
 
159
  def load_example_document(img, question, model):
160
  if img is not None:
161
  document = ImageDocument(Image.fromarray(img))
162
+ preview, answer, answer_text = process_question(question, document, model)
163
+ return document, question, preview, gr.update(visible=True), answer, answer_text
164
  else:
165
  return None, None, None, gr.update(visible=False), None
166
 
210
  """
211
 
212
  with gr.Blocks(css=CSS) as demo:
213
+ gr.Markdown("# DocQuery: Document Query Engine")
214
  gr.Markdown(
215
  "DocQuery uses LayoutLMv1 fine-tuned on DocVQA, a document visual question"
216
  " answering dataset, as well as SQuAD, which boosts its English-language comprehension."
269
  "Submit", variant="primary", elem_id="submit-button"
270
  )
271
  with gr.Column():
272
+ output_text = gr.Textbox(label="Top Answer", visible=False)
273
  output = gr.JSON(label="Output", visible=False)
274
 
275
  img_clear_button.click(
277
  gr.update(visible=False, value=None),
278
  None,
279
  gr.update(visible=False, value=None),
280
+ gr.update(visible=False, value=None),
281
  gr.update(visible=False),
282
  None,
283
  None,
284
  None,
285
  ),
286
  inputs=img_clear_button,
287
+ outputs=[
288
+ image,
289
+ document,
290
+ output,
291
+ output_text,
292
+ img_clear_button,
293
+ example_image,
294
+ upload,
295
+ url,
296
+ ],
297
  )
298
  clear_button.click(
299
  lambda _: (
301
  None,
302
  None,
303
  gr.update(visible=False, value=None),
304
+ gr.update(visible=False, value=None),
305
  None,
306
  None,
307
  None,
308
  ),
309
  inputs=clear_button,
310
+ outputs=[
311
+ image,
312
+ document,
313
+ question,
314
+ output,
315
+ output_text,
316
+ example_image,
317
+ upload,
318
+ url,
319
+ ],
320
  )
321
 
322
  upload.change(
323
  fn=process_upload,
324
  inputs=[upload],
325
+ outputs=[document, image, img_clear_button, output, output_text],
326
  )
327
  url.change(
328
  fn=process_path,
329
  inputs=[url],
330
+ outputs=[document, image, img_clear_button, output, output_text],
331
  )
332
 
333
  question.submit(
334
  fn=process_question,
335
  inputs=[question, document, model],
336
+ outputs=[image, output, output_text],
337
  )
338
 
339
  submit_button.click(
340
  process_question,
341
  inputs=[question, document, model],
342
+ outputs=[image, output, output_text],
343
  )
344
 
345
  model.change(
346
+ process_question,
347
+ inputs=[question, document, model],
348
+ outputs=[image, output, output_text],
349
  )
350
 
351
  example_image.change(
352
  fn=load_example_document,
353
  inputs=[example_image, example_question, model],
354
+ outputs=[document, question, image, img_clear_button, output, output_text],
355
  )
356
 
357
  if __name__ == "__main__":