Support multi-page preview

app.py

@@ -45,8 +45,8 @@ def run_pipeline(model, question, document, 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 lift_word_boxes(document, page):
+    return document.context["image"][page][1]
 
 
 def expand_bbox(word_boxes):
@@ -122,23 +122,26 @@ def process_question(question, document, model=list(CHECKPOINTS.keys())[0]):
         return None, None
 
     predictions = run_pipeline(model, question, document, 3)
-    image = document.preview.copy()
-    draw = ImageDraw.Draw(image, "RGBA")
+    pages = [x.copy() for x in document.preview]
     for i, p in enumerate(ensure_list(predictions)):
         if i > 0:
             # Keep the code around to produce multiple boxes, but only show the top
             # prediction for now
             break
 
-        if "start" in p and "end" in p and p.get("page") == 0:
+        if "start" in p and "end" in p:
+            image = pages[p["page"]]
+            draw = ImageDraw.Draw(image, "RGBA")
             x1, y1, x2, y2 = normalize_bbox(
-                expand_bbox(lift_word_boxes(document)[p["start"] : p["end"] + 1]),
+                expand_bbox(
+                    lift_word_boxes(document, p["page"])[p["start"] : p["end"] + 1]
+                ),
                 image.width,
                 image.height,
             )
             draw.rectangle(((x1, y1), (x2, y2)), fill=(0, 255, 0, int(0.4 * 255)))
 
-    return gr.update(visible=True, value=image), gr.update(
+    return gr.update(visible=True, value=pages), gr.update(
         visible=True, value=predictions
     )
 
@@ -197,7 +200,7 @@ with gr.Blocks(css=CSS) as demo:
         submit_button = gr.Button("Submit", variant="primary", elem_id="submit-button")
 
     with gr.Row():
-        image = gr.Image(visible=False)
+        image = gr.Gallery(visible=False)
         with gr.Column():
             output = gr.JSON(label="Output", visible=False)
 

requirements.txt

@@ -1,4 +1,4 @@
 torch
 git+https://github.com/huggingface/transformers.git@21f6f58721dd9154357576be6de54eefef1f1818
-git+https://github.com/impira/docquery.git@5de147aa7df285829a8d064f3ca2c7ac74efe0f8
+git+https://github.com/impira/docquery.git@3aa3cc6ca6624d6371db9cd4732cbbcb9c8b3ea0
 sentencepiece