如何基于文件夹结构将图像与NPY文件输入VGG16训练

问题背景 #

使用VGG16搭建了双输入模型，需要将图像（存储在按fake/real分类的文件夹中）与对应的20维NPY特征文件同时输入模型训练。图像与NPY文件分属不同文件夹，但文件名一一对应（如pic_test/fake/01_pic.png对应feature_test/fake/01_npy.npy）。尝试过ImageDataGenerator但仅支持图像输入，需解决双输入数据的加载问题。

现有代码 #

train_datagen = ImageDataGenerator(
        rescale=1./255,
        shear_range=0.2,
        zoom_range=0.2,
        horizontal_flip=True
)

test_datagen = ImageDataGenerator(rescale=1./255)

batch_size=8 

train_generator = train_datagen.flow_from_directory(
        './pic_train',      
        target_size=(224,224),
        batch_size=batch_size, 
        class_mode='categorical'
        )

test_generator = test_datagen.flow_from_directory(
        './pic_test',
        target_size=(224,224),
        batch_size=batch_size, 
        class_mode='categorical'
        )

解决方案 #

方法一：自定义Sequence类（Keras原生支持） #

通过继承keras.utils.Sequence实现自定义双输入数据生成器，支持数据增强、批量加载，且兼容Keras的fit方法。

import numpy as np
import os
from PIL import Image
from tensorflow.keras.utils import Sequence
from tensorflow.keras.preprocessing.image import img_to_array, load_img

class DualInputSequence(Sequence):
    def __init__(self, img_dir, npy_dir, target_size=(224,224), batch_size=8, datagen=None, class_mode='categorical'):
        self.img_dir = img_dir
        self.npy_dir = npy_dir
        self.target_size = target_size
        self.batch_size = batch_size
        self.datagen = datagen
        self.class_mode = class_mode
        
        # 建立类别与索引的映射
        self.classes = sorted(os.listdir(img_dir))
        self.class_indices = {cls: idx for idx, cls in enumerate(self.classes)}
        
        # 收集所有图像-NPY对及对应标签
        self.data_pairs = []
        for cls in self.classes:
            img_cls_dir = os.path.join(img_dir, cls)
            npy_cls_dir = os.path.join(npy_dir, cls)
            for img_name in os.listdir(img_cls_dir):
                if img_name.endswith('.png'):
                    # 根据命名规则匹配对应NPY文件
                    npy_name = img_name.replace('_pic.png', '_npy.npy')
                    img_path = os.path.join(img_cls_dir, img_name)
                    npy_path = os.path.join(npy_cls_dir, npy_name)
                    if os.path.exists(npy_path):
                        self.data_pairs.append((img_path, npy_path, self.class_indices[cls]))
        
    def __len__(self):
        # 返回总batch数
        return len(self.data_pairs) // self.batch_size
    
    def __getitem__(self, idx):
        # 加载当前batch的图像、NPY和标签
        batch_data = self.data_pairs[idx*self.batch_size : (idx+1)*self.batch_size]
        
        imgs, npys, labels = [], [], []
        for img_path, npy_path, label in batch_data:
            # 加载并预处理图像
            img = load_img(img_path, target_size=self.target_size)
            img_array = img_to_array(img) / 255.0  # 基础归一化
            
            if self.datagen:
                # 应用数据增强（仅训练集使用）
                img_array = self.datagen.random_transform(img_array)
            
            # 加载NPY特征
            npy_array = np.load(npy_path)
            
            imgs.append(img_array)
            npys.append(npy_array)
            labels.append(label)
        
        # 转换为numpy数组
        imgs = np.array(imgs)
        npys = np.array(npys)
        
        # 处理标签格式
        if self.class_mode == 'categorical':
            labels = np.eye(len(self.classes))[labels]
        
        # 返回双输入和标签
        return [imgs, npys], labels

使用方式

# 训练集数据增强配置
train_datagen = ImageDataGenerator(
        shear_range=0.2,
        zoom_range=0.2,
        horizontal_flip=True
)

# 初始化训练/测试序列
train_sequence = DualInputSequence(
    img_dir='./pic_train',
    npy_dir='./feature_train',
    target_size=(224,224),
    batch_size=8,
    datagen=train_datagen,
    class_mode='categorical'
)

test_sequence = DualInputSequence(
    img_dir='./pic_test',
    npy_dir='./feature_test',
    target_size=(224,224),
    batch_size=8,
    datagen=None,  # 测试集不做数据增强
    class_mode='categorical'
)

# 模型训练
model.fit(
    train_sequence,
    validation_data=test_sequence,
    epochs=10,
    # 其他训练参数...
)

方法二：使用tf.data.Dataset（适合大规模数据） #

利用TensorFlow的tf.data.DatasetAPI构建高效的双输入数据集，支持并行加载和预处理，性能更优。

import tensorflow as tf
import numpy as np
import os

def load_data(img_path, npy_path, label):
    # 加载并预处理图像
    img = tf.io.read_file(img_path)
    img = tf.image.decode_png(img, channels=3)
    img = tf.image.resize(img, (224,224))
    img = img / 255.0  # 归一化
    
    # 加载NPY特征
    npy = tf.numpy_function(np.load, [npy_path], tf.float32)
    npy = tf.reshape(npy, (20,))  # 确保为20维向量
    
    return (img, npy), label

def augment_data(inputs, label):
    # 对图像应用数据增强
    img, npy = inputs
    img = tf.image.random_flip_left_right(img)
    img = tf.image.random_shear(img, 0.2)
    img = tf.image.random_zoom(img, (0.8, 1.2))
    return (img, npy), label

def create_dataset(img_dir, npy_dir, batch_size=8, shuffle=True, augment=False):
    # 建立类别映射
    classes = sorted(os.listdir(img_dir))
    class_indices = {cls: idx for idx, cls in enumerate(classes)}
    
    # 收集所有图像-NPY对及标签
    img_paths, npy_paths, labels = [], [], []
    for cls in classes:
        img_cls_dir = os.path.join(img_dir, cls)
        npy_cls_dir = os.path.join(npy_dir, cls)
        for img_name in os.listdir(img_cls_dir):
            if img_name.endswith('.png'):
                npy_name = img_name.replace('_pic.png', '_npy.npy')
                img_path = os.path.join(img_cls_dir, img_name)
                npy_path = os.path.join(npy_cls_dir, npy_name)
                if os.path.exists(npy_path):
                    img_paths.append(img_path)
                    npy_paths.append(npy_path)
                    labels.append(class_indices[cls])
    
    # 构建数据集
    dataset = tf.data.Dataset.from_tensor_slices((img_paths, npy_paths, labels))
    dataset = dataset.map(load_data, num_parallel_calls=tf.data.AUTOTUNE)
    
    if shuffle:
        dataset = dataset.shuffle(buffer_size=len(img_paths))
    
    if augment:
        dataset = dataset.map(augment_data, num_parallel_calls=tf.data.AUTOTUNE)
    
    # 批量处理并预取
    dataset = dataset.batch(batch_size).prefetch(tf.data.AUTOTUNE)
    return dataset

使用方式

# 创建训练/测试数据集
train_dataset = create_dataset(
    './pic_train', './feature_train',
    batch_size=8, shuffle=True, augment=True
)

test_dataset = create_dataset(
    './pic_test', './feature_test',
    batch_size=8, shuffle=False, augment=False
)

# 模型训练
model.fit(
    train_dataset,
    validation_data=test_dataset,
    epochs=10,
    # 其他训练参数...
)

注意事项 #

1. 文件名匹配：确保代码中文件名替换逻辑与你的实际命名规则一致，若命名规则不同，需修改img_name.replace(...)的部分。
2. 数据增强：仅对图像应用增强，NPY特征不需要做任何变换。
3. 模型输入匹配：模型需包含两个输入层，分别接收(224,224,3)的图像张量和(20,)的NPY特征张量，再通过拼接层合并后进入后续网络。

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