#

Carvana Dataset for the U-Net experiment

You can find the download instructions on Kaggle.

Save the training images inside carvana/train folder and the masks in carvana/train_masks folder.

16from torch import nn
17from pathlib import Path
18
19import torch.utils.data
20import torchvision.transforms.functional
21from PIL import Image
22
23from labml import lab

#

Carvana Dataset

26class CarvanaDataset(torch.utils.data.Dataset):

#

image_path is the path to the images
mask_path is the path to the masks

31    def __init__(self, image_path: Path, mask_path: Path):

#

Get a dictionary of images by id

37        self.images = {p.stem: p for p in image_path.iterdir()}

#

Get a dictionary of masks by id

39        self.masks = {p.stem[:-5]: p for p in mask_path.iterdir()}

#

Image ids list

42        self.ids = list(self.images.keys())

#

Transformations

45        self.transforms = torchvision.transforms.Compose([
46            torchvision.transforms.Resize(572),
47            torchvision.transforms.ToTensor(),
48        ])

#

Get an image and its mask.

idx is index of the image

50    def __getitem__(self, idx: int):

#

Get image id

58        id_ = self.ids[idx]

#

Load image

60        image = Image.open(self.images[id_])

#

Transform image and convert it to a PyTorch tensor

62        image = self.transforms(image)

#

Load mask

64        mask = Image.open(self.masks[id_])

#

Transform mask and convert it to a PyTorch tensor

66        mask = self.transforms(mask)

#

The mask values were not $1$ , so we scale it appropriately.

69        mask = mask / mask.max()

#

Return the image and the mask

72        return image, mask

#

Size of the dataset

74    def __len__(self):

#

78        return len(self.ids)

#

Testing code

82if __name__ == '__main__':
83    ds = CarvanaDataset(lab.get_data_path() / 'carvana' / 'train', lab.get_data_path() / 'carvana' / 'train_masks')