如何将自定义OpenCV图像预处理代码转为TensorFlow Sequential层？

问题背景 #

需要把包含OpenCV操作（自适应阈值、形态学变换、轮廓检测与绘制）的自定义图像处理逻辑，封装成TensorFlow Keras可识别的层，嵌入到现有的图像分类Sequential模型中。

解决方案步骤 #

1. 封装自定义Keras层 #

通过继承tf.keras.layers.Layer，结合tf.py_function将基于OpenCV的numpy操作包装为TensorFlow兼容的层，需处理批量图像输入并确保张量与numpy数组的正确转换：

import cv2
import numpy as np
import tensorflow as tf
from tensorflow import keras
from tensorflow.keras import layers

class OpenCVPreprocessingLayer(layers.Layer):
    def __init__(self, image_size=512, blob_area_thresh=1000, **kwargs):
        super().__init__(**kwargs)
        self.image_size = image_size
        self.blob_area_thresh = blob_area_thresh

    def call(self, inputs):
        return tf.py_function(self._process_batch, [inputs], tf.float32)
    
    def _process_batch(self, batch_imgs):
        processed_batch = []
        for img in batch_imgs:
            # 还原归一化后的图像到OpenCV要求的uint8格式
            img_np = (img.numpy() * 255).astype(np.uint8)
            # 转为灰度图
            gray = cv2.cvtColor(img_np, cv2.COLOR_RGB2GRAY)
            
            # 自适应阈值处理
            thresh = cv2.adaptiveThreshold(gray, 255, cv2.ADAPTIVE_THRESH_MEAN_C, cv2.THRESH_BINARY, 101, 3)
            
            # 形态学开闭操作
            kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (5,5))
            blob = cv2.morphologyEx(thresh, cv2.MORPH_OPEN, kernel)
            kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (9,9))
            blob = cv2.morphologyEx(blob, cv2.MORPH_CLOSE, kernel)
            
            # 反转图像
            blob = 255 - blob
            
            # 轮廓检测与筛选
            cnts = cv2.findContours(blob, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
            cnts = cnts[0] if len(cnts) == 2 else cnts[1]
            
            # 处理无轮廓或轮廓过小的情况
            if len(cnts) == 0:
                result = img_np
            else:
                big_contour = max(cnts, key=cv2.contourArea)
                blob_area = cv2.contourArea(big_contour)
                if blob_area < self.blob_area_thresh:
                    result = img_np
                else:
                    result = img_np.copy()
                    cv2.drawContours(result, [big_contour], -1, (0,0,255), 1)
            
            # 转回归一化的float32格式
            processed_batch.append(result.astype(np.float32) / 255.0)
        
        return np.array(processed_batch, dtype=np.float32)
    
    def compute_output_shape(self, input_shape):
        return input_shape

2. 嵌入到现有模型流水线 #

将自定义层添加到Sequential模型中，放置在resize_and_rescale之后、data_augmentation之前（可根据需求调整顺序）：

# 沿用原有流水线定义
dataset = tf.keras.preprocessing.image_dataset_from_directory(
    "../dataset",
    shuffle = True,
    image_size = (IMAGE_SIZE, IMAGE_SIZE),
    batch_size = BATCH_SIZE
)
train_ds, val_ds, test_ds = get_dataset_partitions_tf(dataset)

train_ds = train_ds.cache().shuffle(1000).prefetch(buffer_size = tf.data.AUTOTUNE)
val_ds = val_ds.cache().shuffle(1000).prefetch(buffer_size = tf.data.AUTOTUNE)
test_ds = test_ds.cache().shuffle(1000).prefetch(buffer_size = tf.data.AUTOTUNE)

resize_and_rescale = tf.keras.Sequential([
    layers.experimental.preprocessing.Resizing(IMAGE_SIZE, IMAGE_SIZE),
    layers.experimental.preprocessing.Rescaling(1.0/255)
])

data_augmentation = tf.keras.Sequential([
    layers.experimental.preprocessing.RandomFlip("horizontal_and_vertical"),
    layers.experimental.preprocessing.RandomRotation(0.2)
])

# 初始化自定义预处理层
custom_preprocessing = OpenCVPreprocessingLayer(image_size=IMAGE_SIZE)

# 重构模型
model = keras.Sequential([
    resize_and_rescale,
    custom_preprocessing,  # 插入自定义层
    data_augmentation,
    layers.Conv2D(32, (3,3), activation='relu', input_shape=(IMAGE_SIZE, IMAGE_SIZE, CHANNELS)),
    layers.MaxPooling2D((2,2)),
    layers.Conv2D(64, (3,3), activation='relu'),
    layers.MaxPooling2D((2,2)),
    layers.Conv2D(64, (3,3), activation='relu'),
    layers.MaxPooling2D((2,2)),
    layers.Conv2D(64, (3,3), activation='relu'),
    layers.MaxPooling2D((2,2)),
    layers.Conv2D(64, (3,3), activation='relu'),
    layers.MaxPooling2D((2,2)),
    layers.Conv2D(64, (3,3), activation='relu'),
    layers.MaxPooling2D((2,2)),
    layers.Flatten(),
    layers.Dense(64, activation='relu'),
    layers.Dense(nClassNum, activation='softmax'),
])

model.build(input_shape=(BATCH_SIZE, IMAGE_SIZE, IMAGE_SIZE, CHANNELS))

model.compile(
    optimizer='adam',
    loss= tf.keras.losses.SparseCategoricalCrossentropy(from_logits=False),
    metrics=['accuracy']
)

3. 关键注意事项 #

• 批量适配：自定义层必须遍历批量中的每张图像处理，避免单张图像逻辑直接作用于批量张量。
• 格式转换：通过tf.py_function实现张量与numpy数组的双向转换，适配OpenCV的输入输出要求。
• 异常兜底：添加无轮廓、轮廓过小的处理逻辑，避免训练时抛出异常。
• 性能优化：tf.py_function会带来一定性能开销，若追求更高效率，可尝试将部分操作替换为TensorFlow原生API（如tfa.image中的形态学操作）。

内容的提问来源于stack exchange，提问作者Mohammad Ishfakur