How to generate multi class segmentation image from trained model?

[ksriharshitha]

I have trained the model with 8 classes. But during prediction, I'm not able to plot the ground truth and predictions.

[JordanMakesMaps]

" I'm not able to plot the ground truth and predictions."

What do you mean by plot? Like literally show the image using matplotlib? The predictions should be coming out in the shape of (batch size height, width, number of classes), so for you to visualize the prediction you need to use something like numpy.argmax(prediction, axis = 2) since we want an array that has only a single channel as opposed to 8. If you used to_categorical() then the ground truth will need to be treated the same way.

[JonnoFTW]

If you have one-hot encoded masks like this in shape (rows, cols, classes) you can make an image with this function:

def mask2img(mask):
    palette = {
        0: (0, 0, 0),
        1: (255, 0, 0),
        2: (0, 255, 0),
        3: (0, 0, 255),
        4: (0, 255, 255),
        .... # repeat for all your classes
    }
    rows = mask.shape[0]
    cols = mask.shape[1]
    image = np.zeros((rows, cols, 3), dtype=np.uint8)
    for j in range(rows):
        for i in range(cols):
            image[j, i] = palette[np.argmax(mask[j, i])]
    return image