import torch


def inverse_one_hot(image_in, bg_label: int = -1, dim: int = -3, threshold: float = 0.1):
    """
    Takes float tensor and compute the argmax and max_value along the specified dimension.
    Returns a integer tensor of the same shape as the input_tensor but with the :attr:`dim` removed.
    If the max_value is less than the threshold the bg_label is assigned.

    Note:
        It can take an image of size :math:`(C, W, H)` and generate an integer mask
        of size :math:`(W, H)`.
        This operation can be thought as the inverse of the one-hot operation which takes an integer tensor of size (*)
        and returns a float tensor with an extra dimension, for example (*, num_classes).

    Args:
        image_in: any float tensor
        bg_label: integer, the value assigned to the entries of which are smaller than the threshold
        dim: int, the dimension along which to compute the max.
            For images this is usually the channel dimension, i.e. -3.
        threshold: float, the value of the threshold. Value smaller than this are set assigned to the background

    Returns:
        out: An integer mask with the same size of the input tensor but with the dim removed.
    """
    assert isinstance(bg_label, int), "Error. bg_label must be an integer. Received {0}".format(bg_label)

    _values, _indices = torch.max(image_in, dim=dim)
    _mask_bg = (_values < threshold)
    _indices[_mask_bg] = bg_label
    return _indices.long()