import torch


def onehot(image, width: int, height: int, col: dict):
    """
    This function takes the target pictures and onehot encode them. The amount
    of classes are defined by the amount of different classes in the col
    dictionary.

    Args:
        image (tensor): target image meant to be onehot encoded.
        width (int): width of the given image.
        height (int): height of the given image.
        col (dict): a mapping of the different colors to differnt layers.

    Return:
        Tensor: A onehot encoded version of the target image.
    """
    map = torch.zeros((len(set(col.values())), height, width), dtype=torch.uint8)
    # make empty tensor

    for r in range(height):
        for j in range(width):
            num = col.get(tuple(image[r][j].numpy()), 0)
            map[num][r][j] = 1

    return map