Python: get boundingbox coordinates for each cluster in segmentation map (2D numpy array)

Question

Python:

I got a segmentation map (2D numpy array) with class values (integer 0 to N) for each pixel of the original img and i want to find the bounding box coordinates for each connected cluster in the segmentation map.

EDIT: there might be more than one cluster per class in the map!

I guess i can use something like skimage.measure.label(seg_map, connectivity=1)

Answer 1

Complete answer

There's an existing function scipy.ndimage.measurements.find_objects that allllmost does exactly what you want. Needs a little help from Numpy and scipy.ndimage.measurements.label, though:

import numpy as np
import scipy.ndimage.measurements as mnts

A = np.array([
    [0, 0, 0, 0, 0, 0, 0],
    [0, 1, 1, 0, 2, 2, 0],
    [0, 1, 1, 0, 2, 2, 0],
    [0, 0, 0, 0, 0, 0, 0],
    [0, 4, 4, 0, 1, 1, 0],
    [0, 4, 4, 0, 1, 1, 0],
    [4, 0, 0, 0, 0, 0, 1]
])

structure = np.array([
    [1,1,1],
    [1,1,1],
    [1,1,1]
])

bboxSlices = {}
for i in range(1, A.max() + 1):
    B = A.copy()
    B[B != i] = 0

    bboxSlices[i] = mnts.find_objects(mnts.label(B, structure=structure)[0])

print(bboxSlices)

output:

{1: [(slice(1, 3, None), slice(1, 3, None)),
     (slice(4, 7, None), slice(4, 7, None))],
 2: [(slice(1, 3, None), slice(4, 6, None))],
 3: [],
 4: [(slice(4, 7, None), slice(0, 3, None))]}

Each entry in the bboxSlices dict is a list of tuples. Each tuple contains two slices, a row slice and a column slice, that each define a bounding box around a cluster of the corresponding class.

details

label(...) finds clusters of features and replaces their values with a label (eg 1 for the 1st cluster, 2 for the 2nd, etc). find_objects(...) then finds the bounding boxes around each label. The problem is that label treats all non-zero values as "features". So for each class value i we need a copy of A with all non-i values zeroed-out.

structure defines the connectivity of the clusters. If you wanted clusters that were not connected along diagonals, you would use a different structure:

structure = np.array([
    [0,1,0],
    [1,1,1],
    [0,1,0]
])

Simple answer

It's easy if there's only 1 cluster per class:

import numpy as np

A = np.array([
    [0, 0, 0, 0, 0, 0, 0],
    [0, 1, 1, 0, 2, 2, 0],
    [0, 1, 1, 0, 2, 2, 0],
    [0, 0, 0, 0, 0, 0, 0],
    [0, 4, 4, 0, 3, 3, 0],
    [0, 4, 4, 0, 3, 3, 0],
    [0, 0, 0, 0, 0, 0, 0]
])

bboxCorners = {}
for i in range(1, A.max()+1):
    B = np.argwhere(A==i)
    bboxCorners[i] = B.min(0), B.max(0)

print(bboxCorners)

output:

{1: (array([1, 1]), array([2, 2])),
 2: (array([1, 4]), array([2, 5])),
 3: (array([4, 4]), array([5, 5])),
 4: (array([4, 1]), array([5, 2]))}