2

I have a 2D array, where I label clusters using the ndimage.label() function like this:

import numpy as np
from scipy.ndimage import label

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

labeled_array, _ = label(input_array)

# Result:
# labeled_array == [[0, 1, 1, 0],
#                   [1, 1, 0, 0],
#                   [0, 0, 0, 2],
#                   [0, 0, 0, 2]]

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):

{
    1: [(0, 1), (0, 2), (1, 0), (1, 1)],  # Coordinates of the elements that have the label "1"
    2: [(2, 3), (3, 3)]  # Coordinates of the elements that have the label "2"
}

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?

Canol Gökel
  • 1,168
  • 2
  • 13
  • 29

2 Answers2

4

You can make a map of the coordinates, sort and split it:

# Get the indexes (coordinates) of the labeled (non-zero) elements
ind = np.argwhere(labeled_array)

# Get the labels corresponding to those indexes above
labels = labeled_array[tuple(ind.T)]

# Sort both arrays so that lower label numbers appear before higher label numbers. This is not for cosmetic reasons,
# but we will use sorted nature of these label indexes when we use the "diff" method in the next step.
sort = labels.argsort()
ind = ind[sort]
labels = labels[sort]

# Find the split points where a new label number starts in the ordered label numbers
splits = np.flatnonzero(np.diff(labels)) + 1

# Create a data structure out of the label numbers and indexes (coordinates).
# The first argument to the zip is: we take the 0th label number and the label numbers at the split points
# The second argument is the indexes (coordinates), split at split points
# so the length of both arguments to the zip function is the same
result = {k: v for k, v in zip(labels[np.r_[0, splits]],
                               np.split(ind, splits))}
Canol Gökel
  • 1,168
  • 2
  • 13
  • 29
Mad Physicist
  • 107,652
  • 25
  • 181
  • 264
  • I am accepting this solution because the heavy weight stuff is done in a vectorized manner and there is a very lightweight comprehension in the end. So it works really fast on big arrays with thousands of labels. Thank you! – Canol Gökel Jul 20 '21 at 09:49
  • I added comments for as much as I understand the code. A question: can we replace the np.r_ with labels[0, *splits] or even np.unique? Both of which are a bit easier to understand to me, what do you think? – Canol Gökel Jul 20 '21 at 10:06
  • @CanolGökel. np.r_ is shorthand for np.concatenate that allows scalars – Mad Physicist Jul 20 '21 at 13:47
  • I suppose you can use `labels[[0] + list(splits)]`, but I don't like converting to list just to convert right back – Mad Physicist Jul 20 '21 at 13:59
1

Method 1:

You can try this, still looping using dictionary comprehension:

{k: list(zip(*np.where(labeled_array == k))) for k in range(1,3)}

Output:

{1: [(0, 1), (0, 2), (1, 0), (1, 1)], 2: [(2, 3), (3, 3)]}

Method 2 (slow):

Here's a way using pandas probably slower that Mad Physicist's method:

(pd.DataFrame(labeled_array)
  .stack() 
  .reset_index()
  .groupby(0).agg(list)[1:]
  .apply(lambda x: list(zip(*x)), axis=1)
).to_dict()

Output:

{1: [(0, 1), (0, 2), (1, 0), (1, 1)], 2: [(2, 3), (3, 3)]}

Timings with this data:

Dictionary comprehension

8.73 µs ± 216 ns per loop (mean ± std. dev. of 7 runs, 100000 loops each)

Using map coordinates, sort and split:

57.3 µs ± 5.55 µs per loop (mean ± std. dev. of 7 runs, 10000 loops each)

pandas

5.16 ms ± 283 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)

Scott Boston
  • 147,308
  • 15
  • 139
  • 187
  • 1
    Thank you for your answer and benchmarks. I am currently trying out both your and @Mad Physicist answers. For big arrays with thousands of labels, answer of Mad Physicist's seems to be working significantly faster, but I haven't finished yet checking out comprehensively. – Canol Gökel Jul 20 '21 at 07:46