Hot to make custom preprocessing layer in tensorflow 2.4?

Question

I want to perform some transformation on the input batch during training. For example, if I have a batch of images of size (number of samples, width, height, channels), I want to replace the 3rd channel with the difference of the first two channels, then resize the image size and finally normalize it. I tried to define a custom layer:

class CustomLayer(tf.keras.layers.Layer):
    def __init__(self):
        super(CustomLayer, self).__init__()
    def build(self, input_shape):
        pass
    def call(self, input_):
    #Loaded images
    self.img_tr = []
    for image in input_:
        img_input = resize(image,(267,400))#from skimage import resize
        img_diff = (img_input[:,:,1]/np.max(img_input[:,:,1]))-((img_input[:,:,0]+img_input[:,:,2])/np.max(img_input[:,:,0]+img_input[:,:,2]))
        img_temp = np.zeros((267,400,3))
        img_temp[:,:,0] = img_input[:,:,0]/np.max(img_input[:,:,0])
        img_temp[:,:,1] = img_input[:,:,1]/np.max(img_input[:,:,1])
        img_temp[:,:,2] = img_diff/np.max(img_diff)
        self.img_tr.append(img_temp)
    self.img_tr= np.asarray(self.img_tr)
    return self.img_tr

Then I used:

input_0 = tf.keras.Input(shape = (None,None,3))
clayer = CustomLayer()
input_1 = clayer(input_0)

x = tf.keras.layers.Conv2D(filters = 16, kernel_size = (7,7), activation = tf.keras.activations.relu)(input_1)
x = tf.keras.layers.MaxPool2D(pool_size = (2,2))(x)
x = tf.keras.layers.Flatten()(x)
x = tf.keras.layers.Dense(units = 64, activation = tf.keras.activations.relu)(x)
output = tf.keras.layers.Dense(units = 12)(x)
model = tf.keras.Model(inputs = input_0, outputs = output)

model.compile(
    optimizer = tf.keras.optimizers.Adam(),
    loss = tf.keras.losses.SparseCategoricalCrossentropy(from_logits = True),
    metrics = tf.keras.metrics.SparseCategoricalAccuracy()
)

model.summary()

I get an error that says:

AttributeError: 'Tensor' object has no attribute 'ndim'

I think the issue related to the fact that my custom layer expects a 4d numpy array but the input is having this format:

<KerasTensor: shape=(None, None, None, 3) dtype=float32 (created by layer 'input_20')>

How can I resolve the issue? I cannot find a way to convert KerasTensor to a numpy array inside my custom layer.

Edit I tried to avoid for loops and numpy so I tried:

class CustomLayer(tf.keras.layers.Layer):
    def __init__(self):
        super(CustomLayer, self).__init__()
    def build(self, input_shape):
        pass
    
    def call(self, input_):
 
        input_ = tf.Variable(input_)
        img_input = tf.image.resize(input_,(267,400))
        img_diff = (img_input[:,:,:,1])-((img_input[:,:,:,0]+img_input[:,:,:,2]))
        img_input[:,:,:,2] = img_diff
        output_img = tf.image.per_image_standardization(img_input)
            
        return input_

However, when I use the custom layer in the functional API I get the error:

ValueError: Tensor-typed variable initializers must either be wrapped in an init_scope or callable (e.g., `tf.Variable(lambda : tf.truncated_normal([10, 40]))`) when building functions. Please file a feature request if this restriction inconveniences you.

It seems has something to do with tf.Variable . Even if I set validate_shape to False I still get the same error.

Answer 1

Simply removing tf.Variable does the job. Below is the full layer:

class CustomLayer(tf.keras.layers.Layer):

    def __init__(self):
        super(CustomLayer, self).__init__()
  
    def build(self, input_shape):
        pass
    
    def call(self, inp):
 
        img_input = tf.image.resize(inp, (267,400))
        img_diff = (img_input[:,:,:,1])-((img_input[:,:,:,0]+img_input[:,:,:,2]))
        img_diff = tf.expand_dims(img_diff, -1)
        img_input = tf.keras.layers.Concatenate()([img_input[:,:,:,:-1], img_diff])
        output_img = tf.image.per_image_standardization(img_input)
            
        return output_img

I used tf.keras.layers.Concatenate to replace the last channel of img_input with img_diff .

Here the running notebook