chore: import upstream snapshot with attribution

labml_nn/unet/carvana.py

@@ -0,0 +1,82 @@
+"""
+---
+title: Carvana dataset for the U-Net experiment
+summary: >
+  Carvana dataset for the U-Net experiment.
+---
+
+# Carvana Dataset for the [U-Net](index.html) [experiment](experiment.html)
+
+You can find the download instructions
+[on Kaggle](https://www.kaggle.com/competitions/carvana-image-masking-challenge/data).
+
+Save the training images inside `carvana/train` folder and the masks in `carvana/train_masks` folder.
+"""
+
+from pathlib import Path
+
+import torchvision.transforms.functional
+from PIL import Image
+
+import torch.utils.data
+from labml import lab
+
+
+class CarvanaDataset(torch.utils.data.Dataset):
+    """
+    ## Carvana Dataset
+    """
+
+    def __init__(self, image_path: Path, mask_path: Path):
+        """
+        :param image_path: is the path to the images
+        :param mask_path: is the path to the masks
+        """
+        # Get a dictionary of images by id
+        self.images = {p.stem: p for p in image_path.iterdir()}
+        # Get a dictionary of masks by id
+        self.masks = {p.stem[:-5]: p for p in mask_path.iterdir()}
+
+        # Image ids list
+        self.ids = list(self.images.keys())
+
+        # Transformations
+        self.transforms = torchvision.transforms.Compose([
+            torchvision.transforms.Resize(572),
+            torchvision.transforms.ToTensor(),
+        ])
+
+    def __getitem__(self, idx: int):
+        """
+        #### Get an image and its mask.
+
+        :param idx: is index of the image
+        """
+
+        # Get image id
+        id_ = self.ids[idx]
+        # Load image
+        image = Image.open(self.images[id_])
+        # Transform image and convert it to a PyTorch tensor
+        image = self.transforms(image)
+        # Load mask
+        mask = Image.open(self.masks[id_])
+        # Transform mask and convert it to a PyTorch tensor
+        mask = self.transforms(mask)
+
+        # The mask values were not $1$, so we scale it appropriately.
+        mask = mask / mask.max()
+
+        # Return the image and the mask
+        return image, mask
+
+    def __len__(self):
+        """
+        #### Size of the dataset
+        """
+        return len(self.ids)
+
+
+# Testing code
+if __name__ == '__main__':
+    ds = CarvanaDataset(lab.get_data_path() / 'carvana' / 'train', lab.get_data_path() / 'carvana' / 'train_masks')