How do we make our outputs to have the same size as the true mask?

Question

When we are doing multi-label segmentation tasks, our y_true (the mask) will be (w, h, 3), but, in our model, at the last layer, we will be getting (w, h, number of classes) as output.

How do we make our outputs to have the same size as the true mask so that to apply the loss function, given that, currently, the shapes are not equal? Also, if we are done with applying the loss function and trained the model, how do I make results in the shape of (w, h, 3) from (w, h, number of classes)?

Answer 1

You can create a mapping from classes to colors for a simple one is:

y # y.shape = (W, H, n_classes)
_, y_color = y.max(dim=-1, keepdim=True) / n_classes # y_color.shape = (W, H, 1)
y_color = torch.cat([y_color] * 3, dim=-1) # y_color.shape = (W, H, 3)

(using pytorch like code)

This mapping is visualizable, of course, you may get nicer visualizations if your class to color mapping is more sophisticated.