Ankur Goyal commited on
Commit
0b2b653
1 Parent(s): 2b1c83d

Highlight the answer with a bounding box

Browse files
Files changed (1) hide show
  1. app.py +12 -5
app.py CHANGED
@@ -49,12 +49,19 @@ def lift_word_boxes(document):
49
  return document.context["image"][0][1]
50
 
51
 
52
- def expand_bbox(word_boxes):
53
  if len(word_boxes) == 0:
54
  return None
55
 
56
  min_x, min_y, max_x, max_y = zip(*[x[1] for x in word_boxes])
57
- return [min(min_x), min(min_y), max(max_x), max(max_y)]
 
 
 
 
 
 
 
58
 
59
 
60
  # LayoutLM boxes are normalized to 0, 1000
@@ -96,7 +103,7 @@ def process_upload(file):
96
  return None, None, None
97
 
98
 
99
- colors = ["blue", "green", "black"]
100
 
101
 
102
  def process_question(question, document, model=list(CHECKPOINTS.keys())[0]):
@@ -105,7 +112,7 @@ def process_question(question, document, model=list(CHECKPOINTS.keys())[0]):
105
 
106
  predictions = run_pipeline(model, question, document, 3)
107
  image = document.preview.copy()
108
- draw = ImageDraw.Draw(image)
109
  for i, p in enumerate(ensure_list(predictions)):
110
  if i > 0:
111
  # 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]):
118
  image.width,
119
  image.height,
120
  )
121
- draw.rectangle(((x1, y1), (x2, y2)), outline=colors[i], width=2)
122
 
123
  return image, predictions
124
 
 
49
  return document.context["image"][0][1]
50
 
51
 
52
+ def expand_bbox(word_boxes, padding=0.1):
53
  if len(word_boxes) == 0:
54
  return None
55
 
56
  min_x, min_y, max_x, max_y = zip(*[x[1] for x in word_boxes])
57
+ min_x, min_y, max_x, max_y = [min(min_x), min(min_y), max(max_x), max(max_y)]
58
+ if padding != 0:
59
+ padding = max((max_x - min_x) * padding, (max_y - min_y) * padding)
60
+ min_x = max(0, min_x - padding)
61
+ min_y = max(0, min_y - padding)
62
+ max_x = max_x + padding
63
+ max_y = max_y + padding
64
+ return [min_x, min_y, max_x, max_y]
65
 
66
 
67
  # LayoutLM boxes are normalized to 0, 1000
 
103
  return None, None, None
104
 
105
 
106
+ colors = ["#64A087", "green", "black"]
107
 
108
 
109
  def process_question(question, document, model=list(CHECKPOINTS.keys())[0]):
 
112
 
113
  predictions = run_pipeline(model, question, document, 3)
114
  image = document.preview.copy()
115
+ draw = ImageDraw.Draw(image, "RGBA")
116
  for i, p in enumerate(ensure_list(predictions)):
117
  if i > 0:
118
  # Keep the code around to produce multiple boxes, but only show the top
 
125
  image.width,
126
  image.height,
127
  )
128
+ draw.rectangle(((x1, y1), (x2, y2)), fill=(0, 255, 0, int(0.4 * 255)))
129
 
130
  return image, predictions
131