If what would be a 4-connectivity in 2D is enough, you can get the neighbouring pixels that are also foreground in n log n time using a nearest neighbour tree.
Then it is matter of constructing the graph and finding the connected components (also n log n, IIRC).
#!/usr/bin/env python
"""
https://stackoverflow.com/questions/66724201/connected-component-labling-for-arrays-quasi-images-with-many-dimension
"""
import numpy as np
import networkx as nx
from scipy.spatial import cKDTree
def get_components(boolean_array):
# find neighbours
coordinates = list(zip(*np.where(boolean_array)))
tree = cKDTree(coordinates)
neighbours_by_pixel = tree.query_ball_tree(tree, r=1, p=1) # p=1 -> Manhatten distance; r=1 -> what would be 4-connectivity in 2D
# create graph and find components
G = nx.Graph()
for ii, neighbours in enumerate(neighbours_by_pixel):
if len(neighbours) > 1:
G.add_edges_from([(ii, jj) for jj in neighbours[1:]]) # skip first neighbour as that is a self-loop
components = nx.connected_components(G)
# create output image
output = np.zeros_like(data, dtype=np.int)
for ii, component in enumerate(components):
for idx in component:
output[coordinates[idx]] = ii+1
return output
if __name__ == '__main__':
shape = (5, 2, 3, 6, 10)
D = len(shape)
data = np.random.rand(*shape) < 0.1
output = get_components(data)
For an array with shape (50, 50, 50, 50), I get the following timings on my laptop:
In [48]: %timeit output = get_components(data)
5.85 s ± 279 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
