What is the simplest way to change class_to_idx attribute?

Hello everyone, i’m defining my data and dataloader in the currently way:

	data = {
		'train':
		datasets.ImageFolder(root=train_path, transform=image_transforms['train'], ),
		'valid':
		datasets.ImageFolder(root=valid_path, transform=image_transforms['valid']),
	}
	
	dataloaders = {
		'train': DataLoader(data['train'], batch_size=batch_size, shuffle=True),
		'val': DataLoader(data['valid'], batch_size=batch_size, shuffle=True)
	}

In my problem i have four classes (A, AB, AC, ABC). When I print data['train'].class_to_idx i get the following output:

{'A': 0, 'AB': 1, 'ABC': 2, 'AC': 3}

All that i wanna to do is change index of AC to 2, and ABC to 3. and get the following result

{'A': 0, 'AB': 1, 'ABC': 3, 'AC': 2}

What is the simplest way to do that, and what other changes i have to do along my code?

Thanks :smiley:

I would recommend to create a custom Dataset and maybe just reuse some parts of the ImageFolder dataset from here.
Changing this attribute after the ImageFolder was created seems to be wrong, since class_to_idx will be used to create the dataset as seen here.

2 Likes

Would you like to share your solution?