如何从语义图像分割输出中生成遮罩？

Question

我最近开始为细分问题创建U-Net。但是，经过训练后，遮罩仅位于一个通道（400,400,1）中，而不是三通道彩色图像（400,400,3）。

我贴了我要指的口罩。左蒙版是我的输入蒙版，右蒙版是我的模型输出的蒙版。

图像链接（无信誉）：https://i.stack.imgur.com/wCQlD.png

要获取输出蒙版，请使用Tensorflow教程进行图像分割。代码：

def create_mask1(pred_mask):
    pred_mask = tf.argmax(pred_mask, axis=-1)
    pred_mask = pred_mask[..., tf.newaxis]
    return pred_mask[0]
....
plt.imshow(create_mask1(segmentation_model.predict(image)))

我还有另一种生成蒙版的方法吗？如果还有另一种方法，我是否必须重新训练模型以使用该方法？