Relabeling (grouping) masks

Deep Learning
1

I am trying to train a model on the ADE20k image segmentation dataset but I want to convert the 150 classes into only 7 classes, i.e. floor = [‘floor’, ‘road’, ‘sidewalk’, ‘path’, ‘runway’].

How can I relabel the masks so that I am only training with 7 classes made from the initial 150 classes?

2

How are your masks encoded? If they’re images, you’ll want to load the images and change the values corresponding to [‘floor’, ‘road’, ‘sidewalk’, ‘path’, ‘runway’] to just “floor”. You can do this reasonably efficiently by indexing into the tensor:

image_mask=torch.randint(0,20,(10,10))  #10x10 image with 20 masks
image_mask[image_mask==5] = 2           # Lets say the mask label of 5 is road and 2 is floor

Then repeat this for the classes you want to merge.

3

Thanks for the reply. Yes. They are PNG files. It turns out I had to get the data property off the opened mask to do that after using open_mask to open it. Like…
mask.data[mask.data==5] = 2

However, now I run into a few issues. This doesn’t seem to change the value of mask.data. When I look at mask.data again, it is the same. If I try doing mask.data = mask.data[mask.data==5] = 2, it says I can’t set that property.

Also, how can I replace all values in an array with a new number. For example, in the below code I want to replace all the values in wall_ind with the first value in wall_ind:

mask.data[mask.data in wall_ind] = wall_ind[0]

Thanks for your help.

EDIT: I figured out the immutable issue by reading the source code. After opening the mask you get an ImageSegment which has a property called data that returns a tensor. I stored the result of mask.data in a new variable and then was able to modify that variable with reconstructed_mask_tensor = mask_tensor[mask_tensor==5] = 2. I used the result of this in the reconstruct function on ImageSegment like so mask.reconstruct(reconstructed_mask_tensor). The source code for ImageSegment can be found here.

I still haven’t figured out my second issue regarding replacing all the values in an array.

EDIT 2: I managed to relabel all the masks with the below code, however, I can’t figure out how to save these masks as labels. It just gets lost in the local variable.

for img in src.items:
    mask = open_mask(get_y_fn(img))
    mask_tensor = mask.data
    mask_tensor.apply_(lambda e: (e in wall_ind and wall_ind[0]) or e)
    mask.reconstruct(mask_tensor)

I also don’t know if the above code can be sped up by vectorizing it somehow. At the moment it is very slow.

4

I managed to fix the issue I was having where I couldn’t save the newly created masks. Using the .save() method on the ImageSegment doesn’t work because it defaults to the Image's .save() method which makes multiplies the pixel values by 255. The value of the masks are not in this format. I discovered this thanks to the below post. The mentioned function for saving worked well for me.