pythonnumpyscikit-image

compare numpy array against list of values


I would like to have your help about a problem I have with a numpy array.

In my code I have a huge array made of integers, this is the result of a skimage labelling process.

The idea is that I want to generate a boolean array with same shape as the input array containing a True value if the cell value belong to a given list of good labels, and false otherwise.

It is much easier with a small example:

import numpy as np
input_array = np.arange(15, dtype=np.int64).reshape(3, 5)
good_labels = [1, 5, 7]
mask_output = np.zeros(input_array.shape).astype(int)
for l in good_labels:
    masked_input = (input_array == l).astype(int)
    mask_output += masked_input
mask_output = mask_output > 0

This code works, but it is extremely slow because input_array is huge and good_labels is also very long, so I need to loop too much.

Is there a way to make it faster?

I was thinking to replace the whole loop with something like

mask_output = input_array in good_labels

But this does not work as expected.

Can you help me?


Solution

  • You could try numpy.isin.

    mask_output = np.isin(input_array,good_labels)