如何使用Python加载图像文件夹并转为training dataset

Hey there! Let's walk through how to load your image dataset folder into a training-ready dataset using Python. I'll cover a few common approaches depending on what framework you're using—whether you're working with PyTorch, TensorFlow, or just need a flexible manual setup.

First, let's assume your dataset follows the standard classification task structure (the most common setup for training image models):

your_dataset_folder/
├── cat/
│   ├── cat1.jpg
│   ├── cat2.png
│   └── ...
├── dog/
│   ├── dogA.jpeg
│   ├── dogB.png
│   └── ...
└── ...

If your folder structure is different (like having a separate label file), feel free to drop a comment with details and I'll adjust the solution.

PyTorch's torchvision.datasets.ImageFolder is built exactly for this use case—it automatically maps folder names to class labels and works seamlessly with data loaders for batching/shuffling.

First, install dependencies if you haven't already:

pip install torch torchvision

Here's the code:

import torch
from torchvision import datasets, transforms
from torch.utils.data import DataLoader

# Define image preprocessing (tweak based on your model's requirements)
train_transform = transforms.Compose([
    transforms.Resize((224, 224)),  # Resize all images to 224x224
    transforms.RandomHorizontalFlip(),  # Data augmentation: random horizontal flip
    transforms.ToTensor(),  # Convert to PyTorch tensor
    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])  # Normalize with ImageNet stats
])

# Load the dataset
train_dataset = datasets.ImageFolder(
    root='path/to/your_dataset_folder',  # Replace with your actual folder path
    transform=train_transform
)

# Check class mappings
print("Class list:", train_dataset.classes)
print("Class-to-index mapping:", train_dataset.class_to_idx)

# Convert to DataLoader for training batches
train_loader = DataLoader(
    train_dataset,
    batch_size=32,  # Load 32 images per batch
    shuffle=True,  # Shuffle data for better training
    num_workers=4  # Use 4 processes to speed up loading
)

# Test a single batch
for images, labels in train_loader:
    print(f"Batch shape: {images.shape}")  # Should be (32, 3, 224, 224)
    print(f"Corresponding labels: {labels}")
    break

TensorFlow's tf.keras.utils.image_dataset_from_directory is another zero-fuss option, with built-in support for data augmentation and preprocessing.

Install dependencies if needed:

pip install tensorflow

Code example:

import tensorflow as tf
from tensorflow.keras import layers

# Load the dataset directly from folders
train_dataset = tf.keras.utils.image_dataset_from_directory(
    'path/to/your_dataset_folder',
    image_size=(224, 224),  # Standardize image size
    batch_size=32,
    shuffle=True,
    seed=42  # Fix random seed for reproducibility
)

# Get class names
class_names = train_dataset.class_names
print("Class list:", class_names)

# Optional: Add data augmentation to boost model generalization
data_augmentation = tf.keras.Sequential([
    layers.RandomFlip("horizontal"),
    layers.RandomRotation(0.1),
    layers.RandomZoom(0.1)
])

# Apply augmentation and normalize pixel values (0-255 → 0-1)
train_dataset = train_dataset.map(
    lambda x, y: (data_augmentation(x, training=True), y)
)
train_dataset = train_dataset.map(
    lambda x, y: (x / 255.0, y)
)

# Test a batch
for images, labels in train_dataset.take(1):
    print(f"Batch shape: {images.shape}")
    print(f"Corresponding labels: {labels}")

If you need custom logic (like handling non-standard labels or image formats), use PIL or OpenCV to manually traverse the folder.

Install dependencies:

pip install pillow opencv-python

Code with PIL:

import os
from PIL import Image
import numpy as np

# Set dataset path and map classes to indexes
dataset_path = 'path/to/your_dataset_folder'
class_names = sorted(os.listdir(dataset_path))
class_to_idx = {cls: idx for idx, cls in enumerate(class_names)}

# Store images and labels
images = []
labels = []

# Traverse each class folder
for cls_name in class_names:
    cls_path = os.path.join(dataset_path, cls_name)
    for img_name in os.listdir(cls_path):
        img_path = os.path.join(cls_path, img_name)
        try:
            # Load and process image
            img = Image.open(img_path).convert('RGB')  # Ensure RGB format
            img = img.resize((224, 224))  # Standardize size
            img_array = np.array(img) / 255.0  # Normalize pixel values
            images.append(img_array)
            labels.append(class_to_idx[cls_name])
        except Exception as e:
            print(f"Failed to load {img_path}: {str(e)}")

# Convert to numpy arrays (easily convertible to tensors later)
images = np.array(images)
labels = np.array(labels)

print(f"Loaded {len(images)} images across {len(class_names)} classes")
print(f"Image array shape: {images.shape}")

If you run into any issues—like weird image formats, non-standard folder structures, or need help adapting this to your specific model—just let me know!

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