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
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)