Data Loading in Pytorch for a dataset having all the classes in same folder
I am new to deep learning and Pytorch. I have data set of 6000 images that have all four classes in a single folder. I used the following snippet to upload my data.


torchvision.datasets.ImageFolder(root='/content/drive/My Drive/DFU/base_dir/train_dir', transform=None)



I read that for ImageFolder, the images should be organized into sub-folders based on class labels. However, my dataset has all four class images in a single folder. I have a .csv file that contains the one-hot-encoded class label for each image. How to load my dataset to Pytorch?


The simplest solution would be to reorganise the images into class-subfolders based on the csv file, and load as intended by ImageFolder:
import pandas as pd
from pathlib import Path

root = '/content/drive/My Drive/DFU/base_dir/train_dir'
my_csv_file = ...

# Loading csv as {image:class,...} format
df = pd.read_csv(my_csv_file).set_index('images')
class_dict = df.idxmax(axis=&quot;columns&quot;).to_dict()

# Moving files to class-named subfolders
for path in Path(root).iterdir():
    if path.is_file() and path.name in class_dict.keys():
        path.rename(Path(path.parent, class_dict[path.name], path.name)

# Loading dataset
dataset = torchvision.datasets.ImageFolder(root=root, transform=None)

