Charlie Li commited on
Commit
b8048df
1 Parent(s): fc269ec

filter bboxes based on iou

Browse files
Files changed (1) hide show
  1. app.py +40 -5
app.py CHANGED
@@ -2,20 +2,55 @@ import gradio as gr
2
  import pytesseract
3
  from PIL import Image, ImageDraw
4
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
5
  def extract_text_and_boxes(image):
6
  data = pytesseract.image_to_data(image, output_type=pytesseract.Output.DICT)
7
- draw = ImageDraw.Draw(image)
8
-
9
  boxes_and_words = []
10
 
11
  for i in range(len(data['text'])):
12
- if data['text'][i].strip() != '': # Filters out empty text results
13
  x, y, w, h = data['left'][i], data['top'][i], data['width'][i], data['height'][i]
14
  word = data['text'][i]
15
  boxes_and_words.append({'box': (x, y, w, h), 'word': word})
16
- draw.rectangle([x, y, x + w, y + h], outline='red', width=2)
17
 
18
- return image, boxes_and_words
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
19
 
20
  iface = gr.Interface(fn=extract_text_and_boxes,
21
  inputs=gr.Image(type='pil'),
 
2
  import pytesseract
3
  from PIL import Image, ImageDraw
4
 
5
+ def calculate_iou(box1, box2):
6
+ xA = max(box1[0], box2[0])
7
+ yA = max(box1[1], box2[1])
8
+ xB = min(box1[0] + box1[2], box2[0] + box2[2])
9
+ yB = min(box1[1] + box1[3], box2[1] + box2[3])
10
+
11
+ intersection_area = max(0, xB - xA) * max(0, yB - yA)
12
+
13
+ box1_area = box1[2] * box1[3]
14
+ box2_area = box2[2] * box2[3]
15
+
16
+ iou = intersection_area / float(box1_area + box2_area - intersection_area)
17
+ return iou
18
+
19
  def extract_text_and_boxes(image):
20
  data = pytesseract.image_to_data(image, output_type=pytesseract.Output.DICT)
 
 
21
  boxes_and_words = []
22
 
23
  for i in range(len(data['text'])):
24
+ if data['text'][i].strip() != '':
25
  x, y, w, h = data['left'][i], data['top'][i], data['width'][i], data['height'][i]
26
  word = data['text'][i]
27
  boxes_and_words.append({'box': (x, y, w, h), 'word': word})
 
28
 
29
+ # Remove overlapping boxes based on IoU
30
+ to_remove = set()
31
+ for i in range(len(boxes_and_words)):
32
+ for j in range(i + 1, len(boxes_and_words)):
33
+ box1 = boxes_and_words[i]['box']
34
+ box2 = boxes_and_words[j]['box']
35
+ iou = calculate_iou(box1, box2)
36
+ if iou > 0.5:
37
+ # Remove the box with the smaller area
38
+ area1 = box1[2] * box1[3]
39
+ area2 = box2[2] * box2[3]
40
+ if area1 > area2:
41
+ to_remove.add(j)
42
+ else:
43
+ to_remove.add(i)
44
+
45
+ filtered_boxes_and_words = [bw for i, bw in enumerate(boxes_and_words) if i not in to_remove]
46
+
47
+ # Draw the filtered boxes
48
+ draw = ImageDraw.Draw(image)
49
+ for bw in filtered_boxes_and_words:
50
+ x, y, w, h = bw['box']
51
+ draw.rectangle([x, y, x + w, y + h], outline='red', width=2)
52
+
53
+ return image, filtered_boxes_and_words
54
 
55
  iface = gr.Interface(fn=extract_text_and_boxes,
56
  inputs=gr.Image(type='pil'),