CODEWITHSUNDEEP
pythonarraysnumpy
GaryO
GaryO

Reputation: 6338

Set masked pixels in a 3D RGB numpy array

I'd like to set all pixels matching some condition in a 3d numpy array (RGB image) using a mask. I have something like this:

def make_dot(img, color, radius):
    """Make a dot of given color in the center of img (rgb numpy array)"""
    (ydim,xdim,dummy) = img.shape
    # make an open grid of x,y
    y,x = np.ogrid[0:ydim, 0:xdim, ]
    y -= ydim/2                 # centered at the origin
    x -= xdim/2
    # now make a mask
    mask = x**2+y**2 <= radius**2 # start with 2d
    mask.shape = mask.shape + (1,) # make it 3d
    print img[mask].shape
    img[mask] = color

img = np.zeros((100, 200, 3))
make_dot(img, np.array((.1, .2, .3)), 25)

but that gives ValueError: array is not broadcastable to correct shape in this line:

img[mask] = color

because the shape of img[mask] is (1961,); i.e. it's flattened to contain only the "valid" pixels, which makes sense; but how can I make it "write through the mask" as it were to set only the pixels where the mask is 1? Note that I want to write three values at once to each pixel (the last dim).

Upvotes: 4

Views: 3325

Answers (1)

kith
kith

Reputation: 5566

You almost have it right. 

(ydim,xdim,dummy) = img.shape
# make an open grid of x,y
y,x = np.ogrid[0:ydim, 0:xdim, ]
y -= ydim/2                 # centered at the origin
x -= xdim/2
# now make a mask
mask = x**2+y**2 <= radius**2 # start with 2d
img[mask,:] = color

the extra ",:" at the end of the assignment lets you assign the color throughout the 3 channels in one shot.

Upvotes: 6

Related Questions