What is the correct way to implement a basic GLCM-Layer in Tensorflow/Keras?

Question

I am trying to get a GLCM implementation running in a custom Keras Layer in a reasonable fast time. So far I took the _glcm_loop from skimage-implementation, reduced it to what I needed and put it into a basic layer, like this:

import numpy as np
import tensorflow as tf
from time import time
from tensorflow import keras
from tensorflow.keras.preprocessing import image
from tensorflow.keras import layers
from skimage.feature import *
from numpy import array
from math import sin, cos
from time import time
import matplotlib.pyplot as plt

class GLCMLayer(keras.layers.Layer):
    def __init__(self, greylevels=32, angles=[0], distances=[1], name=None, **kwargs):
        self.greylevels = greylevels
        self.angles = angles
        self.distances = distances
        super(GLCMLayer, self).__init__(name=name, **kwargs)

    def _glcm_loop(self, image, distances, angles, levels, out):
        rows = image.shape[0]
        cols = image.shape[1]

        for a_idx in range(len(angles)):
            angle = angles[a_idx]
            for d_idx in range(len(distances)):
                distance = distances[d_idx]
                offset_row = round(sin(angle) * distance)
                offset_col = round(cos(angle) * distance)
                start_row = max(0, -offset_row)
                end_row = min(rows, rows - offset_row)
                start_col = max(0, -offset_col)
                end_col = min(cols, cols - offset_col)
                for r in range(start_row, end_row):
                    for c in range(start_col, end_col):
                        i = image[r, c]
                        row = r + offset_row
                        col = c + offset_col
                        j = image[row, col]
                        out[i, j, d_idx, a_idx] += 1

    def call(self, inputs):
        P = np.zeros((self.greylevels, self.greylevels, len(self.distances), len(self.angles)), dtype=np.uint32, order='C')
        self._glcm_loop(inputs, self.distances, self.angles, self.greylevels, P)
        return P

    def get_config(self):
        config = {
            'angle': self.angle,
            'distance': self.distance,
            'greylevel': self.greylevel,
        }
        base_config = super(GLCMLayer, self).get_config()
        return dict(list(base_config.items()) + list(config.items()))

My execution code looks like this:

def quant(img, greylevels):
    return array(img)//(256//greylevels)

if __name__ == "__main__":
    source_file = "<some sour file>"
    img_raw = image.load_img(source_file, target_size=(150,150), color_mode="grayscale")
    img = quant(img_raw, 32)

    layer = GLCMLayer()
    start = time()
    aug = layer(img)
    tf.print(time()-start)

This is my first step to create it as a preprocessing layer. The second step then will be to modify it to run it also as hidden layer inside a model. That is why I didn't put it to a complete model yet, but I feel like there will be additional changes required when doing so.

For some reason the execution time is about 15-20 seconds long. Executing the code on the CPU without the layer takes about 0.0009 seconds. Obviously, something is going wrong here.

I am fairly new to tf and keras, so I fear I am missing something in the way on how to use the framework. In order to resolve it, I read about (which doesn't mean I understood):

• do not use np-functions inside tensorflow, but tf-functions instead,
• use tf.Variable,
• use tf.Data,
• unfolding is not possible in some way (whatever that means)

I tried a little here and there, but couldn't get them running, instead finding various different exceptions. So my questions are:

1. What is the correct way to use tf-functions in a GLCM to get the best performance on the GPU?
2. What do I need to take care of when using the layer in a complete model?

From that point on, I should hopefully be able to then implement the GLCM properties. Any help is greatly appreciated.

(Disclaimer: I assume that there is a lot of other stuff not optimal yet, if anything comes to your mind just add it.)