Mask of an image with a list of pixel values

Question

I want to create a mask of an image with the values in a list. For example I have an RGB image with dimension (2, 5):

a = (np.random.rand(2, 5, 3) * 10).astype(int)

array([[[0, 5, 8],
        [9, 0, 2],
        [2, 2, 9],
        [9, 2, 4],
        [2, 5, 3]],

       [[7, 5, 7],
        [1, 9, 3],
        [4, 3, 3],
        [9, 1, 1],
        [9, 5, 5]]]

b = np.array([[0, 5, 8], [7, 5, 7], [4, 3, 3]])

array([[[0, 5, 8],
        [7, 5, 7],
        [4, 3, 3]]])

What I want to do is I want to create a mask of image a retaining the pixel values (each pixel value is in a list of pixel value in list b) across 3 RGB channel (third dimension). Example result for given a and b is:

array([[[True, True, True],          # check if [0, 5, 8] is in b
        [False, False, False],       # check if [9, 0, 2] is in b
        [False, False, False],       # check if [2, 2, 9] is in b
        [False, False, False],       # check if [9, 2, 4] is in b
        [False, False, False]],      # check if [2, 5, 3] is in b

       [[True, True, True],
        [False, False, False],
        [True, True, True],
        [False, False, False],
        [False, False, False]]]

or the output can be this because we are masking an image with pixel values so the third dimension can be omitted

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

       [True,
        False,
        True,
        False,
        False]]

I've tried these but they are too slow for my need:

# A naive way using list comprehension. Slowest I've tried
mask = [[True if np.any(b == a[r, c, :]) else False
          for c in range(a.shape[1])] for r in range(a.shape[0])]

# I've tried this but it's probably a wrong solution since in1d flatten the array to compare
#them. I need the pixel to be compared across the 3rd dimension
mask = np.in1d(a[:, :,  (?)], b, invert=True).reshape(image.shape[:2])

# Fastest I've tried. Took around 4 seconds on my machine on an image with dimension
# (3000, 3000, 3). But apparently it's not fast enough
mask = np.zeros(image.shape[:2], dtype=np.bool)
for val in b:
    mask |= (a == b).all(-1)

Any help is appreciated:D

Answer 1

Try this solution using broadcasting.

aa = a[:, :, None, :]
bb = b[None, None]
mask = (aa == bb).any(axis=2).all(axis=-1)

We get:

In [57]: mask
Out[57]:
array([[ True, False, False, False, False],
       [ True, False,  True, False, False]])

Explanation

Broadcasting is useful here since we want to replicate information across dimensions without looping. We first take your image in a and introduce a singleton dimension in the axis=2 dimension, then move the channels over to the axis=3 dimension. This makes our array four dimensions. Similarly for b , we make the array four dimensions by introducing singleton dimensions for the first and second dimensions. Once you do this, doing aa == bb will do an element-wise equality check over every RGB pixel in a with every RGB pixel in b . The any computation with axis=2 will check to see if any pixel in a matches all elements in b . If this happens, every single RGB channel for this pixel should equal True . To finally complete the check, we should output True if every single RGB pixel is True , so we finish this off with axis=-1 combined with all . This should be fast for a large image since you are only changing the views. If not, it should be at least faster than the for loop approaches you have created.

Answer 2

You could maybe try taking the dot product of each RGB pixel with [1,256,65536] to "flatten" each pixel to a single 24-bit integer then you can use np.in1d() :

# Get some deterministic randomness ;-)
np.random.seed(42)

# Synthesize (hopefully) representative image and 16 colours
a = np.random.randint(0,256,(3000, 3000, 3), np.uint8)
b = np.random.randint(0,256,(16, 3), np.uint8)

# "Flatten" third dimension to single 24-bit number
a24 = np.dot(a.astype(np.uint32),[1,256,65536])
b24 = np.dot(b.astype(np.uint32),[1,256,65536])

# Now use np.in1d()
mask = np.in1d(a24,b24)

On my Mac, the timings are as follows:

%timeit a24 = np.dot(a.astype(np.uint32),[1,256,65536])
153 ms ± 1.43 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

%timeit mask = np.in1d(a24,b24)
93.4 ms ± 400 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)

Answer 3

If I'm correct you want to check if any of the pixel values of a are in b. As a quick solution, you can do it with a simple trick.

[[str(j) in str(b) for j in i] for i in a]


[[True, False, False, False, False], [True, False, True, False, False]]