如何构建类似tf.keras.datasets.cifar10.load_data()返回格式的自定义手写数学图像数据集

If you want your custom handwritten math dataset to match the format returned by tf.keras.datasets.cifar10.load_data()—which uses numpy arrays for images and integer labels—here are two straightforward methods to achieve this:

Method 1: Using TensorFlow's image_dataset_from_directory (Recommended) #

This is the modern, TensorFlow-native approach that handles loading and formatting with minimal code.

Step 1: Load the dataset as a TensorFlow Dataset #

import tensorflow as tf
import numpy as np

# Load training data
train_ds = tf.keras.utils.image_dataset_from_directory(
    'train_test_data/train/',
    image_size=(45, 45),  # Match your model's input size
    batch_size=None,      # Load all images at once instead of batches
    label_mode='int'      # Get integer labels (like CIFAR-10)
)

# Load validation/test data
test_ds = tf.keras.utils.image_dataset_from_directory(
    'train_test_data/val/',
    image_size=(45, 45),
    batch_size=None,
    label_mode='int'
)

Step 2: Convert to NumPy arrays #

# Convert training data to numpy arrays
x_train = np.concatenate([x for x, y in train_ds], axis=0)
y_train = np.concatenate([y for x, y in train_ds], axis=0)
y_train = y_train.reshape(-1, 1)  # Match CIFAR-10's 2D label shape

# Convert test data to numpy arrays
x_test = np.concatenate([x for x, y in test_ds], axis=0)
y_test = np.concatenate([y for x, y in test_ds], axis=0)
y_test = y_test.reshape(-1, 1)

Step 3: Rescale pixel values (optional, match your original preprocessing) #

If you want to normalize pixels to the [0, 1] range like your original ImageDataGenerator setup:

x_train = x_train / 255.0
x_test = x_test / 255.0

Method 2: Manual Loading with PIL and NumPy #

This method gives you full control over the loading process, useful if you need custom preprocessing.

Step 1: Define a helper function to load data from directories #

import os
import numpy as np
from PIL import Image

def load_custom_data(directory, target_size=(45, 45)):
    images = []
    labels = []
    # Get sorted class names (to ensure consistent label indexing)
    class_names = sorted([d for d in os.listdir(directory) if os.path.isdir(os.path.join(directory, d))])
    class_to_label = {name: idx for idx, name in enumerate(class_names)}
    
    for class_name in class_names:
        class_dir = os.path.join(directory, class_name)
        for img_file in os.listdir(class_dir):
            img_path = os.path.join(class_dir, img_file)
            try:
                # Open image and convert to RGB (matches your model's 3-channel input)
                img = Image.open(img_path).convert('RGB')
                img = img.resize(target_size)
                # Convert to numpy array
                img_array = np.array(img)
                images.append(img_array)
                labels.append(class_to_label[class_name])
            except Exception as e:
                print(f"Skipping {img_path}: {str(e)}")
    
    # Convert lists to numpy arrays
    x = np.array(images)
    y = np.array(labels).reshape(-1, 1)  # 2D label array like CIFAR-10
    return x, y

Step 2: Load train and test data #

x_train, y_train = load_custom_data('train_test_data/train/')
x_test, y_test = load_custom_data('train_test_data/val/')

# Rescale to [0,1] if needed
x_train = x_train / 255.0
x_test = x_test / 255.0

Using the Data with Your CNN Model #

Now you can use the data exactly like you did with CIFAR-10:

# Compile the model (use sparse_categorical_crossentropy for integer labels)
cnn.compile(optimizer='adam', loss='sparse_categorical_crossentropy', metrics=['accuracy'])

# Train the model
cnn.fit(x_train, y_train, epochs=15, batch_size=32, validation_data=(x_test, y_test))

Key Notes: #

• Channel Count: Both methods convert images to RGB (3 channels) to match your model's input_shape=[45,45,3]. If your original images are grayscale, you could modify the code to load as single-channel and expand dimensions (e.g., x_train = np.expand_dims(x_train, axis=-1)), but you'd need to adjust your model's input shape to (45,45,1) too.
• Label Format: If you prefer one-hot encoded labels (like your original class_mode='categorical'), use tf.keras.utils.to_categorical(y_train, num_classes=82) and switch the loss function to categorical_crossentropy.
• Data Augmentation: If you want to keep using augmentation (shear, zoom, flip), you can apply it to the numpy arrays using ImageDataGenerator.flow() after loading, or use TensorFlow's tf.keras.layers.Random* augmentation layers directly in your model (recommended for better performance).

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