Problem applying binary mask to an RGB image with numpy

Question

I'm trying to apply a binary mask to an RGB image with numpy

I found this https://stackoverflow.com/a/26843467/4628384 but I don't have permissions to write a comment yet. Anyway I'm running into a problem; Any help much appreciated.

def masktoRGB(self,image,image_mask):

        # create mask with same dimensions as image
        mask = np.zeros_like(image)
        # copy your image_mask to all dimensions (i.e. colors) of your image
        for i in range(3):
            mask[:,:,i] = image_mask.copy()

        # apply the mask to your image

        # tried to swap axes, not a solution 
        #image = image.swapaxes(0,1)
        #this gives the error:
        masked_image = image[mask]

        print(mask.shape)
        print(image.shape)
        print(image_mask.shape)



        return masked_image

this gives me:

IndexError: index 213 is out of bounds for axis 0 with size 212

print output:

(188, 212, 3) (188, 212, 3) (188, 212)

image and image_mask are the same image, except the first is RGB and the second is mode L

Answer 1

Try to use broadcasting and multiplication:

image * image_mask[..., None]

I assume that the type of image_mask is bool that maps to numbers 0 and 1. Therefore pairwise multiplication of image and mask will set masked values to zero.

Similar efect can be achieved with np.where() or & operator.

The issue is that shapes of image and image_mask are not compatible. Numpy will first add extra dimensions at the head of shape until both tensor have the same shape. Thus image_mask is reshaped from (188, 212) to (1,188, 212). This new shape is not compatible with the shape of image (188, 212, 3).

The trick is to reshape image_mask using fancy indexing. You can use None as index to add a dummy dimension of size 1 at the end of shape. Operation image_mask[..., None] will reshape it from (188, 212) to (188, 212,1).

Broadcast rules tells that dimensions of size 1 are expanded by repeating all values along broadcast dimension. Thus numoy will automatically reshape tensir from (188, 212,1) to (188, 212,3). Operation is very fast because no copy is created.

Now bith tensors can be multiplied producing desired result.