indexing a matrix from a vector array-CodePudding

I have two images, one is a RGB image and the other is a mask image that contains 0 and 1 to segment a specified object. (both images are of the same object) I want to extract the RBG values of the initial image only at the indexes where the second matrix is 1, so that the final value is an image of just the object with a black background. is there a simple way to achieve this in numpy?

I would like to solve this problem without using too many for loops, I think there should be a straight forward way in numpy but I have not had any luck so far

CodePudding user response：

Yes.

You can use numpy.multiply, from here.

For example

img = np.array([
                 [[1,2],[3,4]],
                 [[5,6],[7,8]],
                 [[9,10],[11,12]]
               ])

mask = np.array(
                 [[0,1],[0,1]]
               )

print(np.array([
        np.multiply(img[0],mask),
        np.multiply(img[1],mask),
        np.multiply(img[2],mask)]
    ))
# Res: 


    #[[[ 0  2]
    #  [ 0  4]]
    #
    # [[ 0  6]
    #   [ 0  8]]
    #
    #  [[ 0 10]
    #  [ 0 12]]]

CodePudding user response：

You can use numpys built-in broadcasting, and then just straight-out multiply the two in "pythonic" form.

import numpy as np

img = np.array([[[ 1,  2,  3],
                 [ 4,  5,  6]],
               [[ 7,  8,  9],
                [10, 11, 12]]])  # shape (2, 2, 3)

mask = np.array([[0,1],[0,1]])  # shape (2, 2)

masked_img = img * np.expand_dims(mask, -1)

Alternatively, you can expand the mask dimensions via np.newaxis:

masked_img = img * mask[..., np.newaxis]