使用 Tensorflow 计算像素邻域中的像素 colors

Question

I have an image, of shape [1024, 1024, 3] in a reduced color space (like ~100 colors). I need to count, for every possible couple of colors, how many times one appears in the neighborhood of the other. This is quite simple in python, the problem I have is to write it using Tensorflow functions. I will make a simple example to make it easier to understand:

Let`s consider a color space with only 4 colors (named 1,2,3,4 ) and a 4x4 image:

[ [2 4 1 4]  
  [2 4 2 1]  
  [2 4 3 1]  
  [2 1 1 2] ]

From this I have to compute the neighborhood color histogram , that explain what is written above: how many pixel of every color appears in the neighborhood of every pixel of every other color. The NCH of the above image is the following:

[ [8 7 4 6]  
  [7 6 2 12]  
  [4 2 0 2]  
  [6 12 2 4] ]

To make it clear, the 8 in the first row means that color 1 appears in the neighborhoods of color 1 8 times. The 7 indicates that color 2 appears in the neighborhoods of color 1 7 times and so on.

In this simple example I consider a neighborhood of size 3x3 , but the code I am working on should be extendable to a generic NxN size.

The only step I was able to implement with Tensorflow is to split each image in the batch (I am implementing this as the first step of a loss function, so it has to work with batches) in NxN patches, centered on every pixel, using the following function:

patches = tf.image.extract_patches(image, sizes=(1, D_size, D_size, 1), strides=(1, 1, 1, 1),
                                         padding='SAME', rates=[1, 1, 1, 1])

What I have to do now is to iterate over every patch and increment, according to the central color of the patch, the NCH matrix. This is very simple using loops, but I struggle to find the right Tensorflow functions to execute them in a more parallel way.

I am pretty new to the tensorflow world, so I understand it may seems an obvious question, but I really do not know what else to try and I will prefer to not use loops. Thank you all in advance.

Answer 1

Here's a way of doing it. I was expecting this to some out at about 5 lines of code, but it's turned out a lot longer. (I've left in some diagnostic print statements, might help understand what it is doing) If anyone knows a more elegant way I'd be very interested.

As Alberto has said, probably don't expect to be able to backprop through this.

def count_neighbours(image, ncolours=5, adjacency=1):
    '''
    Note colours have to be numbers 1 to ncolours, NOT zero
    Zero values are created during the calc, both in padding and in tensor products,
    and must be discarded
    '''
    f = 2 * adjacency + 1      # filter size
    ic = adjacency * (f + 1)   # Index of central pixel in flattened patch
    print(f"{adjacency=} {ic=} {f=}")
    print(f"{image.shape=}")

    # Get patches. Returns location of patch in dims 1, 2, and (flattened) patch in 3
    patches = tf.image.extract_patches(image, sizes=(1, f, f, 1), strides=(1, 1, 1, 1),
                                         padding='SAME', rates=[1, 1, 1, 1])
    print(f"{patches.shape=}") 
    
    # A sort of outer product with one-hot encoding of the central pixel in a patch
    # So if a patch [i, j, k, :] has central pixel l, it is copied to [i, j, k, :, l]
    # First, create a one-hot encoding of the central pixel in each patch
    oh_central = tf.one_hot(patches[:,:,:,ic], axis=-1, depth=ncolours + 1, dtype=tf.int32)  # one-hot of central pixel in each patch
    print(f"{oh_central.shape=}") 

    # Want oh1[m, h, w, ipixel, colour] = patches[m, h, w, ipixel] * oh_central[m, h, w, colour]
    oh1 = tf.einsum('ijkl,ijkm->ijklm',patches, oh_central)
    print(f"{oh1.shape=}")
    
    # One-hot encode the patches
    oh2 = tf.one_hot(oh1, axis=-1, depth=ncolours + 1, dtype=tf.int32)
    print(f"{oh2.shape=}")
    
    # Set central pixels to zero, since a pixel is not counted as being adjacent to itself
    mask = tf.concat([
            tf.ones((oh2.shape[0], oh2.shape[1], oh2.shape[2], oh2.shape[3] //2, oh2.shape[4], oh2.shape[5]), dtype=tf.int32),
            tf.zeros((oh2.shape[0], oh2.shape[1], oh2.shape[2], 1, oh2.shape[4], oh2.shape[5]), dtype=tf.int32),
            tf.ones((oh2.shape[0], oh2.shape[1], oh2.shape[2], oh2.shape[3] //2, oh2.shape[4], oh2.shape[5]), dtype=tf.int32),
                ], axis=3)
    assert mask.shape == oh2.shape
    oh3 = oh2 * mask

    # finally, reduce sum, discard counts of zeros and return result
    result = tf.reduce_sum(oh3, axis=[1,2,3])[:,1:,1:]
    return result

Then this can be called like this, and gives the results you posted.

x = tf.constant([ [2, 4, 1, 4]  ,
  [2, 4, 2, 1]  ,
  [2, 4, 3, 1]  ,
  [2, 1, 1, 2] ], dtype=tf.int32)
x = tf.expand_dims(tf.expand_dims(x,0),3)  # add dummy batch and channel dimensions
count_neighbours(x, ncolours=4)