Getting the connected components from a 3D segmentation mask with SimpleITK

components.py

import SimpleITK as sitk

def get_components(mask):
    f = sitk.ConnectedComponentImageFilter()
    labeled = f.Execute(mask)
    count = f.GetObjectCount()
    return count, labeled

Get Connected Components of 3D Mask with SimpleITK.md

The SimpleITK library has an easy to use class for computing connected components in a segmentation mask: sitk.ConnectedComponentImageFilter.

Here is a sample usage (here the “mask” argument would be a binary mask (0 = background, 1 = foreground)

def get_components(mask):
    f = sitk.ConnectedComponentImageFilter()
    labeled = f.Execute(mask)
    count = f.GetObjectCount()
    return count, labeled

That function returns the number of connected components and a mask image with the components separated (i.e. background is 0, first component is 1, second component is 2, etc, up to count components). You can then get the mask for each one like

for i in range(1, count+1):
    mask_i = labeled==i

sample_usage.py

from components import get_components

def yield_all_components(mask):
    """Yield all connected components in mask, sequentially."""
    count, labeled = get_components(mask)
    
    for i in range(1, i+1):
        yield labeled == i