Hi,
I have a given cmap for labels in a dataset with 35 classes such as a 35x3 matrix, where each tuple is a value corresponding to the r,g,b value of pixel for some class.
My final semantic segmentation output is of size <batch, 35, height, width>
So doing the following I get the indexes of the pixels having the highest probability
num_classes = 35
output = output_from_segmentation_model()
_, pred = torch.max(output, dim=1)
Hence pred is a tensor with the index values where pred.shape = [batch, height, width] and pred.min = 0, pred.max = 34
Now I want to cast this an 3 channel image :
img = torch.zeroes(output.size[0], output.size[2], output.size[3])
such that for the respective value of pixel in pred the image has r,g,b frrom that class in cmap, where cmap is given as follows :
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]])
something in pseudocode :
img[…] = img[cmap[pred[x]]]
Please advice on correct syntax…