May remove the fields variable

app.py

@@ -6,6 +6,7 @@ from PIL import Image, ImageDraw
 import traceback
 
 import gradio as gr
+from gradio import processing_utils
 
 import torch
 from docquery import pipeline
@@ -99,16 +100,17 @@ FIELDS = {
     "Payment Terms": ["Payment Terms?"],
 }
 
-EMPTY_TABLE = dict(
-    headers=["Field", "Value"], value=[[name, None] for name in FIELDS.keys()]
-)
 
+def empty_table(fields):
+    return {"value": [[name, None] for name in fields.keys()], "interactive": False}
 
-def process_document(document, model, error=None):
+
+def process_document(document, fields, model, error=None):
     if document is not None and error is None:
-        preview, json_output, table = process_fields(document, model)
+        preview, json_output, table = process_fields(document, fields, model)
         return (
             document,
+            fields,
             preview,
             gr.update(visible=True),
             gr.update(visible=False, value=None),
@@ -118,6 +120,7 @@ def process_document(document, model, error=None):
     else:
         return (
             None,
+            fields,
             None,
             gr.update(visible=False),
             gr.update(visible=True, value=error) if error is not None else None,
@@ -129,6 +132,7 @@ def process_document(document, model, error=None):
 def process_path(path, model):
     error = None
     document = None
+    fields = {**FIELDS}
     if path:
         try:
             document = load_document(path)
@@ -136,7 +140,7 @@ def process_path(path, model):
             traceback.print_exc()
             error = str(e)
 
-    return process_document(document, model, error)
+    return process_document(document, fields, model, error)
 
 
 def process_upload(file, model):
@@ -159,40 +163,36 @@ def annotate_page(prediction, pages, document):
         draw.rectangle(((x1, y1), (x2, y2)), fill=(0, 255, 0, int(0.4 * 255)))
 
 
-def process_question(question, document, model=list(CHECKPOINTS.keys())[0]):
+def process_question(
+    question, document, img_gallery, model, fields, output, output_table
+):
     if not question or document is None:
         return None, None, None
 
     text_value = None
-    predictions = run_pipeline(model, question, document, 3)
-    pages = [x.copy().convert("RGB") for x in document.preview]
-    for i, p in enumerate(ensure_list(predictions)):
-        if i == 0:
-            text_value = p["answer"]
-        else:
-            # Keep the code around to produce multiple boxes, but only show the top
-            # prediction for now
-            break
-
-        annotate_page(p, pages, document)
-
+    pages = [processing_utils.decode_base64_to_image(p) for p in img_gallery]
+    prediction = run_pipeline(model, question, document, 1)
+    annotate_page(prediction, pages, document)
+
+    field_name = question.rstrip("?")
+    fields = {**FIELDS, field_name: [question]}
+    output[field_name] = prediction
+    table = output_table.values.tolist() + [[field_name, prediction.get("answer")]]
     return (
         gr.update(visible=True, value=pages),
-        gr.update(visible=True, value=predictions),
-        gr.update(
-            visible=True,
-            value=text_value,
-        ),
+        fields,
+        output,
+        gr.update(value=table, interactive=False),
     )
 
 
-def process_fields(document, model=list(CHECKPOINTS.keys())[0]):
+def process_fields(document, fields, model=list(CHECKPOINTS.keys())[0]):
     pages = [x.copy().convert("RGB") for x in document.preview]
 
     ret = {}
     table = []
 
-    for (field_name, questions) in FIELDS.items():
+    for (field_name, questions) in fields.items():
         answers = [run_pipeline(model, q, document, top_k=1) for q in questions]
         answers.sort(key=lambda x: -x.get("score", 0) if x else 0)
         top = answers[0]
@@ -208,23 +208,22 @@ def process_fields(document, model=list(CHECKPOINTS.keys())[0]):
 
 
 def load_example_document(img, title, model):
+    document = None
+    fields = {**FIELDS}
     if img is not None:
         if title in QUESTION_FILES:
-            print("using document")
             document = load_document(QUESTION_FILES[title])
         else:
             document = ImageDocument(Image.fromarray(img), ocr_reader=get_ocr_reader())
-    else:
-        document = None
 
-    return process_document(document, model)
+    return process_document(document, fields, model)
 
 
 CSS = """
 #question input {
     font-size: 16px;
 }
-#url-textbox {
+#url-textbox, #question-textbox {
     padding: 0 !important;
 }
 #short-upload-box .w-full {
@@ -327,6 +326,7 @@ with gr.Blocks(css=CSS) as demo:
     )
 
     document = gr.Variable()
+    fields = gr.Variable(value={**FIELDS})
     example_question = gr.Textbox(visible=False)
     example_image = gr.Image(visible=False)
 
@@ -364,13 +364,16 @@ with gr.Blocks(css=CSS) as demo:
             )
 
         with gr.Column() as col:
+            with gr.Tabs():
+                with gr.TabItem("Table"):
+                    output_table = gr.Dataframe(
+                        headers=["Field", "Value"], **empty_table(fields.value)
+                    )
+
+                with gr.TabItem("JSON"):
+                    output = gr.JSON(label="Output", visible=False)
+
             gr.Markdown("## 2. Ask a question")
-            question = gr.Textbox(
-                label="Question",
-                placeholder="e.g. What is the invoice number?",
-                lines=1,
-                max_lines=1,
-            )
             model = gr.Radio(
                 choices=list(CHECKPOINTS.keys()),
                 value=list(CHECKPOINTS.keys())[0],
@@ -379,24 +382,27 @@ with gr.Blocks(css=CSS) as demo:
             )
 
             with gr.Row():
-                clear_button = gr.Button("Clear", variant="secondary")
+                question = gr.Textbox(
+                    label="Question",
+                    show_label=False,
+                    placeholder="e.g. What is the invoice number?",
+                    lines=1,
+                    max_lines=1,
+                    elem_id="question-textbox",
+                )
+                clear_button = gr.Button("Clear", variant="secondary", visible=False)
                 submit_button = gr.Button(
-                    "Submit", variant="primary", elem_id="submit-button"
+                    "Add", variant="primary", elem_id="submit-button"
                 )
-            with gr.Tabs():
-                with gr.TabItem("Table"):
-                    output_table = gr.Dataframe(**EMPTY_TABLE)
-
-                with gr.TabItem("JSON"):
-                    output = gr.JSON(label="Output", visible=False)
 
     for cb in [img_clear_button, clear_button]:
         cb.click(
             lambda _: (
-                gr.update(visible=False, value=None), # image
-                None, # document 
-                gr.update(visible=False, value=None), # output
-                gr.update(**EMPTY_TABLE), # output_table
+                gr.update(visible=False, value=None),  # image
+                None,  # document
+                {**FIELDS},  # fields
+                gr.update(visible=False, value=None),  # output
+                gr.update(**empty_table(FIELDS)),  # output_table
                 gr.update(visible=False),
                 None,
                 None,
@@ -408,6 +414,7 @@ with gr.Blocks(css=CSS) as demo:
             outputs=[
                 image,
                 document,
+                fields,
                 output,
                 output_table,
                 img_clear_button,
@@ -419,22 +426,32 @@ with gr.Blocks(css=CSS) as demo:
             ],
         )
 
+    submit_outputs = [
+        document,
+        fields,
+        image,
+        img_clear_button,
+        url_error,
+        output,
+        output_table,
+    ]
+
     upload.change(
         fn=process_upload,
         inputs=[upload, model],
-        outputs=[document, image, img_clear_button, url_error, output, output_table],
+        outputs=submit_outputs,
     )
 
     submit.click(
         fn=process_path,
         inputs=[url, model],
-        outputs=[document, image, img_clear_button, url_error, output, output_table],
+        outputs=submit_outputs,
     )
 
     question.submit(
         fn=process_question,
-        inputs=[question, document, model],
-        outputs=[image, output, output_table],
+        inputs=[question, document, image, model, fields, output, output_table],
+        outputs=[image, fields, output, output_table],
     )
 
     submit_button.click(
@@ -452,7 +469,7 @@ with gr.Blocks(css=CSS) as demo:
     example_image.change(
         fn=load_example_document,
         inputs=[example_image, example_question, model],
-        outputs=[document, image, img_clear_button, url_error, output, output_table],
+        outputs=submit_outputs,
     )
 
 if __name__ == "__main__":