import torch


def mask_to_label(mask, classes):
    """
    Returns the combined masks, i.e. label for a mask based on the color classes.

    Arguments:
        mask {torch.Tensor} -- mask tensor [shape = (B, N, H, W) ; B = Batch Size, N = Number of Classes]
        classes {torch.Tensor} -- color classes tensor [shape = (N, 3)]

    Returns:
        torch.Tensor -- combined masks / label [shape = (B, 3, H, W)]
    """

    mask_indices = torch.argmax(mask, dim=1)

    return (
        classes[mask_indices.view(-1)].view(*mask_indices.shape, 3).permute(0, 3, 1, 2)
    )