keras custom layer unknown output shape-CodePudding

I am trying to build a custom keras layer that does Canny edge detection with OpenCV. Here's my code:

class CannyEdgeDetectorLayer(layers.Layer):
    def __init__(self, threshold1=60, threshold2=120, **kwargs):
        super(CannyEdgeDetectorLayer, self).__init__(**kwargs)
        self.threshold1 = threshold1
        self.threshold2 = threshold2
    
    def call(self, inputs):
        return tf.py_function(func=self.canny_edge_detector, inp=[inputs], Tout=tf.float32)

    def canny_edge_detector(self, inputs):
        inputs = inputs.numpy()
        edges = [cv2.Canny(img, self.threshold1, self.threshold2).reshape(inputs.shape[1], inputs.shape[2], -1) / 255 for img in inputs]
        return tf.reshape(tf.convert_to_tensor(edges, dtype=tf.float32), (inputs.shape[0], inputs.shape[1], inputs.shape[2], 1))
        # return np.array(edges).reshape(inputs.shape[0], inputs.shape[1], inputs.shape[2], 1)
    
    def compute_output_shape(self, input_shape):
        return (input_shape[0], input_shape[1], input_shape[2], 1)

    def get_config(self):
        config = super().get_config().copy()
        config.update({
            'threshold1': self.threshold1,
            'threshold2': self.threshold2
        })
        return config
    
    def build(self, input_shape):
        return super().build(input_shape)

And my model is as follows:

inputs = keras.Input(shape=(255, 255, 3))
x = CannyEdgeDetectorLayer(60, 120)(inputs)
x = layers.RandomFlip('horizontal')(x)
x = layers.RandomRotation(1./12)(x)
x = layers.Conv2D(filters=32, kernel_size=(3, 3), activation='gelu')(x)
x = layers.MaxPooling2D(pool_size=2)(x)
x = layers.Conv2D(filters=64, kernel_size=(3, 3), activation='gelu')(x)
x = layers.Conv2D(filters=64, kernel_size=(3, 3), activation='gelu')(x)
x = layers.MaxPooling2D(pool_size=2)(x)
x = layers.Conv2D(filters=128, kernel_size=(3, 3), activation='gelu')(x)
x = layers.Conv2D(filters=128, kernel_size=(3, 3), activation='gelu')(x)
x = layers.MaxPooling2D(pool_size=2)(x)
x = layers.Conv2D(filters=256, kernel_size=(3, 3), activation='gelu')(x)
x = layers.Conv2D(filters=256, kernel_size=(3, 3), activation='gelu')(x)
x = layers.MaxPooling2D(pool_size=2)(x)
x = layers.Flatten()(x)
x = layers.Dense(128, activation='gelu')(x)
x = layers.Dropout(0.5)(x)
x = layers.Dense(64, activation='gelu')(x)
x = layers.Dropout(0.5)(x)
outputs = layers.Dense(1, activation='sigmoid')(x)

model = Model(inputs=inputs, outputs=outputs)

I have tried my custom layers on some test images, and it worked fine and successfully outputted a batch of tensors with shape (n, h, w, 1). But when I am trying to build my model I get the following error:

Image augmentation layers are expecting inputs to be rank 3 (HWC) or 4D (NHWC) tensors. Got shape: <unknown>

Call arguments received by layer "random_flip_25" (type RandomFlip):
  • inputs=tf.Tensor(shape=<unknown>, dtype=float32)
  • training=True

What went wrong and how should I properly specify the output shape of my custom layer?

CodePudding user response：

I solved the problem by specify the output's shape in call:

def call(self, inputs):
    out = tf.py_function(func=self.canny_edge_detector, inp=[inputs], Tout=tf.float32)
    out.set_shape((inputs.shape[0], inputs.shape[1], inputs.shape[2], 1))
    return out

Turns out it's a problem of py_function and EagerTensor.