如何使用带有多标签掩码的 ImageDataGenerator 进行多类图像分割？

Question

为了进行多类分割，需要对掩码进行单热编码。例如，如果我有 100 张形状为 224x224x3 且有 5 个不同类别的图像，我将拥有一组形状为 (100, 224, 224, 5) 的掩码，即最后一个维度（通道）是指像素的类别。取一个包含 6 个类的灰度蒙版，其中每个像素的标签为 1-6，我可以使用 tf.keras.utils.to_categorical 轻松将其转换为我需要的分类蒙版。

如果我使用与 keras 一起提供的 ImageDataGenerator，我知道我可以为图像和蒙版创建一个生成器，然后将它们压缩在一起以解决问题（如下代码所示），但我感到困惑的是如何将蒙版转换为这种分类蒙版- 使用 ImageDataGenerator 时的热编码结构？ImageDataGenerator 只在保存为图像的目录中找到文件，因此我无法转换掩码，然后将它们保存为 numpy 数组（单热编码掩码）供生成器拾取，因为图像不能有有 4 个以上的频道吧？是否有某种方式告诉生成器进行这种转换？或者这是否因此限制了我在问题中可以拥有的课程数量？

一种解决方案是使用我已经完成的序列类编写我自己的自定义生成器，但我很想知道这是否可以与 Keras 内置的 ImageDataGenerator 一起使用？在网络上编写我的 lambda 层是解决方案吗？

mask_categorical = tf.keras.utils.to_categoricl(mask) #converts 224x224 grayscale mask to one-hot encoding version



imgDataGen = ImageDataGenerator(rescale=1/255.)

maskDataGen = ImageDataGenerator()

imageGenerator =imageDataGen.flow_from_directory("dataset/image/",
                                                class_mode=None, seed=40)

maskGenerator = maskDataGen.flow_from_directory("dataset/mask/",
                                               class_mode=None, seed=40)

trainGenerator = zip(imageGenerator, maskGenerator)