custom_dataset.py

from torch.utils.data import Dataset as BaseDataset
import cv2
import os
import numpy as np


class CustomDataset(BaseDataset):
    """CustomDataset. Read images, apply augmentation and preprocessing transformations.

    Args:
        images_dir (str): path to images folder
        masks_dir (str): path to segmentation masks folder
        class_values (list): values of classes to extract from segmentation mask
        transforms (albumentations.Compose): data transfromation pipeline
            (e.g. flip, scale, etc.)
        preprocessing (albumentations.Compose): data preprocessing
            (e.g. noralization, shape manipulation, etc.)

    """

    CLASSES = ["background", "foreground"]

    def __init__(
        self,
        images_dir,
        masks_dir,
        classes=None,
        transforms=None,
        preprocessing=None,
        split="train",
    ):
        pre_image_names = os.listdir(images_dir)
        self.ids = [f.split(".")[0] for f in pre_image_names[0:10]]
        self.ids = os.listdir(images_dir)
        self.images_fps = [
            os.path.join(images_dir, image_id + ".jpg") for image_id in self.ids
        ]
        self.masks_fps = [
            os.path.join(masks_dir, image_id + ".png") for image_id in self.ids
        ]

        # convert str names to class values on masks
        self.class_values = [self.CLASSES.index(cls.lower()) for cls in classes]

        self.augmentation = transforms
        self.preprocessing = preprocessing
        self.split = split

    def __getitem__(self, i):

        # read data
        image = cv2.imread(self.images_fps[i])
        image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
        mask = cv2.imread(self.masks_fps[i], 0)

        # extract certain classes from mask (e.g. cars)
        masks = [(mask == 255) for v in self.class_values]
        mask = np.stack(masks, axis=-1).astype("float")

        # apply augmentations
        if self.augmentation:
            sample = self.augmentation(image=image, mask=mask)
            image, mask = sample["image"], sample["mask"]

        # apply preprocessing
        if self.preprocessing:
            sample = self.preprocessing(image=image, mask=mask)
            image, mask = sample["image"], sample["mask"]
        if self.split == "test":
            return image, mask, os.path.basename(self.images_fps[i])
        return image, mask

    def __len__(self):
        return len(self.ids)