如何在Keras中使用flow_from_directory进行多类语义分割?

时间:2019-08-14 00:48:44

标签: python keras

假设我有100个训练灰度图像和100个RGB训练蒙版,每个蒙版的大小均为512x512。我可以在下面的Keras中使用to_categorical对掩膜进行一次热编码

numclasses=3
masks_one_hot=to_categorical(maskArr,numclasses)

其中maskArr是100x512x512x1,而masks_one_hot是100x512x512x3。

但是,要使用https://github.com/zhixuhao/unet/blob/master/data.py中的ImageDataGeneratorflow_from_directory来使用trainGeneratortrainGenerator,我尝试保存一键编码的训练图像,然后使用{{1 }}。但是,我注意到在对它们使用imwrite然后用imread读取它们之后,它们从单热编码的512x512x3变为512x512x3 RGB图像。也就是说,现在不是每个通道的值为0或1,它们的范围是0-255

结果,如果我这样做:

myGenerator = trainGeneratorOneHot(20,'data/membrane/train','image','label',data_gen_args,save_to_dir = "data/membrane/train/aug", flag_multi_class = True,
num_class = 3, target_size=(512,512,3))

num_batch=3
for i,batch in enumerate(myGenerator):
    if(i >= num_batch):
        break

其中trainGeneratorOneHot在下面:

def trainGeneratorOneHot(batch_size,...class_mode=None, image_class_mode=None):

    image_datagen = ImageDataGenerator(**aug_dict)
    mask_datagen = ImageDataGenerator(**aug_dict)
    image_generator = image_datagen.flow_from_directory(train_path,classes = [image_folder], class_mode = image_class_mode, color_mode = image_color_mode,target_size = target_size, ...)
    mask_generator = mask_datagen.flow_from_directory(train_path, classes = [mask_folder], class_mode = class_mode, target_size = target_size,...)
    train_generator = zip(image_generator, mask_generator)

    for (img,mask) in train_generator:
        img,mask = adjustDataOneHot(img,mask)
        yield (img,mask)

def adjustDataOneHot(img,mask):
    return (img,mask)

然后我得到`ValueError:无法将输入数组从形状(512,512,1)广播到形状(512,512,3,1)

我该如何解决?

1 个答案:

答案 0 :(得分:0)

几天前正在处理同一问题。我发现使自己的数据生成器类处理从数据帧中获取数据,对其进行扩充并在将其传递给我的模型之前进行一次热编码非常必要。我从来没有让Keras ImageDataGenerator来解决多个类的语义分割问题。

下面是一个数据生成器类,以防可能对您有所帮助:

def one_hot_encoder(mask, num_classes = 8):

    hot_mask = np.zeros(shape = mask.shape, dtype = 'uint8')

    for _ in range(8):
        temp = np.zeros(shape = mask.shape[0:2], dtype = 'uint8')
        temp[mask[:, :, _] != 0] = 1
        hot_mask[:, :, _] = temp

    return hot_mask

# Image data generator class
class DataGenerator(keras.utils.Sequence):

    def __init__(self, dataframe, batch_size, n_classes = 8, augment = False):
        self.dataframe = dataframe
        self.batch_size = batch_size
        self.n_classes = n_classes
        self.augment = augment


    # Steps per epoch    
    def __len__(self):
        return len(self.dataframe) // self.batch_size

    # Shuffles and resets the index at the end of training epoch
    def on_epoch_end(self):
        self.dataframe = self.dataframe.reset_index(drop = True)


    # Generates data, feeds to training
    def __getitem__(self, index):

        processed_images = []
        processed_masks = []

        for _ in range(self.batch_size):

            the_image = io.imread(self.dataframe['Images'][index])
            the_mask = io.imread(self.dataframe['Masks'][index]).astype('uint8');
            one_hot_mask = one_hot_encoder(the_mask, 8)


            if(self.augment):
                # Resizing followed by some augmentations
                processed_image = augs_for_images(image = the_image) / 255.0
                processed_mask = augs_for_masks(image = one_hot_mask)


            else:
                # Still resizing but no augmentations   
                processed_image = resize(image = the_image) / 255.0
                processed_mask = resize(image = one_hot_mask)

            processed_images.append(processed_image)
            processed_masks.append(processed_mask)


        batch_x = np.array( processed_images )
        batch_y = np.array( processed_masks )

        return (batch_x, batch_y)

此外,这是到仓库的链接,其中包含一些您可能感兴趣的语义细分模型。 notebook本身说明了作者如何处理多类语义分割。