def __getitem__(self, index):
    'Generate one batch of data'
    # Generate indexes of the batch
    indexes = self.indexes[index*self.batch_size: (index+1)*self.batch_size]

    # Find list of IDs
    list_IDs_temp = [self.list_IDs[k] for k in indexes]

    # Generate data
    X = self._generate_X(list_IDs_temp)

    if self.to_fit:
        y = self._generate_y(list_IDs_temp)
        return X, y
    else
        return X

def _generate_X(self, list_IDs_temp):
    'Generates data containing batch_size images'
    # Initialization
    X = np.empty((self.batch_size, *self.dim, self.n_channels))

    # Generate data
    for i, ID in enumerate(list_IDs_temp):
        # Store sample
        X[i,] = _load_grayscale_image(self.image_path + self.labels[ID])

    return X

def _generate_y(self, list_IDs_temp):
    'Generates data containing batch_size masks'
    y = np.empty((self.batch_size, *self.dim), dtype=int)

    # Generate data
    for i, ID in enumerate(list_IDs_temp):
        # Store sample
        y[i,] = _load_grayscale_image(self.mask_path + self.labels[ID])

    return y

def _load_grayscale_image(image_path):
    'Load grayscale image'
    img = cv2.imread(image_path)
    img = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
    img = img / 255

    return img