I have a 2D array, where I label clusters using the ndimage.label()
function like this:
JavaScript
x
16
16
1
import numpy as np
2
from scipy.ndimage import label
3
4
input_array = np.array([[0, 1, 1, 0],
5
[1, 1, 0, 0],
6
[0, 0, 0, 1],
7
[0, 0, 0, 1]])
8
9
labeled_array, _ = label(input_array)
10
11
# Result:
12
# labeled_array == [[0, 1, 1, 0],
13
# [1, 1, 0, 0],
14
# [0, 0, 0, 2],
15
# [0, 0, 0, 2]]
16
I can get the element counts, the centroids or the bounding box of the labeled clusters. But I would like to also get the coordinates of each element in clusters. Something like this (the data structure doesn’t have to be like this, any data structure is okay):
JavaScript
1
5
1
{
2
1: [(0, 1), (0, 2), (1, 0), (1, 1)], # Coordinates of the elements that have the label "1"
3
2: [(2, 3), (3, 3)] # Coordinates of the elements that have the label "2"
4
}
5
I can loop over the label list and call np.where()
for each one of them but I wonder if there is a way to do this without a loop, so that it would be faster?
Advertisement
Answer
You can make a map of the coordinates, sort and split it:
JavaScript
1
22
22
1
# Get the indexes (coordinates) of the labeled (non-zero) elements
2
ind = np.argwhere(labeled_array)
3
4
# Get the labels corresponding to those indexes above
5
labels = labeled_array[tuple(ind.T)]
6
7
# Sort both arrays so that lower label numbers appear before higher label numbers. This is not for cosmetic reasons,
8
# but we will use sorted nature of these label indexes when we use the "diff" method in the next step.
9
sort = labels.argsort()
10
ind = ind[sort]
11
labels = labels[sort]
12
13
# Find the split points where a new label number starts in the ordered label numbers
14
splits = np.flatnonzero(np.diff(labels)) + 1
15
16
# Create a data structure out of the label numbers and indexes (coordinates).
17
# The first argument to the zip is: we take the 0th label number and the label numbers at the split points
18
# The second argument is the indexes (coordinates), split at split points
19
# so the length of both arguments to the zip function is the same
20
result = {k: v for k, v in zip(labels[np.r_[0, splits]],
21
np.split(ind, splits))}
22