indexing the last dimension of numpy array

Question

I have two images called image and mask with the following shapes:

image shape: (876, 1020, 3)
mask shape: (876, 1020)

What I am trying to do is divide each of the three elements of image with the mask where the mask is non-zero. So, I do:

import numpy as np

index = np.nonzero(mask)
image[index, :] = image[index, :]/mask[index]

This however causes the program to take a long time and then it fails with;

ValueError: operands could not be broadcast together with shapes (2,302793,1020,3) (302793,)

Answer 1

If you don't have memory problems with the creation of a copy of the mask, you could set the zero elements to 1 and divide through directly:

mask2 = mask.copy()
mask2[mask2==0]=1
image /= mask2[...,np.newaxis]