DataLoader not reading in masks correctly?

[eqy]

The fact that mask is only two dimensional is an issue if it should have color information. Can you try something like

mask = torch.zeros(3, h, w, dtype = torch.long)

and then

mask[0,:,:][validx] = torch.tensor(k[0], dtype=torch.long)
mask[1,:,:][validx] = torch.tensor(k[1], dtype=torch.long)
mask[2,:,:][validx] = torch.tensor(k[2], dtype=torch.long)

There might be a way to clean up the indexing here but I think this illustrates the intent.

The issue is that mapping is from colors to classes, but classes are being assigned to the mask here rather than colors. So here we are using the “key” of the mapping for the mask as that actually contains the colors.

On the other hand, if mask is supposed to have a format like 1000x1000 or 1x1000x1000 where each position is supposed to store the class id, then you might want to reuse part of the code here

How to save multi-class segmentation prediction as image?

Ah, curr_color should be used instead: class_to_color = [torch.tensor([1.0, 0.0, 0.0]), ...] output = torch.zeros(1, 3, out.size(-2), out.size(-1), dtype=torch.float) for class_idx, color in enumerate(class_to_color): mask = out[:,class_idx,:,:] == torch.max(out, dim=1)[0] mask = mask.unsqueeze(1) # should have shape 1, 1, 100, 100 curr_color = color.reshape(1, 3, 1, 1) segment = mask*curr_color # should have shape 1, 3, 100, 100 output += seg…

to generate something that has 3 colors and is viewable