#

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 pathlib import Path
17
18import torchvision.transforms.functional
19from PIL import Image
20
21import torch.utils.data
22from labml import lab

#

Carvana Dataset

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

#

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

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

#

Get a dictionary of images by id

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

#

Get a dictionary of masks by id

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

#

Image ids list

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

#

Transformations

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

#

Get an image and its mask.

idx is index of the image

49    def __getitem__(self, idx: int):

#

Get image id

57        id_ = self.ids[idx]

#

Load image

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

#

Transform image and convert it to a PyTorch tensor

61        image = self.transforms(image)

#

Load mask

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

#

Transform mask and convert it to a PyTorch tensor

65        mask = self.transforms(mask)

#

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

68        mask = mask / mask.max()

#

Return the image and the mask

71        return image, mask

#

Size of the dataset

73    def __len__(self):

#

77        return len(self.ids)

#

Testing code

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