如何将图像预处理逻辑集成到TFLite模型中 自动适配输入要求

完全可以将图像预处理逻辑集成到TFLite模型内部，无需推理侧单独实现预处理代码，最常用的实现方案是将预处理逻辑和原模型合并为新的SavedModel后再转换为TFLite格式，具体实现步骤如下：

步骤1：封装包含预处理的新模型 #

将原SavedModel和你需要的预处理逻辑合并成一个新的TensorFlow模型，示例代码如下：

import tensorflow as tf

# 加载原有SavedModel
original_model = tf.saved_model.load(SAVED_MODEL_PATH)
# 获取原模型的推理签名
infer = original_model.signatures['serving_default']

# 定义带预处理逻辑的包装模型
class PreprocessWrappedModel(tf.Module):
    def __init__(self, original_infer):
        super().__init__()
        self.original_infer = original_infer
    
    # 指定输入签名：接受任意尺寸的3通道float32图像输入
    @tf.function(input_signature=[tf.TensorSpec(shape=[None, None, 3], dtype=tf.float32, name='input_image')])
    def __call__(self, image):
        # 此处直接写入你需要的所有预处理逻辑，和手动实现的逻辑完全对应即可
        # 1.  resize到模型要求的320*320尺寸
        resized_image = tf.image.resize(image, [320, 320])
        # 2. 新增batch维度，适配模型输入要求
        input_data = tf.expand_dims(resized_image, 0)
        # 调用原模型执行推理
        infer_result = self.original_infer(input_data)
        # 按需返回需要的输出字段即可
        return infer_result

# 实例化包装后的模型
wrapped_model = PreprocessWrappedModel(infer)

# 导出合并了预处理逻辑的新SavedModel
NEW_SAVED_MODEL_PATH = './saved_model_with_preprocess'
tf.saved_model.save(
    wrapped_model,
    NEW_SAVED_MODEL_PATH,
    signatures={'serving_default': wrapped_model.__call__.get_concrete_function()}
)

步骤2：转换新模型为TFLite格式 #

用新导出的带预处理的SavedModel执行TFLite转换：

converter = tf.lite.TFLiteConverter.from_saved_model(NEW_SAVED_MODEL_PATH)
# 如有量化、算子兼容等配置可在此处添加
tflite_model = converter.convert()

# 保存最终的TFLite模型
with open('model_with_preprocess.tflite', 'wb') as f:
    f.write(tflite_model)

步骤3：验证推理效果 #

转换完成后，推理侧无需再单独实现预处理逻辑，直接传入原始图像即可得到推理结果：

import cv2
import numpy as np
import tensorflow as tf

# 加载带预处理的TFLite模型
interpreter = tf.lite.Interpreter(model_path='model_with_preprocess.tflite')
interpreter.allocate_tensors()
input_details = interpreter.get_input_details()
output_details = interpreter.get_output_details()

# 直接读取原始图像，不需要手动执行resize、加维度等操作
img = cv2.imread(IMAGE_PATH)
image_np = np.array(img, dtype=np.float32)

# 直接传入原始图像执行推理
interpreter.set_tensor(input_details[0]['index'], image_np)
interpreter.invoke()
output_data = interpreter.get_tensor(output_details[0]['index'])
print(output_data)

注意事项 #

• 所有TensorFlow原生支持的图像处理算子都可以直接集成到预处理逻辑中，包括自定义的转换操作（如归一化、通道顺序转换、减均值等），只需把对应逻辑写到__call__方法内即可。
• 如你使用的是TensorFlow 1.x版本，也可以通过在原计算图中插入预处理节点的方式实现，整体逻辑和上述方案一致。

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