从文件夹读取多图并保存为训练用数据集技术求助

Got it, let's fix that slow image loading issue once and for all! Saving your preprocessed images and one-hot labels into a reusable dataset format (just like MNIST) will cut down your startup time drastically. Here are a few practical, developer-friendly approaches you can implement right away:

1. Simple & Fast: NumPy .npz Format #

This is the easiest option for small to medium-sized datasets. NumPy's savez lets you bundle multiple arrays into a single compressed file, which loads in milliseconds.

Step-by-Step Code: #

First, convert your loaded images and labels into NumPy arrays (assuming you already have them stored in lists like image_list and onehot_label_list):

import numpy as np

# Convert lists to NumPy arrays
images = np.array(image_list)  # Shape: (num_samples, height, width, channels) or (num_samples, height, width) for grayscale
labels = np.array(onehot_label_list)  # Shape: (num_samples, num_classes)

# Save to a single .npz file
np.savez("my_custom_dataset.npz", images=images, labels=labels)

Loading Later: #

# Load the saved dataset
dataset = np.load("my_custom_dataset.npz")
loaded_images = dataset["images"]
loaded_labels = dataset["labels"]

# Now you can use these arrays directly for training!

2. Flexible: Pickle for Python Objects #

If you need to save additional metadata (like class names, image paths, etc.) along with your data, Pickle is a great choice. It serializes Python objects so you can store entire dictionaries or custom classes.

Saving Code: #

import pickle

# Create a dictionary to hold your data and metadata
dataset_dict = {
    "images": np.array(image_list),
    "labels": np.array(onehot_label_list),
    "class_names": ["cat", "dog", "bird"]  # Example metadata
}

# Save to a .pkl file
with open("my_custom_dataset.pkl", "wb") as f:
    pickle.dump(dataset_dict, f, protocol=4)  # Protocol 4 is efficient for large objects

Loading Later: #

with open("my_custom_dataset.pkl", "rb") as f:
    loaded_data = pickle.load(f)

loaded_images = loaded_data["images"]
loaded_labels = loaded_data["labels"]
class_names = loaded_data["class_names"]

3. Scalable: TFRecord (for TensorFlow Users) #

If you're working with TensorFlow and have a large dataset, TFRecord is the industry standard. It's optimized for fast streaming during training and works seamlessly with tf.data.Dataset.

Saving Code: #

import tensorflow as tf

def _bytes_feature(value):
    return tf.train.Feature(bytes_list=tf.train.BytesList(value=[value]))

def _float_feature(value):
    return tf.train.Feature(float_list=tf.train.FloatList(value=value))

# Write to TFRecord file
with tf.io.TFRecordWriter("my_dataset.tfrecord") as writer:
    for img, lbl in zip(image_list, onehot_label_list):
        # Convert image to bytes
        img_bytes = tf.io.serialize_tensor(img).numpy()
        # Convert label to a flat list of floats
        lbl_flat = lbl.flatten().tolist()
        
        feature = {
            "image": _bytes_feature(img_bytes),
            "label": _float_feature(lbl_flat)
        }
        example = tf.train.Example(features=tf.train.Features(feature=feature))
        writer.write(example.SerializeToString())

Loading Later: #

def parse_example(example_proto):
    feature_description = {
        "image": tf.io.FixedLenFeature([], tf.string),
        "label": tf.io.FixedLenFeature([num_classes], tf.float32)  # Replace num_classes with your number of classes
    }
    parsed_features = tf.io.parse_single_example(example_proto, feature_description)
    # Deserialize image back to array
    image = tf.io.parse_tensor(parsed_features["image"], out_type=tf.float32)  # Adjust dtype as needed
    image = tf.reshape(image, (height, width, channels))  # Replace with your image dimensions
    return image, parsed_features["label"]

# Load and prepare the dataset
dataset = tf.data.TFRecordDataset("my_dataset.tfrecord")
dataset = dataset.map(parse_example).batch(32)  # Add batching, shuffling, etc. as needed

4. PyTorch-Friendly: Custom Dataset + Torch Save #

If you're using PyTorch, you can create a custom Dataset class and save it, or use torch.save to store your tensors directly.

Saving Code: #

import torch

# Convert to PyTorch tensors
images_tensor = torch.tensor(np.array(image_list))
labels_tensor = torch.tensor(np.array(onehot_label_list))

# Save tensors
torch.save({"images": images_tensor, "labels": labels_tensor}, "my_dataset.pt")

Loading Later: #

loaded_data = torch.load("my_dataset.pt")
loaded_images = loaded_data["images"]
loaded_labels = loaded_data["labels"]

# Create a custom Dataset if needed
from torch.utils.data import Dataset, DataLoader

class CustomImageDataset(Dataset):
    def __init__(self, images, labels):
        self.images = images
        self.labels = labels
    
    def __len__(self):
        return len(self.images)
    
    def __getitem__(self, idx):
        return self.images[idx], self.labels[idx]

dataset = CustomImageDataset(loaded_images, loaded_labels)
dataloader = DataLoader(dataset, batch_size=32, shuffle=True)

Pro Tips: #

• Preprocess First: Make sure you resize, normalize, or apply any other preprocessing before saving—this way you don't have to redo it every time you load the dataset.
• Check Dtypes: Use uint8 for raw images (saves space) or float32 if you've normalized them. For labels, float32 works fine for one-hot encoding.
• Large Datasets: If your dataset is too big to fit in memory, go with TFRecord or HDF5 (using h5py library) to enable streaming loading without hogging RAM.

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