Python equivalent of label2idx MATLAB function

Given an array of labels label_arr containing all labels from 1 to max(label_arr), you can do:

def label2idx(label_arr):
    return [
        np.where(label_arr.ravel() == i)[0]
        for i in range(1, np.max(label_arr) + 1)]

If you want to relax the requirement of all labels being contained you can add a simple if, i.e.:

def label2idx(label_arr):
    return [
        np.where(label_arr.ravel() == i)[0]
            if i in label_arr else np.array([], dtype=int)
        for i in range(1, np.max(label_arr) + 1)]

Just to replicate the example in the MATLAB docs:

import numpy as np
import scipy as sp
import scipy.ndimage

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

label_arr, num_labels = sp.ndimage.label(struct_arr)
# label_arr:
# [[1 1 1 0 0 0 0 0]
#  [1 1 1 0 2 2 0 0]
#  [1 1 1 0 2 2 0 0]
#  [1 1 1 0 0 0 0 0]
#  [1 1 1 0 0 0 3 0]
#  [1 1 1 0 0 0 3 0]
#  [1 1 1 0 0 3 3 0]
#  [1 1 1 0 0 0 0 0]]

def label2idx(label_arr):
    return [
        np.where(label_arr.ravel() == i)[0]
        for i in range(1, np.max(label_arr) + 1)]

pixel_idxs = label2idx(label_arr)

for pixel_idx in pixel_idxs:
    print(pixel_idx)

# [ 0  1  2  8  9 10 16 17 18 24 25 26 32 33 34 40 41 42 48 49 50 56 57 58]
# [12 13 20 21]
# [38 46 53 54]

Note, however that you do not get the very same results because of the differences between MATLAB and NumPy, notably:

• MATLAB: FORTRAN-style matrix indexing and 1-based indexing
• Python+NumPy: C-style matrix indexing and 0-based indexing

and if you want to get the very same numbers you get in MATLAB you may use this instead (note the extra .T and the + 1):

def label2idx_MATLAB(label_arr):
    return [
        np.where(label_arr.T.ravel() == i)[0] + 1
        for i in range(1, np.max(label_arr) + 1)]

MATLAB's label2idx outputs the flattened indices (column-major ordered) given the labeled image.

We can use scikit-image's built-in regionprops to get those indices from the labeled image. Scikit-image also provides for us a built-in to get the labeled image, so it all works out with that same package. The implementation would look something like this -

from skimage.measure import label,regionprops

def label2idx(L):
    # Get region-properties for all labels
    props = regionprops(L)

    # Get XY coordinates/indices for each label
    indices = [p.coords for p in props]

    # Get flattened-indices for each label, similar to MATLAB version
    # Note that this is row-major ordered.
    flattened_indices = [np.ravel_multi_index(idx.T,L.shape) for idx in indices]
    return indices, flattened_indices

Sample run -

# Input array
In [62]: a
Out[62]: 
array([[1, 1, 1, 0, 0, 0, 0, 0],
       [1, 1, 1, 0, 1, 1, 0, 0],
       [1, 1, 1, 0, 1, 1, 0, 0],
       [1, 1, 1, 0, 0, 0, 0, 0],
       [1, 1, 1, 0, 0, 0, 1, 0],
       [1, 1, 1, 0, 0, 0, 1, 0],
       [1, 1, 1, 0, 0, 1, 1, 0],
       [1, 1, 1, 0, 0, 0, 0, 0]])

# Get labeled image
In [63]: L = label(a)

In [64]: idx,flat_idx = label2idx(L)

In [65]: flat_idx
Out[65]: 
[array([ 0,  1,  2,  8,  9, 10, 16, 17, 18, 24, 25, 26, 32, 33, 34, 40, 41,
        42, 48, 49, 50, 56, 57, 58]),
 array([12, 13, 20, 21]),
 array([38, 46, 53, 54])]

If you need the indices in column-major order like in MATLAB, simply transpose the image and then feed in -

In [5]: idx,flat_idx = label2idx(L.T)

In [6]: flat_idx
Out[6]: 
[array([ 0,  1,  2,  3,  4,  5,  6,  7,  8,  9, 10, 11, 12, 13, 14, 15, 16,
        17, 18, 19, 20, 21, 22, 23]),
 array([33, 34, 41, 42]),
 array([46, 52, 53, 54])]

Note that the indexing still starts from 0, unlike in MATLAB where it's from 1.

Alternative to getting labeled image with SciPy

SciPy also has a built-in to get labeled-image : scipy.ndimage.label -

from scipy.ndimage import label

L = label(a)[0]