Преобразовать изображение с размерами высоты, ширины, количества каналов в n_masks, image_height, image_width

Question

У меня есть изображение RGB, оно содержит маски разных цветов, каждый цвет представляет определенный класс.Я хочу преобразовать его в формат - n_masks, image_height, image_width, где n_masks - количество масок, присутствующих на изображении.и каждый фрагмент матрицы вдоль 0-й оси представляет одну двоичную маску.
До сих пор я был в состоянии преобразовать ее в формат image_height, image_width, где каждое значение массива представляет, к какому классу он принадлежит, но я вродеударил после этого.Ниже приведен мой код для преобразования его в image_height,image_width формат-

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

. Он преобразует изображение, скажем, формата (512,512,3) в (512,512), где значения каждого пикселя представляют класс, к которому он принадлежит,но я понятия не имею, как действовать дальше.
PS. Я кодирую его в pytorch, но также приветствуется любой подход, включающий numpy.
Прикрепление изображения для наглядности

Answer 1

Предположим, у вас уже есть (512,512) маска.Сначала вы можете использовать mask.unique(), чтобы получить значения пикселей всех классов.Тогда для каждого значения класса torch.where(mask==cls_val, torch.tensor(1), torch.tensor(0)) вернет маску одного определенного класса.В конце вы складываете все выходные данные.

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]]])