Search code examples
pythonnumpytensorflowkerasmask

Is it possible to mask tensor pixels in graph mode using a 2D array?


I want to mask an array of pixels with an array of ones and zeros. I'm using Keras in graph mode, so only using operations that can be performed on tensors.

Most symbolic operators seem to work in graph mode, and there is a subset of numpy operations in keras.backend. But there is no support for no iterating or enumerating.

So for example I have:

    arr = [[(1,2,3),(2,3,4),(4,5,6)],[(5,6,7),(6,7,8),(7,8,9)]]
    mask = [[1, 0, 1],[0, 0, 1]]
    # and I want
    arr * mask == [[(1,2,3),(0,0,0),(4,5,6)],[(0,0,0),(0,0,0),(7,8,9)]]

The actual arrays are images, so much bigger than this example. Is there a reasonable way to mask points like this?


Solution

  • From the code your provide, one of easiest way to do is using broadcasting, like @user1318499 mentioned in the comments. Since mask is in shape of (2,3), and arr is in shape of (2,3,3), mask can be expanded one dim to do broadcast with arr.

    import tensorflow as tf
    
    arr = tf.reshape(tf.range(2*3*3), [2,3,3])
    # array([[[ 0,  1,  2],
    #         [ 3,  4,  5],
    #         [ 6,  7,  8]],
    
    #        [[ 9, 10, 11],
    #         [12, 13, 14],
    #         [15, 16, 17]]], dtype=int32)
    mask = tf.constant([[1, 0, 1],[0, 0, 1]])
    res = arr * tf.expand_dims(mask, axis=-1)
    # array([[[ 0,  1,  2],
    #         [ 0,  0,  0],
    #         [ 6,  7,  8]],
    
    #        [[ 0,  0,  0],
    #         [ 0,  0,  0],
    #         [15, 16, 17]]], dtype=int32)