TensorFlow模型单图识别正常，批量识别文件夹图片报错求助

错误原因 #

cv2.imread(file)返回None，因为os.walk返回的file只是单纯的文件名，不是完整文件路径，程序找不到目标图片，导致后续tf.image.resize尝试将None转换为Tensor时触发报错。

修复后的代码 #

修改遍历部分的逻辑，构建图片的完整路径，同时增加异常处理避免单个图片读取失败导致程序中断：

import tensorflow as tf
import os
import cv2
import imghdr
import numpy as np

# GPU配置（原代码保留）
gpus = tf.config.experimental.list_physical_devices('GPU')
for gpu in gpus:
    tf.config.experimental.set_memory_growth(gpu, True)

# 数据集清理（原代码保留）
data_dir = 'dataset'
image_exts = ['jpeg','jpg', 'bmp', 'png']
for image_class in os.listdir(data_dir):
    for image in os.listdir(os.path.join(data_dir, image_class)):
        image_path = os.path.join(data_dir, image_class, image)
        try:
            img = cv2.imread(image_path)
            tip = imghdr.what(image_path)
            if tip not in image_exts:
                print('Image not in ext list {}'.format(image_path))
                os.remove(image_path)
        except Exception as e:
            print('Issue with image {}'.format(image_path))

# 数据集加载与预处理（原代码保留）
data = tf.keras.utils.image_dataset_from_directory('dataset')
data_iterator = data.as_numpy_iterator()
batch = data_iterator.next()
data = data.map(lambda x,y: (x/255, y))
data.as_numpy_iterator().next()

train_size = int(len(data)*.7)
val_size = int(len(data)*.2)
test_size = int(len(data)*.1)

train = data.take(train_size)
val = data.skip(train_size).take(val_size)
test = data.skip(train_size+val_size).take(test_size)

# 模型构建与训练（原代码保留）
from tensorflow.keras.models import Sequential
from tensorflow.keras.layers import Conv2D, MaxPooling2D, Dense, Flatten, Dropout

model = Sequential()
model.add(Conv2D(16, (3,3), 1, activation='relu', input_shape=(256,256,3)))
model.add(MaxPooling2D())
model.add(Conv2D(32, (3,3), 1, activation='relu'))
model.add(MaxPooling2D())
model.add(Conv2D(16, (3,3), 1, activation='relu'))
model.add(MaxPooling2D())
model.add(Flatten())
model.add(Dense(256, activation='relu'))
model.add(Dense(1, activation='sigmoid'))

model.compile('adam', loss=tf.losses.BinaryCrossentropy(), metrics=['accuracy'])
logdir='logs'
tensorboard_callback = tf.keras.callbacks.TensorBoard(log_dir=logdir)
hist = model.fit(train, epochs=10, validation_data=val, callbacks=[tensorboard_callback])

# 修复后的批量预测逻辑
folder_path = "test"
for root, dirs, files in os.walk(folder_path):
    for file in files:
        # 构建完整图片路径
        img_path = os.path.join(root, file)
        print(f"Processing: {img_path}")
        # 读取图片
        img = cv2.imread(img_path)
        # 检查图片是否读取成功
        if img is None:
            print(f"Failed to read image: {img_path}")
            continue
        # 预处理图片
        resize = tf.image.resize(img, (256,256))
        yhat = model.predict(np.expand_dims(resize/255, 0), verbose=0)
        # 输出预测结果
        if yhat > 0.5:
            print(f'Predicted class is Sneakers ')
        else:
            print(f'Predicted class is Flats')

额外优化建议 #

• 添加图片格式校验：在批量预测时，可复用之前的imghdr工具校验图片格式，避免非图片文件导致报错
• 使用TensorFlow原生工具读取图片：可以用tf.io.read_file和tf.image.decode_image替代cv2.imread，更好地兼容TensorFlow的张量流程
• 批量处理图片：如果图片数量较多，可将图片整理成批量张量后再预测，提升效率

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