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
20import torchvision.transforms.functional
21from PIL import Image
23from labml import lab

Carvana Dataset

26class CarvanaDataset(
  • 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())


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 =[id_])

Transform image and convert it to a PyTorch tensor

62        image = self.transforms(image)

Load mask

64        mask =[id_])

Transform mask and convert it to a PyTorch tensor

66        mask = self.transforms(mask)

The mask values were not , 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')