import torch


def get_number_of_voxels_per_class(labels: torch.Tensor) -> torch.Tensor:
    """
    Computes the number of voxels for each class in a one-hot label map.
    :param labels: one-hot label map in shape Batches x Classes x Z x Y x X or Classes x Z x Y x X
    :return: A tensor of shape [Batches x Classes] containing the number of non-zero voxels along Z, Y, X
    """
    if not len(labels.shape) in [5, 4]:
        raise Exception("labels must have either 4 (Classes x Z x Y x X) "
                        "or 5 dimensions (Batches x Classes x Z x Y x X), found:{}"
                        .format(len(labels.shape)))

    if len(labels.shape) == 4:
        labels = labels[None, ...]

    return torch.count_nonzero(labels, dim=(2, 3, 4))