data_loader.py

import numpy as np
import cv2
import os
import random
import glob

from keras.utils import Sequence
from keras.applications.imagenet_utils import preprocess_input as pinput


class HairGenerator(Sequence):

    def __init__(self,
                 transformer,
                 root_dir,
                 mode='Training',
                 nb_classes=3,
                 batch_size=4,
                 backbone=None,
                 shuffle=False):

        # backbone fit for segmentation_models，have been deleted now...
        assert mode in ['Training', 'Testing'], "Data set selection error..."

        self.image_path_list = sorted(
            glob.glob(os.path.join(root_dir, 'Original', mode, '*')))
        self.mask_path_list = sorted(
            glob.glob(os.path.join(root_dir, 'GT', mode, '*')))
        self.transformer = transformer
        self.batch_size = batch_size
        self.nb_classes = nb_classes
        self.shuffle = shuffle
        self.mode = mode
        self.backbone = backbone

    def __getitem__(self, idx):
        images, masks = [], []

        for (image_path, mask_path) in zip(self.image_path_list[idx * self.batch_size: (idx + 1) * self.batch_size],
                                           self.mask_path_list[idx * self.batch_size: (idx + 1) * self.batch_size]):
            image = cv2.imread(image_path, 1)
            mask = cv2.imread(mask_path, 0)

            image = self._padding(image)
            mask = self._padding(mask)

            # augumentation
            augmentation = self.transformer(image=image, mask=mask)
            image = augmentation['image']
            mask = self._get_result_map(augmentation['mask'])

            images.append(image)
            masks.append(mask)

        images = np.array(images)
        masks = np.array(masks)
        images = pinput(images)

        return images, masks

    def __len__(self):
        """Steps required per epoch"""
        return len(self.image_path_list) // self.batch_size

    def _padding(self, image):
        shape = image.shape
        h, w = shape[:2]
        width = np.max([h, w])
        padd_h = (width - h) // 2
        padd_w = (width - w) // 2
        if len(shape) == 3:
            padd_tuple = ((padd_h, width - h - padd_h),
                          (padd_w, width - w - padd_w), (0, 0))
        else:
            padd_tuple = ((padd_h, width - h - padd_h), (padd_w, width - w - padd_w))
        image = np.pad(image, padd_tuple, 'constant')
        return image

    def on_epoch_end(self):
        """Shuffle image order"""
        if self.shuffle:
            c = list(zip(self.image_path_list, self.mask_path_list))
            random.shuffle(c)
            self.image_path_list, self.mask_path_list = zip(*c)

    def _get_result_map(self, mask):
        """Processing mask data"""

        # mask.shape[0]: row, mask.shape[1]: column
        result_map = np.zeros((mask.shape[1], mask.shape[0], self.nb_classes))
        # 0 (background pixel), 128 (face area pixel) or 255 (hair area pixel).
        skin = (mask == 128)
        hair = (mask == 255)

        if self.nb_classes == 2:
            # hair = (mask > 128)
            background = np.logical_not(hair)
            result_map[:, :, 0] = np.where(background, 1, 0)
            result_map[:, :, 1] = np.where(hair, 1, 0)
        elif self.nb_classes == 3:
            background = np.logical_not(hair + skin)
            result_map[:, :, 0] = np.where(background, 1, 0)
            result_map[:, :, 1] = np.where(skin, 1, 0)
            result_map[:, :, 2] = np.where(hair, 1, 0)
        else:
            raise Exception("error...")

        return result_map