remove background of RGB image (3D array) based on another array (boolean mask)

I’m stuck on what should be fairly straight forward, but other opened questions don’t see to address exactly the same issue i’m having.

I’m trying to crop an image based on boolean mask. I can do this with:

# crop image based on segmentation mask
def get_segment_crop(img,tol=0, mask=None):
    for i in range(mask.shape[2]):  
        crops = []
        if mask is None:
            mask = img > tol
        for i in range(mask.shape[2]):
            crops.append(img[np.ix_(mask[:,:,i].any(1), mask[:,:,i].any(0))])
        return crops

(This function inputs RGB image (w, h, c) and True/False mask of shape (w, h, n) and crops a bounding box around each segmentation mask). But then I want to try to remove the background (ie. set everything that isn’t True to 0) around each segmentation mask of the cropped outputs (the mask segmentations are shapes — not perfect rectangles). I thought that removing the background first then cropping could work, but i’m stuck on the removing the background part because my image has 3 channels.

Here’s a working example and what i’ve tried so far with the issues i’m facing embedded as comments.

import numpy as np
from scipy import misc
import matplotlib.pyplot as plt

file = misc.face()
img = np.array(file)

#create random masks
tf_mask = np.full((img.shape[0], img.shape[1], 2), False)
tf_mask[500:600,500:600,0] = True
tf_mask[500:650, 550:700, 0] = True
tf_mask[100:200, 100:200, 1] = True
tf_mask[200:300, 100:150, 1] = True
# np.where(tf_mask[:,:,0])[0] #sanity check True values were inputed

# Trying to set all values in img where masks=False to 0. Problem: it saves over original img array so the next mask doesn't display region of interest from original image
for i in range(tf_mask.shape[2]):
    img[tf_mask[:,:,i]==False] = 0
    plt.imshow(img)
    plt.show()

# issue with this: image is 3 channels so need to iterate over all channels but how to then merge back?
for i in range(tf_mask.shape[2]):
    for j in range(3):
        new_img = np.where(tf_mask[:,:, i]==False, 0, img[:,:,j])
        plt.imshow(new_img)
        plt.show()

Thanks for ideas to get around this or if there’s a neater way to do this.

Answer

If you have multiple masks in a single 3d array where the first two dimensions are the same as the image and the individual masks are stacked in the third dimension, first collapse the mask(s) along the third dimension then multiply it with the image.

Assuming you want to keep a pixel if it is True in any the masks:

mask = np.any(mask,axis=-1)
masked = mask[:,:,None] * img

---

import numpy as np
rng = np.random.default_rng()
a = np.ones((3,3,3))
b = rng.integers(0,1,(3,3,2),endpoint=True,dtype=bool)

>>> a
array([[[1., 1., 1.],
        [1., 1., 1.],
        [1., 1., 1.]],

       [[1., 1., 1.],
        [1., 1., 1.],
        [1., 1., 1.]],

       [[1., 1., 1.],
        [1., 1., 1.],
        [1., 1., 1.]]])
>>> b
array([[[ True, False],
        [False,  True],
        [False, False]],

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

       [[False,  True],
        [ True,  True],
        [ True,  True]]])
>>> b[0,0,:]
array([ True, False])
>>> b[0,1,:]
array([False,  True])
>>> b[0,2,:]
array([False, False])
>>> b = np.any(b,axis=-1)
>>> b[0,0]
True
>>> b[0,1]
True
>>> b[0,2]
False
>>> a * b[:,:,None]
array([[[1., 1., 1.],
        [1., 1., 1.],
        [0., 0., 0.]],

       [[1., 1., 1.],
        [0., 0., 0.],
        [1., 1., 1.]],

       [[1., 1., 1.],
        [1., 1., 1.],
        [1., 1., 1.]]])