Draw a box over the answer

app.py

@@ -2,6 +2,7 @@ import os
 
 os.environ["TOKENIZERS_PARALLELISM"] = "false"
 
+from PIL import ImageDraw
 import streamlit as st
 
 import torch
@@ -24,8 +25,28 @@ def construct_pipeline():
 
 
 @st.cache
-def run_pipeline(question, document):
-    return construct_pipeline()(question=question, **document.context)
+def run_pipeline(question, document, top_k):
+    return construct_pipeline()(question=question, **document.context, top_k=top_k)
+
+
+# TODO: Move into docquery
+# TODO: Support words past the first page (or window?)
+def lift_word_boxes(document):
+    return document.context["image"][0][1]
+
+
+def expand_bbox(word_boxes):
+    if len(word_boxes) == 0:
+        return None
+
+    min_x, min_y, max_x, max_y = zip(*[x[1] for x in word_boxes])
+    return [min(min_x), min(min_y), max(max_x), max(max_y)]
+
+
+# LayoutLM boxes are normalized to 0, 1000
+def normalize_bbox(box, width, height):
+    pct = [c / 1000 for c in box]
+    return [pct[0] * width, pct[1] * height, pct[2] * width, pct[3] * height]
 
 
 st.markdown("# DocQuery: Query Documents w/ NLP")
@@ -75,16 +96,30 @@ question = st.text_input("QUESTION", "")
 document = st.session_state.document
 loading_placeholder = st.empty()
 if document is not None:
-    col1, col2 = st.columns(2)
-    col1.image(document.preview, use_column_width=True)
+    col1, col2 = st.columns([3, 1])
+    image = document.preview
 
-if document is not None and question is not None and len(question) > 0:
-    predictions = run_pipeline(question=question, document=document)
 
+colors = ["blue", "red", "green"]
+if document is not None and question is not None and len(question) > 0:
     col2.header("Answers")
-    for p in ensure_list(predictions):
-        col2.subheader(f"{ p['answer'] }: ({round(p['score'] * 100, 1)}%)")
 
+    predictions = run_pipeline(question=question, document=document, top_k=1)
+
+    word_boxes = lift_word_boxes(document)
+    image = image.copy()
+    draw = ImageDraw.Draw(image)
+    for i, p in enumerate(ensure_list(predictions)):
+        col2.markdown(f"#### { p['answer'] }: ({round(p['score'] * 100, 1)}%)")
+        x1, y1, x2, y2 = normalize_bbox(
+            expand_bbox(word_boxes[p["start"] : p["end"] + 1]),
+            image.width,
+            image.height,
+        )
+        draw.rectangle(((x1, y1), (x2, y2)), outline=colors[i])
+
+if document is not None:
+    col1.image(image, use_column_width=True)
 
 "DocQuery uses LayoutLMv1 fine-tuned on DocVQA, a document visual question answering dataset, as well as SQuAD, which boosts its English-language comprehension. To use it, simply upload an image or PDF, type a question, and click 'submit', or click one of the examples to load them."