Loading images using file name as label

data
1

Hi, I’m new to machine learning and looking to load a dataset from a file structure that looks something like:

root
-star
–img1.png
– img2.png
– img3.png
…

-planet
–img20.png
–img21.png
…

Every tutorial I find uses a preconstructed dataset and isn’t very helpful, so I would appreciate a pointer

2

Hi. I think something like this would do:

import os
from torchvision import datasets, transforms

def load_image_dataset(root_dir, image_size=224):
    """
    Load image dataset from a directory structure.
    
    Args:
        root_dir (str): Path to root directory containing class subdirectories
        image_size (int): Resize images to this dimension
    
    Returns:
        torchvision.datasets.ImageFolder: Loaded dataset
    """
    # Define image transformations (data augmentation)
    data_transforms = transforms.Compose([
        transforms.Resize((image_size, image_size)),
        transforms.ToTensor(),
        transforms.Normalize(mean=[0.485, 0.456, 0.406], 
                             std=[0.229, 0.224, 0.225])
    ])
    
    # Load dataset
    train_data = datasets.ImageFolder(
        root=root_dir, 
        transform=data_transforms
    )
    
    return train_data

Then a example of usage, should look like this:

# You can choose the image size you want
train_data = load_image_dataset('path/to/root/directory', image_size=512)