How to visualize segmentation output - multiclass feature map to rgb image?

vision
2

You could get the class predictions using torch.argmax(output, dim=1) as this will get the argmax of your class dimension.
Then you could use this tensor to index directly into your colormap and visualize the images in a lib like matplotlib:

batch_size = 4
nb_classes = len(labels)
h, w = 96, 96

x = torch.randn(batch_size, nb_classes, h, w)
pred = torch.argmax(x, dim=1)

pred_imgs = [cityscapes_map[p] for p in pred]

for pred_img in pred_imgs:
    plt.imshow(pred_img)
    plt.show()