import torch.nn.functional as F
import torch


def get_class_count(box_class, n_class=-1):
    """
    Return a tensor where each element of the last dimension represent a class and its value is the number of
    bounding boxes with given class.

    :param box_class: A [*, d1] tensor
    :param n_class: Total number of classes
    :return: A [*, n_class] tensor
    """

    class_one_hot = F.one_hot(box_class, num_classes=n_class)
    class_one_hot = torch.transpose(class_one_hot, dim0=-1, dim1=-2)
    return class_one_hot.sum(dim=-1)