Connected Components for Bounding Boxes

… I recently came across a small issue at work where I began to use a transformer based model for object detection that did not have non-maximal suppression baked into the inference pipeline. running inference on a model that was not fine-tuned gave back a boatload of overlapping bounding boxes. the problem got me thinking about connected components and the union-find algorithm and I realized that I could frame collapsing boxes within the confines of graph theory…

import matplotlib.pyplot as plt
import matplotlib.patches as patches

# bounding boxes as tuples (x1, y1, x2, y2)
boxes = [
    (10, 10, 30, 30),
    (15, 15, 35, 35),
    (20, 20, 40, 40),
    (25, 25, 45, 45),

    (60, 10, 80, 30),
    (67, 17, 87, 37),

    (10, 70, 25, 85),
    (70, 70, 85, 85)
]

def plot_boxes(bboxes):
    fig, ax = plt.subplots(figsize=(8, 8))
    ax.set_xlim(0, 100)
    ax.set_ylim(0, 100)
    ax.set_aspect('equal')

    colors = ['red', 'blue', 'green', 'orange', 'purple', 'brown', 'cyan', 'magenta']

    for i, box in enumerate(bboxes):
        x1, y1, x2, y2 = box
        width, height = x2 - x1, y2 - y1
        rect = patches.Rectangle(
            (x1, y1), width, height,
            linewidth=2, edgecolor=colors[i % len(colors)], facecolor='none', alpha=0.7
            )
        ax.add_patch(rect)
        ax.text(x1, y2 + 1, f'Box {i}', color=colors[i % len(colors)], fontsize=9, fontweight='bold')

    plt.title('Diagonal Boxes (Tuples)')
    plt.grid(True, linestyle='--', alpha=0.5)
    plt.show()

plot_boxes(boxes)

def get_iou(box_a, box_b):
    # Determine the coordinates of the intersection rectangle
    x_a = max(box_a[0], box_b[0])
    y_a = max(box_a[1], box_b[1])
    x_b = min(box_a[2], box_b[2])
    y_b = min(box_a[3], box_b[3])

    # Compute the area of intersection
    inter_area = max(0, x_b - x_a) * max(0, y_b - y_a)

    # Compute the area of both the prediction and ground-truth boxes
    box_a_area = (box_a[2] - box_a[0]) * (box_a[3] - box_a[1])
    box_b_area = (box_b[2] - box_b[0]) * (box_b[3] - box_b[1])

    # Compute the intersection over union
    iou = inter_area / float(box_a_area + box_b_area - inter_area)

    # Return the IoU value
    return iou

from itertools import combinations

# IOU threshold for merging two boxes
threshold = 0.25

# adjacency list graph representation
graph = {v: set() for v in range(len(boxes))}

# create edges.  edge is if any combination overlaps
for u, v in combinations(range(len(boxes)), 2):
    if get_iou(boxes[u], boxes[v]) > threshold:
        # add to adjacency list
        graph[u].add(v)
        graph[v].add(u)

print(f"Adjacency List Graph: {graph}")

# parents for union-find
parents = [i for i in range(len(boxes))]

# find goes until a root is found
def find(v):
    while v != parents[v]:
      v = parents[v]
    return v

# union assigns one vertex's parent as the other's parent
def union(u, v):
    root_u = find(u)
    root_v = find(v)
    if root_u != root_v:
        parents[root_v] = root_u

# for every edge in graph perform union to connect components
# NOTE
# ideally this should be done in edge creation step but more explicit here
seen = set()
for u in graph:
    for v in graph[u]:
        if (u, v) not in seen and (v, u) not in seen:
            union(u, v)
            seen.add((u, v))

# for each connected component keep track of merged box as
# [min x1, min y1, max x2, max y2]
components = {}
for i in range(len(boxes)):
    root = find(i)
    box = boxes[i]
    if root not in components:
        # init as a list to allow mutation
        components[root] = list(box)
    else:
        components[root][0] = min(components[root][0], box[0])  # x1
        components[root][1] = min(components[root][1], box[1])  # y1
        components[root][2] = max(components[root][2], box[2])  # x2
        components[root][3] = max(components[root][3], box[3])  # y2

collapsed_boxes = [tuple(b) for b in components.values()]

print(f"Reduced {len(boxes)} boxes to {len(collapsed_boxes)} clusters.")
plot_boxes(collapsed_boxes)

Adjacency List Graph: {0: {1}, 1: {0, 2}, 2: {1, 3}, 3: {2}, 4: {5}, 5: {4}, 6: set(), 7: set()}
Reduced 8 boxes to 4 clusters.