import torch


def one_hot_encoding(input, c):
    """
    One-hot encoder: Converts NxHxW label image to NxCxHxW, where each label is stored in a separate channel
    :param input: input image (NxHxW)
    :param c: number of channels/labels
    :return: output image  (NxCxHxW)
    """
    assert input.dim() == 3
    N, H, W = input.size()
    result = torch.zeros((N, c, H, W))
    # torch.Tensor.scatter_(dim, index, src) -> Tensor
    # eg: For 4d tensor
    #    self[i][index[i][j][k][h]][k][h] = src[i][j][k][h]    # if dim == 1
    result.scatter_(1, input.unsqueeze(1), 1)
    return result