How to zero out row slice in 2 dimensional array in numpy?

Question

I have a numpy array representing an image. I want to zero out all indexes that are below a certain row in each column (based on a external data). I can't seem to figure out how to slice/broadcast/arrange the data to do this "the numpy way".

def first_nonzero(arr, axis, invalid_val=-1):
    mask = arr!=0
    return np.where(mask.any(axis=axis), mask.argmax(axis=axis), invalid_val)

# Find first non-zero pixels in a processed image
# Note, I might have my axes switched here... I'm not sure.
rows_to_zero = first_nonzero(processed_image, 0, processed_image.shape[1])

# zero out data in image below the rows found
# This is the part I'm stuck on.
image[:, :rows_to_zero, :] = 0  # How can I slice along an array of indexes?

# Or in plain python, I'm trying to do this:
for x in range(image.shape[0]):
    for y in range(rows_to_zero, image.shape[1]):
        image[x,y] = 0

Divakar · Accepted Answer · 2017-12-12 04:11:14Z

Create a mask leveraging broadcasting and assign -

mask = rows_to_zero <= np.arange(image.shape[0])[:,None]
image[mask] = 0

Or multiply with the inverted mask : image *= ~mask.

Sample run to showcase mask setup -

In [56]: processed_image
Out[56]: 
array([[1, 0, 1, 0],
       [1, 0, 1, 1],
       [0, 1, 1, 0],
       [0, 1, 0, 1],
       [1, 1, 1, 1],
       [0, 1, 0, 1]])

In [57]: rows_to_zero
Out[57]: array([0, 2, 0, 1])

In [58]: rows_to_zero <= np.arange(processed_image.shape[0])[:,None]
Out[58]: 
array([[ True, False,  True, False],
       [ True, False,  True,  True],
       [ True,  True,  True,  True],
       [ True,  True,  True,  True],
       [ True,  True,  True,  True],
       [ True,  True,  True,  True]], dtype=bool)

Also, for setting per column basis, I think you meant :

rows_to_zero = first_nonzero(processed_image, 0, processed_image.shape[0]-1)

If you meant to zero out on per row basis, you would have the indices first non-zero indices per row, let's call it idx. So, then do -

mask = idx[:,None] <= np.arange(image.shape[1])
image[mask] = 0

Sample run -

In [77]: processed_image
Out[77]: 
array([[1, 0, 1, 0],
       [1, 0, 1, 1],
       [0, 1, 1, 0],
       [0, 1, 0, 1],
       [1, 1, 1, 1],
       [0, 1, 0, 1]])

In [78]: idx = first_nonzero(processed_image, 1, processed_image.shape[1]-1)

In [79]: idx
Out[79]: array([0, 0, 1, 1, 0, 1])

In [80]: idx[:,None] <= np.arange(image.shape[1])
Out[80]: 
array([[ True,  True,  True,  True],
       [ True,  True,  True,  True],
       [False,  True,  True,  True],
       [False,  True,  True,  True],
       [ True,  True,  True,  True],
       [False,  True,  True,  True]], dtype=bool)

Thanks, this helped so much! And you were right, I needed first_nonzero(image, 0, image.shape[0]-1)

Collectives™ on Stack Overflow

How to zero out row slice in 2 dimensional array in numpy?

1 Answer 1

1 Comment

Your Answer

Hot Network Questions

Collectives™ on Stack Overflow

1 Answer 1

1 Comment

Your Answer

Sign up or log in

Post as a guest

Related