如何使用tf.image.resize_with_crop_or_pad处理子目录图片适配TensorFlow ImageDataGenerator

函数基础说明 #

tf.image.resize_with_crop_or_pad 是TensorFlow内置的无插值尺寸调整工具，执行逻辑如下：

• 若输入图片尺寸大于目标尺寸：自动从居中位置裁剪超出的边缘部分
• 若输入图片尺寸小于目标尺寸：自动在四周补0（黑边）填充到目标尺寸
  不需要额外指定填充规则，默认适配绝大多数分类任务的padding需求。

与ImageDataGenerator结合的实现步骤 #

1. 导入依赖 #

import tensorflow as tf
from tensorflow.keras.preprocessing.image import ImageDataGenerator

2. 定义全局参数与自定义预处理函数 #

# 自定义统一输出的图片尺寸，可根据任务需求调整
TARGET_HEIGHT = 224
TARGET_WIDTH = 224

def resize_with_padding(image):
    # 入参image是ImageDataGenerator传入的uint8格式numpy数组，形状为(原始高, 原始宽, 通道数)
    processed_img = tf.image.resize_with_crop_or_pad(image, TARGET_HEIGHT, TARGET_WIDTH)
    # 转回numpy数组适配后续处理逻辑
    return processed_img.numpy()

3. 初始化ImageDataGenerator并加载数据集 #

# 初始化数据生成器，其他增强参数（如旋转、翻转）照常添加即可
datagen = ImageDataGenerator(
    rescale=1./255,
    preprocessing_function=resize_with_padding
)

# 从TRAIN目录加载数据
train_generator = datagen.flow_from_directory(
    directory="./TRAIN",
    target_size=(TARGET_HEIGHT, TARGET_WIDTH),
    batch_size=32,
    class_mode="categorical",
    shuffle=True
)

注意事项 #

ImageDataGenerator的flow_from_directory方法默认会先将所有图片用插值方式缩放到target_size指定的尺寸，再执行自定义预处理函数。如果要完全避免插值缩放，仅保留裁剪/填充逻辑，可以先统计TRAIN目录下所有图片的最大高度与最大宽度，将target_size设置为该最大尺寸，再在预处理函数中将图片调整为最终需要的目标尺寸即可。

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