Numpy filtering using array

Boolean indexing like obs[obs != [0, 0, 0]] return a 1D array with all the elements from obs that satisfy the given condition. Look at the follwoing example:

obs = np.array([
 [[88, 0,99],
  [ 0, 0, 0]],
 [[ 0, 0, 0],
  [88,77,66]]
])

obs != [0, 0, 0] returns a boolean array:

array([[[ True, False,  True],
        [False, False, False]],
       [[False, False, False],
        [ True,  True,  True]]])

and obs[obs != [0, 0, 0]] then returns a 1D array with all the elements where the mask is True: array([88, 99, 88, 77, 66]).

So what you need is where to test if there's any color component not equal 0:

np.where(obs.any(axis=-1, keepdims=True), 255, obs)

Result:

array([[[255, 255, 255],
        [  0,   0,   0]],
       [[  0,   0,   0],
        [255, 255, 255]]])

Note that you need keepdims=True to enable broadcasting to the original shape of obs. Otherwise you'd have to add the lost dimension by np.where(obs.any(-1)[...,np.newaxis], 255, obs) or np.where(np.atleast_3d(obs.any(-1)), 255, obs) which is less elegant.

There are a number of possibilities depending what you actually want to do. Let's do the set-up code (which is common to all possibilities) first so you can see what I mean.

#!/usr/bin/env python3

import numpy as np

# Make a repeatable random image
np.random.seed(764)
obs = np.random.randint(0,32,(252,288,3), dtype=np.uint8)

This image happens to have pure black pixels at the following locations for test purposes:

obs[ 21, 267]
obs[ 28, 252]
obs[ 69, 127]
obs[ 98,   0]
obs[124, 210]
obs[133,  98]
obs[160,  81]
obs[167,  48]
obs[217, 237]

Now, suppose you want a new, pure True/False boolean mask of black pixels, you can use:

mask = obs.any(axis=-1)

That solution has the following characteristics:

time: 876 µs
mask.shape: (252,288)
mask.nbytes: 72576
mask.dtype: 'bool'

You can subsequently use and re-use that mask like this:

# Make masked pixels red
obs[mask,:] = [255,0,0]

# Make unmasked pixels cyan
obs[~mask,:] = [0,255,255]

Now let's suppose you want a new, greyscale image, with black and white pixels, you can use:

grey = obs.any(axis=-1) * np.uint8(255)

That solution has the following characteristics:

time: 887 µs
grey.shape: (252,288)
grey.nbytes: 72576
grey.dtype: np.uint8

Now suppose you want in-place alteration of your already existing "obs" to pure black and white (but still RGB):

obs[obs.any(axis=-1),:] = [255,255,255]

That solution has the following characteristics:

time: 1.98 ms
obs.shape: (252,288,3)
obs.nbytes: 217728
obs.dtype: np.uint8