我试图建立一个自定义的keras层，用OpenCV做Canny边缘检测。

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)

我的模型如下：

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)

我已经在一些测试图像上尝试了我的自定义图层，它工作正常，并成功地输出了一批形状为（n，h，w，1）的Tensor。但当我试图建立我的模型时，我得到了以下错误：

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

哪里出了问题，我应该如何正确地指定我的自定义图层的输出形状？