如何为Tensor专门创建图像的分类掩码?或者将NumPy函数正确地移植到Dataset.map函数

我正试图从NumPy数组作为我的数据集转移到tensorflow.dataset。

现在,我创建了一个管道来训练分类问题的模型。在某个时刻,我只是使用map函数对所有图像进行归一化:

dataset['train'] = dataset['train'].map(pre_pr, num_parallel_calls=tf.data.experimental.AUTOTUNE)

功能描述如下:

@tf.function
def normalize(input_image: tf.Tensor, input_mask: tf.Tensor) -> tuple:
    input_image = tf.cast(input_image, tf.float32) / 255.0
    input_mask= tf.cast(input_mask, tf.float32) / 255.0
    return input_image, input_mask

@tf.function
def pre_pr(datapoint: dict) -> tuple:
    input_image = tf.image.resize(datapoint['image'], (IMG_SIZE, IMG_SIZE))
    input_mask = tf.image.resize(datapoint['y_mask'], (IMG_SIZE, IMG_SIZE))
    return normalize(input_image, input_mask)

这很好用。

在所有映射之后,它转换 <MapDataset shapes: {image: (None, None, 3), y_mask: (None, None, 3)}, types: {image: tf.uint8, y_mask: tf.uint8}> 到 <PrefetchDataset shapes: ((None, 128, 128, 3), (None, 128, 128, 3)), types: (tf.float32, tf.float32)>

但问题在于任何分割问题,当我试图创建分类掩码时,它会引起问题。

我有一个函数,可以用给定的调色板从图像中创建分割掩码,如下所示:

def flat_labels(label):
    label_seg = np.zeros(label.shape, dtype=np.float32)
    for i in range(len(encodings)):
        label_seg[np.all(label == encodings[i], axis=-1)] = i
    return label_seg[:, :, 0]

这里的编码是一个列表/数组,类似于 [255,255,0],[0,0,255]...

当我试着这样称呼它时:

input_mask = tf.numpy_function(flat_labels, [input_mask], tf.uint8)

它将数据类型更改为其他类型:

<PrefetchDataset shapes: ((None, 128, 128, 3), <unknown>), types: (tf.float32, tf.uint8)>

我真的不想更改其余的代码,我只想调用 flat_labels 基于NumPy的函数到张量。或者以某种方式创建一个函数,为张量创建一个分类掩码。

提前谢谢。