filter bboxes based on iou

app.py

@@ -2,20 +2,55 @@ import gradio as gr
 import pytesseract
 from PIL import Image, ImageDraw
 
+def calculate_iou(box1, box2):
+    xA = max(box1[0], box2[0])
+    yA = max(box1[1], box2[1])
+    xB = min(box1[0] + box1[2], box2[0] + box2[2])
+    yB = min(box1[1] + box1[3], box2[1] + box2[3])
+
+    intersection_area = max(0, xB - xA) * max(0, yB - yA)
+
+    box1_area = box1[2] * box1[3]
+    box2_area = box2[2] * box2[3]
+
+    iou = intersection_area / float(box1_area + box2_area - intersection_area)
+    return iou
+
 def extract_text_and_boxes(image):
     data = pytesseract.image_to_data(image, output_type=pytesseract.Output.DICT)
-    draw = ImageDraw.Draw(image)
-    
     boxes_and_words = []
     
     for i in range(len(data['text'])):
-        if data['text'][i].strip() != '':  # Filters out empty text results
+        if data['text'][i].strip() != '':
             x, y, w, h = data['left'][i], data['top'][i], data['width'][i], data['height'][i]
             word = data['text'][i]
             boxes_and_words.append({'box': (x, y, w, h), 'word': word})
-            draw.rectangle([x, y, x + w, y + h], outline='red', width=2)
     
-    return image, boxes_and_words
+    # Remove overlapping boxes based on IoU
+    to_remove = set()
+    for i in range(len(boxes_and_words)):
+        for j in range(i + 1, len(boxes_and_words)):
+            box1 = boxes_and_words[i]['box']
+            box2 = boxes_and_words[j]['box']
+            iou = calculate_iou(box1, box2)
+            if iou > 0.5:
+                # Remove the box with the smaller area
+                area1 = box1[2] * box1[3]
+                area2 = box2[2] * box2[3]
+                if area1 > area2:
+                    to_remove.add(j)
+                else:
+                    to_remove.add(i)
+    
+    filtered_boxes_and_words = [bw for i, bw in enumerate(boxes_and_words) if i not in to_remove]
+
+    # Draw the filtered boxes
+    draw = ImageDraw.Draw(image)
+    for bw in filtered_boxes_and_words:
+        x, y, w, h = bw['box']
+        draw.rectangle([x, y, x + w, y + h], outline='red', width=2)
+    
+    return image, filtered_boxes_and_words
 
 iface = gr.Interface(fn=extract_text_and_boxes,
                      inputs=gr.Image(type='pil'),