How to convert a 55-color PNG segmentation mask in a [55,h,w] tensor?

vision
2

This post gives you an example of a color to class index mapping for the targets of a segmentation use case.