Search code examples
pythonimageperformancenumpyvectorization

Python mask image pixels from a set of values


Provided an image with labels (the value of a pixel corresponds to its label), and the list of labels that are accepted, I am trying to create a "mask" image with 255 value if the pixels label is accepted, 0 otherwise.

I know that this is a slow approach as it iterates over the image at python-speed (but it demonstrates the idea well):

mask = numpy.zeros(labels.shape[:2], dtype = "uint8")

for i in xrange(mask.shape[0]):
    for j in xrange(mask.shape[1]):
        if labels[i][j] in accepted:
            mask[i][j] = 255

I know that it is much faster to use python slicing and masking, but I do not know how to compose a complicated condition. I still get a tremendous speed-up when I mask the pixels one-by-one accepted label, like so:

for value in accepted:
    mask[labels == value] = 255

Can I somehow make a one-liner doing what I want? My python knowledge is rusty (read: almost no python in the last few years), so while I tried composing this using some examples I found, this is the closest I got:

mask[(labels in accepted).all()] = 255

And in this case I get the following error: ValueError: The truth value of an array with more than one element is ambiguous. Use a.any() or a.all()

I've looked at similar SO questions (e.g.here or here and more) but they all seem to cover the cases where either the values are from a range or lower/higher than a threshold (<10) , or where the image slice to change is continuous.

Any suggestion on how check "is value among accepted values" would be great.


Solution

  • In the meantime, I found an acceptable-speed solution to my own question:

    mask = numpy.zeros(labels.shape[:2], dtype = "uint8")
    mask[numpy.in1d(labels, accepted).reshape(mask.shape)] = 255
    

    It consists in first using numpy.in1d to get a boolean array from the labels array, and check which ones are present in accepted (element-wise function of the python keyword "in").

    As this apparently necessarily returns a 1D array, even if it can be applied to a 2D array (it simply unravels the array), so I follow by using reshape() to make the boolean array dimensions correspond to that of the mask.

    Finally, I use this boolean array to index the required elements of mask and set them to a desired value.