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

Question

假设我有 100 个训练灰度图像和 100 个 RGB 训练掩码，每个大小为 512x512。我能够在 Keras 中使用 to_categorical 对掩码进行一次热编码

numclasses=3
masks_one_hot=to_categorical(maskArr,numclasses)

其中maskArr 是 100x512x512x1，masks_one_hot 是 100x512x512x3。

但是，为了使用 ImageDataGenerator 和 flow_from_directory 使用来自 https://github.com/zhixuhao/unet/blob/master/data.py 的 trainGenerator，我尝试保存 one-hot 编码的训练图像，然后使用 trainGenerator 读取它们。然而，我注意到在它们上使用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: could not broadcast input array from shape (512,512,1) into shape (512,512,3,1)

我该如何解决这个问题？

Answer 1

几天前正在处理同样的问题。我发现创建自己的数据生成器类来处理从数据帧中获取数据、对其进行扩充、然后在将其传递给我的模型之前对其进行单热编码非常重要。我永远无法让 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 本身就展示了作者是如何处理多类语义分割的。