File size: 686 Bytes
41eee5f |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 |
from PIL import Image
from torch.utils.data import Dataset
# We have to make a custom dataset class to load them with the torch DataLoader
# Custom dataset class for CIFAR-10 images
class CustomCIFAR10Dataset(Dataset):
def __init__(self, images, labels, transform=None):
self.images = images
self.labels = labels
self.transform = transform
def __len__(self):
return len(self.images)
def __getitem__(self, index):
image = self.images[index]
label = self.labels[index]
# Apply the transformations (if any)
if self.transform is not None:
image = self.transform(image)
return image, label
|