data.py

from __future__ import print_function
from keras.preprocessing.image import ImageDataGenerator
import matplotlib.pyplot as plt
import numpy as np
import os
import glob
import skimage.io as io
import skimage.transform as trans

Sky = [128, 128, 128]
Building = [128, 0, 0]
Pole = [192, 192, 128]
Road = [128, 64, 128]
Pavement = [60, 40, 222]
Tree = [128, 128, 0]
SignSymbol = [192, 128, 128]
Fence = [64, 64, 128]
Car = [64, 0, 128]
Pedestrian = [64, 64, 0]
Bicyclist = [0, 128, 192]
Unlabelled = [0, 0, 0]

COLOR_DICT = np.array([
    Sky,
    Building,
    Pole,
    Road,
    Pavement,
    Tree,
    SignSymbol,
    Fence,
    Car,
    Pedestrian,
    Bicyclist,
    Unlabelled,
])


def adjustData(img, mask, flag_multi_class, num_class):
    if flag_multi_class:
        img = img / 255
        mask = mask[:, :, :, 0] if (len(mask.shape) == 4) else mask[:, :, 0]
        new_mask = np.zeros(mask.shape + (num_class, ))
        for i in range(num_class):
            # for one pixel in the image, find the class in mask and convert it into one-hot vector
            # index = np.where(mask == i)
            # index_mask = (index[0],index[1],index[2],np.zeros(len(index[0]),dtype = np.int64) + i) if (len(mask.shape) == 4) else (index[0],index[1],np.zeros(len(index[0]),dtype = np.int64) + i)
            # new_mask[index_mask] = 1
            new_mask[mask == i, i] = 1
        new_mask = (np.reshape(
            new_mask,
            (
                new_mask.shape[0],
                new_mask.shape[1] * new_mask.shape[2],
                new_mask.shape[3],
            ),
        ) if flag_multi_class else np.reshape(
            new_mask,
            (new_mask.shape[0] * new_mask.shape[1], new_mask.shape[2])))
        mask = new_mask
    elif np.max(img) > 1:
        img = img / 255
        mask = mask / 255
        mask[mask > 0.5] = 1
        mask[mask <= 0.5] = 0
    return (img, mask)


def trainGenerator(
        batch_size,
        train_path,
        image_folder = None,
        mask_folder = None,
        aug_dict = None,
        class_mode = None,
        image_color_mode="rgb",
        mask_color_mode="grayscale",
        image_save_prefix="image",
        mask_save_prefix="mask",
        flag_multi_class=False,
        num_class=2,
        save_to_dir=None,
        target_size=(256, 256),
        seed=1,
):
    """
    can generate image and mask at the same time
    use the same seed for image_datagen and mask_datagen to ensure the transformation for image and mask is the same
    if you want to visualize the results of generator, set save_to_dir = "your path"
    """
    image_datagen = ImageDataGenerator(**aug_dict)
    mask_datagen = ImageDataGenerator(**aug_dict)
    image_generator = image_datagen.flow_from_directory(
        train_path,
        classes=image_folder,
        class_mode=class_mode,
        color_mode=image_color_mode,
        target_size=target_size,
        batch_size=batch_size,
        save_to_dir=save_to_dir,
        save_prefix=image_save_prefix,
        seed=seed,
    )
    mask_generator = mask_datagen.flow_from_directory(
        train_path,
        classes=mask_folder,
        class_mode=class_mode,
        color_mode=mask_color_mode,
        target_size=target_size,
        batch_size=batch_size,
        save_to_dir=save_to_dir,
        save_prefix=mask_save_prefix,
        seed=seed,
    )
    # 只要简单地zip起来，就能生成图像对！
    train_generator = zip(image_generator, mask_generator)
    for (img, mask) in train_generator:
        img, mask = adjustData(img, mask, flag_multi_class, num_class)
        yield (img, mask)


def testGenerator(
    test_path,
    num_image=30,
    target_size=(256, 256),
    flag_multi_class=False,
    as_gray=False,
    centercrop=False,
):
    while 1:
        for filename in glob.glob(os.path.join(test_path, "chips/*.tif")):
            img = io.imread(filename, as_gray=as_gray)
            img = img / 255
            if centercrop:
                pass
            img = trans.resize(img, target_size)
            # img = np.reshape(img, img.shape+(1,)
            #                  ) if (not flag_multi_class) else img
            img = np.reshape(img, (1, ) + img.shape)
            yield img


def indexTestGenerator(
        batch_size,
        train_path,
        image_folder,
        mask_folder,
        nuclei_folder,
        aug_dict,
        image_color_mode="rgb",
        mask_color_mode="grayscale",
        nuclei_color_mode="rgb",
        image_save_prefix="image",
        mask_save_prefix="mask",
        flag_multi_class=False,
        num_class=2,
        save_to_dir=None,
        target_size=(256, 256),
        seed=1,
):
    image_datagen = ImageDataGenerator(**aug_dict)
    mask_datagen = ImageDataGenerator(**aug_dict)
    nuclei_datagen = ImageDataGenerator(**aug_dict)
    image_generator = image_datagen.flow_from_directory(
        train_path,
        classes=[image_folder],
        class_mode=None,
        color_mode=image_color_mode,
        target_size=target_size,
        batch_size=batch_size,
        save_to_dir=save_to_dir,
        save_prefix=image_save_prefix,
        seed=seed,
    )
    mask_generator = mask_datagen.flow_from_directory(
        train_path,
        classes=[mask_folder],
        class_mode=None,
        color_mode=mask_color_mode,
        target_size=target_size,
        batch_size=batch_size,
        save_to_dir=save_to_dir,
        save_prefix=mask_save_prefix,
        seed=seed,
    )
    nuclei_generator = nuclei_datagen.flow_from_directory(
        train_path,
        classes=[nuclei_folder],
        class_mode=None,
        color_mode=nuclei_color_mode,
        target_size=target_size,
        batch_size=batch_size,
        save_to_dir=save_to_dir,
        save_prefix=mask_save_prefix,
        seed=seed,
    )
    train_generator = zip(image_generator, mask_generator, nuclei_generator)
    for (img, mask, nuclei) in train_generator:
        img, mask = adjustData(img, mask, flag_multi_class, num_class)
        if nuclei_color_mode != "rgb":
            img, nuclei = adjustData(img, nuclei, flag_multi_class, num_class)
        else:
            nuclei /= 255
        yield img, mask, nuclei


def geneTrainNpy(
    image_path,
    mask_path,
    flag_multi_class=False,
    num_class=2,
    image_prefix="image",
    mask_prefix="mask",
    image_as_gray=True,
    mask_as_gray=True,
):
    image_name_arr = glob.glob(
        os.path.join(image_path, "%s*.png" % image_prefix))
    image_arr = []
    mask_arr = []
    for index, item in enumerate(image_name_arr):
        img = io.imread(item, as_gray=image_as_gray)
        img = np.reshape(img, img.shape + (1, )) if image_as_gray else img
        mask = io.imread(
            item.replace(image_path,
                         mask_path).replace(image_prefix, mask_prefix),
            as_gray=mask_as_gray,
        )
        mask = np.reshape(mask, mask.shape + (1, )) if mask_as_gray else mask
        img, mask = adjustData(img, mask, flag_multi_class, num_class)
        image_arr.append(img)
        mask_arr.append(mask)
    image_arr = np.array(image_arr)  #  TODO: 这里有输出范围的隐患
    mask_arr = np.array(mask_arr)
    return image_arr, mask_arr


def labelVisualize(num_class, color_dict, img):
    img = img[:, :, 0] if len(img.shape) == 3 else img
    img_out = np.zeros(img.shape + (3, ))
    for i in range(num_class):
        img_out[img == i, :] = color_dict[i]
    return img_out / 255


def saveResult(save_path, npyfile, flag_multi_class=False, num_class=2):
    for i, item in enumerate(npyfile):
        img = (labelVisualize(num_class, COLOR_DICT, item)
               if flag_multi_class else item[:, :, 0])
        plt.imshow(img > 0.5, cmap="gray")
        plt.show()
        io.imsave(os.path.join(save_path, "%d_predict.tif" % i), img)