sal:

For semantic segmentation outputs how do we visualize the output feature map, a tensor of shape <B x Number_of_Classes x H x W> to <B X 3 X H X W> given a color map palette corresponding to  each of the class :

labels = ['unlabeled',	'ego vehicle',	'rectification border',	'out of roi',	'static',	'dynamic',	'ground',	'road',	'sidewalk',	'parking',	'rail track',	'building',	'wall',	'fence',	'guard rail',	'bridge',	'tunnel',	'pole',	'polegroup',	'traffic light',	'traffic sign',	'vegetation',	'terrain',	'sky',	'person',	'rider',	'car',	'truck',	'bus',	'caravan',	'trailer',	'train',	'motorcycle',	'bicycle',	'license plate']
    cityscapes_map = np.array([[0.        , 0.        , 0.        ],
       [0.        , 0.        , 0.        ],
       [0.        , 0.        , 0.        ],
       [0.        , 0.        , 0.        ],
       [0.07843137, 0.07843137, 0.07843137],
       [0.43529412, 0.29019608, 0.        ],
       [0.31764706, 0.        , 0.31764706],
       [0.50196078, 0.25098039, 0.50196078],
       [0.95686275, 0.1372549 , 0.90980392],
       [0.98039216, 0.66666667, 0.62745098],
       [0.90196078, 0.58823529, 0.54901961],
       [0.2745098 , 0.2745098 , 0.2745098 ],
       [0.4       , 0.4       , 0.61176471],
       [0.74509804, 0.6       , 0.6       ],
       [0.70588235, 0.64705882, 0.70588235],
       [0.58823529, 0.39215686, 0.39215686],
       [0.58823529, 0.47058824, 0.35294118],
       [0.6       , 0.6       , 0.6       ],
       [0.6       , 0.6       , 0.6       ],
       [0.98039216, 0.66666667, 0.11764706],
       [0.8627451 , 0.8627451 , 0.        ],
       [0.41960784, 0.55686275, 0.1372549 ],
       [0.59607843, 0.98431373, 0.59607843],
       [0.2745098 , 0.50980392, 0.70588235],
       [0.8627451 , 0.07843137, 0.23529412],
       [1.        , 0.        , 0.        ],
       [0.        , 0.        , 0.55686275],
       [0.        , 0.        , 0.2745098 ],
       [0.        , 0.23529412, 0.39215686],
       [0.        , 0.        , 0.35294118],
       [0.        , 0.        , 0.43137255],
       [0.        , 0.31372549, 0.39215686],
       [0.        , 0.        , 0.90196078],
       [0.46666667, 0.04313725, 0.1254902 ],
       [0.        , 0.        , 0.55686275]])

In this case output classes = 35  and I have list of 35 class names and a matrix of 35X3 values that I assume are the r,g,b values of pixels for each class.
Output of my final classification layer is <B x 35 x H x W>
It is a conventional image segmentation model with (multiple *conv2d) -> (multiple * convTranspose2d) -> 1x1 convolution (having filters = number of classes)
So naively I could write a function to loop over the 35 filters and for each 2d matrix save those values corresponding to the pixels in the output feature map with corresponding r,g,b values from the palette to a single tensor.
But how do I know which filter index matches from the original class list ?
Is there a better way to do this.
In general how are such outputs visualized ?


Semantic segmentation works by learning the color as the label for a class.
The 2d array is number_of_classes X 3
That is an unique rgb value for each class.
But since your objective function is a classification loss over multiple classes you can not use an rgb value as the label itself.
Thus you use the index of the same.
The output feature map after your classification layer gives probabilities.
That is every pixel value is the probability if it belongs to that class.
And you have N channels where N = number of classes.
Thus for each class you get a 2d  feature map of probabilities.
Thus taking a max gives which pixel belongs to which index.
Then you need to map out these index values back the original corresponding rgb dictionary to visualize.
An analogy would be say a word classifier where you have a fixed vocabulary and the label classified is the index of word in vocab. To infer the actual output from the class probabilities you map it back from your vocab.
Hope this helps !