通过2D数组索引3D数组

Question

I have a 3D color image im (shape 512 512 3), and a 2D array mask(512 512). I want to annotate this color image by the mask:

im = im[mask>threshold] + im[mask<threshold] * 0.2 + (255,0,0) * [mask<threshold].

How do I write this in Python efficiently?

Answer 1

This works:

mask3 = numpy.dstack(mask,mask,mask)
im = im * (mask3>threshold) + im * (mask3<threshold) * 0.2
im[:,:,0] += 255 * (mask<threshold)

It relies on the fact that the numeric value of true is 1 and false is 0.

It may not be the clearest or the most efficient, but it will still likely be much faster than indexing by a boolean array, like im[ mask3 < threshold ] *= 0.2 ( unless the index has a very small number of true values, anyway ).