Modifing segmentation masks for only a few specific classes

vision
1

Hello,

Right now I have pseudo-labels for an image dataset from DeepLabV3+ that are of shape [H W, 1] where each pixel is an integer from 0 to 18 representing the class of the pixel (using cityscapes pretraied models).

Now, what I want to do is focus on only a few classes (say for example, 10, 12, 13) and map the rest of the classes to some background class. What is the best way to do this in PyTorch?

I guess numpy would also work.

Thanks!

2

Hi Dryden!

You therefore want four classes, “background”, “class 10”, “class 12”, and
“class 13.” You should then label them with with the four integer class
labels 0, 1, 2, 3.

Set up a translation table that maps the integers {0, 1, 2, ..., 17, 18}
to your desired smaller set of class labels, {0, 1, 2, 3}:

translate_table = torch.tensor ([0] * 19)
translate_table[10] = 1
translate_table[12] = 2
translate_table[13] = 3

Let’s say that labels is your pseudo-labels tensor (in your example
of shape [H, W, 1]). Then:

mapped_labels = translate_table[labels]

will be the class labels for your reduced four-class problem (and will
still have shape [H, W, 1]).

Best.

K. Frank

1 Like