Convert image of dimension height,width,number of channels to n_masks, image_height, image_width
I have an RGB image, it contains masks of different colors, each color represents a particular class.
I want to convert it into the format - n_masks, image_height, image_width where n_masks is number of masks present in the image. and each slice of the matrix along the 0th axis represents one binary mask.
So far, I have been able to convert it into the format of image_height, image_width where each array value represents which class it belongs to but I am kind of struck after it.
Below is my code to convert it into image_height,image_width format-

def mask_to_class(mask):
    target = torch.from_numpy(mask)
    h,w = target.shape[0],target.shape[1]
    masks = torch.empty(h, w, dtype=torch.long)
    colors = torch.unique(target.view(-1,target.size(2)),dim=0).numpy()
    target = target.permute(2, 0, 1).contiguous()
    mapping = {tuple(c): t for c, t in zip(colors.tolist(), range(len(colors)))}
    for k in mapping:
        idx = (target==torch.tensor(k, dtype=torch.uint8).unsqueeze(1).unsqueeze(2))
        validx = (idx.sum(0) == 3) 
        masks[validx] = torch.tensor(mapping[k], dtype=torch.long)
    return masks


It converts an image of let's say format (512,512,3) to (512,512) where each pixel values represents the class it belongs to, but I don't have any idea how to proceed further.
P.S- I am coding it in pytorch but any approach involving numpy is also welcomed.

Suppose you already have (512,512) mask. You can first use mask.unique() to get all classes pixel values. Then for each class value, torch.where(mask==cls_val, torch.tensor(1), torch.tensor(0)) will return the mask of one certain class. In the end, you stack all outputs. 

mask = torch.tensor([[1,2,3],
                     [2,4,5],
                     [1,2,3]])
cls = mask.unique()
res = torch.stack([torch.where(mask==cls_val, torch.tensor(1), torch.tensor(0)) for cls_val in cls])

#tensor([[[1, 0, 0],
#         [0, 0, 0],
#         [1, 0, 0]],
#
#        [[0, 1, 0],
#         [1, 0, 0],
#         [0, 1, 0]],
#
#        [[0, 0, 1],
#         [0, 0, 0],
#         [0, 0, 1]],
#
#        [[0, 0, 0],
#         [0, 1, 0],
#         [0, 0, 0]],
#
#        [[0, 0, 0],
#         [0, 0, 1],
#         [0, 0, 0]]])

