Highlight the answer with a bounding box

app.py

@@ -49,12 +49,19 @@ def lift_word_boxes(document):
     return document.context["image"][0][1]
 
 
-def expand_bbox(word_boxes):
+def expand_bbox(word_boxes, padding=0.1):
     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)]
+    min_x, min_y, max_x, max_y = [min(min_x), min(min_y), max(max_x), max(max_y)]
+    if padding != 0:
+        padding = max((max_x - min_x) * padding, (max_y - min_y) * padding)
+        min_x = max(0, min_x - padding)
+        min_y = max(0, min_y - padding)
+        max_x = max_x + padding
+        max_y = max_y + padding
+    return [min_x, min_y, max_x, max_y]
 
 
 # LayoutLM boxes are normalized to 0, 1000
@@ -96,7 +103,7 @@ def process_upload(file):
         return None, None, None
 
 
-colors = ["blue", "green", "black"]
+colors = ["#64A087", "green", "black"]
 
 
 def process_question(question, document, model=list(CHECKPOINTS.keys())[0]):
@@ -105,7 +112,7 @@ def process_question(question, document, model=list(CHECKPOINTS.keys())[0]):
 
     predictions = run_pipeline(model, question, document, 3)
     image = document.preview.copy()
-    draw = ImageDraw.Draw(image)
+    draw = ImageDraw.Draw(image, "RGBA")
     for i, p in enumerate(ensure_list(predictions)):
         if i > 0:
             # Keep the code around to produce multiple boxes, but only show the top
@@ -118,7 +125,7 @@ def process_question(question, document, model=list(CHECKPOINTS.keys())[0]):
                 image.width,
                 image.height,
             )
-            draw.rectangle(((x1, y1), (x2, y2)), outline=colors[i], width=2)
+            draw.rectangle(((x1, y1), (x2, y2)), fill=(0, 255, 0, int(0.4 * 255)))
 
     return image, predictions