Msp/Document_Parser
7
1import os2 3os.environ["TOKENIZERS_PARALLELISM"] = "false"4 5from PIL import Image, ImageDraw6import traceback7 8import gradio as gr9 10import torch11from docquery import pipeline12from docquery.document import load_document, ImageDocument13from docquery.ocr_reader import get_ocr_reader14 15 16def ensure_list(x):17 if isinstance(x, list):18 return x19 else:20 return [x]21 22 23CHECKPOINTS = {24 "LayoutLMv1 🦉": "impira/layoutlm-document-qa",25 "LayoutLMv1 for Invoices 💸": "impira/layoutlm-invoices",26 "Donut 🍩": "naver-clova-ix/donut-base-finetuned-docvqa",27}28 29PIPELINES = {}30 31 32def construct_pipeline(task, model):33 global PIPELINES34 if model in PIPELINES:35 return PIPELINES[model]36 37 device = "cuda" if torch.cuda.is_available() else "cpu"38 ret = pipeline(task=task, model=CHECKPOINTS[model], device=device)39 PIPELINES[model] = ret40 return ret41 42 43def run_pipeline(model, question, document, top_k):44 pipeline = construct_pipeline("document-question-answering", model)45 return pipeline(question=question, **document.context, top_k=top_k)46 47 48# TODO: Move into docquery49# TODO: Support words past the first page (or window?)50def lift_word_boxes(document, page):51 return document.context["image"][page][1]52 53 54def expand_bbox(word_boxes):55 if len(word_boxes) == 0:56 return None57 58 min_x, min_y, max_x, max_y = zip(*[x[1] for x in word_boxes])59 min_x, min_y, max_x, max_y = [min(min_x), min(min_y), max(max_x), max(max_y)]60 return [min_x, min_y, max_x, max_y]61 62 63# LayoutLM boxes are normalized to 0, 100064def normalize_bbox(box, width, height, padding=0.005):65 min_x, min_y, max_x, max_y = [c / 1000 for c in box]66 if padding != 0:67 min_x = max(0, min_x - padding)68 min_y = max(0, min_y - padding)69 max_x = min(max_x + padding, 1)70 max_y = min(max_y + padding, 1)71 return [min_x * width, min_y * height, max_x * width, max_y * height]72 73 74examples = [75 [76 "invoice.png",77 "What is the invoice number?",78 ],79 [80 "contract.jpeg",81 "What is the purchase amount?",82 ],83 [84 "statement.png",85 "What are net sales for 2020?",86 ],87 88]89 90question_files = {91 "What are net sales for 2020?": "statement.pdf",92 "How many likes does the space have?": "https://huggingface.co/spaces/impira/docquery",93 "What is the title of post number 5?": "https://news.ycombinator.com",94}95 96 97def process_path(path):98 error = None99 if path:100 try:101 document = load_document(path)102 return (103 document,104 gr.update(visible=True, value=document.preview),105 gr.update(visible=True),106 gr.update(visible=False, value=None),107 gr.update(visible=False, value=None),108 None,109 )110 except Exception as e:111 traceback.print_exc()112 error = str(e)113 return (114 None,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 gr.update(visible=True, value=error) if error is not None else None,120 None,121 )122 123 124def process_upload(file):125 if file:126 return process_path(file.name)127 else:128 return (129 None,130 gr.update(visible=False, value=None),131 gr.update(visible=False),132 gr.update(visible=False, value=None),133 gr.update(visible=False, value=None),134 None,135 )136 137 138colors = ["#64A087", "green", "black"]139 140 141def process_question(question, document, model=list(CHECKPOINTS.keys())[0]):142 if not question or document is None:143 return None, None, None144 145 text_value = None146 predictions = run_pipeline(model, question, document, 3)147 pages = [x.copy().convert("RGB") for x in document.preview]148 for i, p in enumerate(ensure_list(predictions)):149 if i == 0:150 text_value = p["answer"]151 else:152 # Keep the code around to produce multiple boxes, but only show the top153 # prediction for now154 break155 156 if "word_ids" in p:157 image = pages[p["page"]]158 draw = ImageDraw.Draw(image, "RGBA")159 word_boxes = lift_word_boxes(document, p["page"])160 x1, y1, x2, y2 = normalize_bbox(161 expand_bbox([word_boxes[i] for i in p["word_ids"]]),162 image.width,163 image.height,164 )165 draw.rectangle(((x1, y1), (x2, y2)), fill=(0, 255, 0, int(0.4 * 255)))166 167 return (168 gr.update(visible=True, value=pages),169 gr.update(visible=True, value=predictions),170 gr.update(171 visible=True,172 value=text_value,173 ),174 )175 176 177def load_example_document(img, question, model):178 if img is not None:179 if question in question_files:180 document = load_document(question_files[question])181 else:182 document = ImageDocument(Image.fromarray(img), get_ocr_reader())183 preview, answer, answer_text = process_question(question, document, model)184 return document, question, preview, gr.update(visible=True), answer, answer_text185 else:186 return None, None, None, gr.update(visible=False), None, None187 188 189CSS = """190#question input {191 font-size: 16px;192}193#url-textbox {194 padding: 0 !important;195}196#short-upload-box .w-full {197 min-height: 10rem !important;198}199/* I think something like this can be used to re-shape200 * the table201 */202/*203.gr-samples-table tr {204 display: inline;205}206.gr-samples-table .p-2 {207 width: 100px;208}209*/210#select-a-file {211 width: 100%;212}213#file-clear {214 padding-top: 2px !important;215 padding-bottom: 2px !important;216 padding-left: 8px !important;217 padding-right: 8px !important;218 margin-top: 10px;219}220.gradio-container .gr-button-primary {221 background: linear-gradient(180deg, #CDF9BE 0%, #AFF497 100%);222 border: 1px solid #B0DCCC;223 border-radius: 8px;224 color: #1B8700;225}226.gradio-container.dark button#submit-button {227 background: linear-gradient(180deg, #CDF9BE 0%, #AFF497 100%);228 border: 1px solid #B0DCCC;229 border-radius: 8px;230 color: #1B8700231}232 233table.gr-samples-table tr td {234 border: none;235 outline: none;236}237 238table.gr-samples-table tr td:first-of-type {239 width: 0%;240}241 242div#short-upload-box div.absolute {243 display: none !important;244}245 246gradio-app > div > div > div > div.w-full > div, .gradio-app > div > div > div > div.w-full > div {247 gap: 0px 2%;248}249 250gradio-app div div div div.w-full, .gradio-app div div div div.w-full {251 gap: 0px;252}253 254gradio-app h2, .gradio-app h2 {255 padding-top: 10px;256}257 258#answer {259 overflow-y: scroll;260 color: white;261 background: #666;262 border-color: #666;263 font-size: 20px;264 font-weight: bold;265}266 267#answer span {268 color: white;269}270 271#answer textarea {272 color:white;273 background: #777;274 border-color: #777;275 font-size: 18px;276}277 278#url-error input {279 color: red;280}281"""282 283with gr.Blocks(css=CSS) as demo:284 gr.Markdown("# Document Parser: Document Parser Engine")285 gr.Markdown(286 "Document_Parser is built on top of DocQuery library)"287 " uses LayoutLMv1 fine-tuned on DocVQA, a document visual question"288 " answering dataset, as well as SQuAD, which boosts its English-language comprehension."289 290 )291 292 document = gr.Variable()293 example_question = gr.Textbox(visible=False)294 example_image = gr.Image(visible=False)295 296 with gr.Row(equal_height=True):297 with gr.Column():298 with gr.Row():299 gr.Markdown("## 1. Select a file", elem_id="select-a-file")300 img_clear_button = gr.Button(301 "Clear", variant="secondary", elem_id="file-clear", visible=False302 )303 image = gr.Gallery(visible=False)304 with gr.Row(equal_height=True):305 with gr.Column():306 with gr.Row():307 url = gr.Textbox(308 show_label=False,309 placeholder="URL",310 lines=1,311 max_lines=1,312 elem_id="url-textbox",313 )314 submit = gr.Button("Get")315 url_error = gr.Textbox(316 visible=False,317 elem_id="url-error",318 max_lines=1,319 interactive=False,320 label="Error",321 )322 gr.Markdown("— or —")323 upload = gr.File(label=None, interactive=True, elem_id="short-upload-box")324 gr.Examples(325 examples=examples,326 inputs=[example_image, example_question],327 )328 329 with gr.Column() as col:330 gr.Markdown("## 2. Ask a question")331 question = gr.Textbox(332 label="Question",333 placeholder="e.g. What is the invoice number?",334 lines=1,335 max_lines=1,336 )337 model = gr.Radio(338 choices=list(CHECKPOINTS.keys()),339 value=list(CHECKPOINTS.keys())[0],340 label="Model",341 )342 343 with gr.Row():344 clear_button = gr.Button("Clear", variant="secondary")345 submit_button = gr.Button(346 "Submit", variant="primary", elem_id="submit-button"347 )348 with gr.Column():349 output_text = gr.Textbox(350 label="Top Answer", visible=False, elem_id="answer"351 )352 output = gr.JSON(label="Output", visible=False)353 354 for cb in [img_clear_button, clear_button]:355 cb.click(356 lambda _: (357 gr.update(visible=False, value=None),358 None,359 gr.update(visible=False, value=None),360 gr.update(visible=False, value=None),361 gr.update(visible=False),362 None,363 None,364 None,365 gr.update(visible=False, value=None),366 None,367 ),368 inputs=clear_button,369 outputs=[370 image,371 document,372 output,373 output_text,374 img_clear_button,375 example_image,376 upload,377 url,378 url_error,379 question,380 ],381 )382 383 upload.change(384 fn=process_upload,385 inputs=[upload],386 outputs=[document, image, img_clear_button, output, output_text, url_error],387 )388 submit.click(389 fn=process_path,390 inputs=[url],391 outputs=[document, image, img_clear_button, output, output_text, url_error],392 )393 394 question.submit(395 fn=process_question,396 inputs=[question, document, model],397 outputs=[image, output, output_text],398 )399 400 submit_button.click(401 process_question,402 inputs=[question, document, model],403 outputs=[image, output, output_text],404 )405 406 model.change(407 process_question,408 inputs=[question, document, model],409 outputs=[image, output, output_text],410 )411 412 example_image.change(413 fn=load_example_document,414 inputs=[example_image, example_question, model],415 outputs=[document, question, image, img_clear_button, output, output_text],416 )417 418if __name__ == "__main__":419 demo.launch(enable_queue=False)420