import torch


def calc_init_centroid(images, num_spixels_width, num_spixels_height):
    """
    calculate initial superpixels

    Args:
        images: torch.Tensor
            A Tensor of shape (B, C, H, W)
        num_spixels_width: int
            A number of superpixels in each column
        num_spixels_height: int
            A number of superpixels int each row

    Return:
        centroids: torch.Tensor
            A Tensor of shape (B, C, H * W)
        init_label_map: torch.Tensor
            A Tensor of shape (B, H * W)

    """
    batchsize, channels, height, width = images.shape
    device = images.device

    # 自适应平均池化
    centroids = torch.nn.functional.adaptive_avg_pool2d(images, (num_spixels_height, num_spixels_width))

    with torch.no_grad():
        num_spixels = num_spixels_width * num_spixels_height        # 一共多少个superpixel
        labels = torch.arange(num_spixels, device=device).reshape(1, 1, *centroids.shape[-2:]).type_as(centroids)   # 假如有20个superpixel，那么从上到下从左到右，依次标号
        init_label_map = torch.nn.functional.interpolate(labels, size=(height, width), mode="nearest")      # 每个像素标记初始superpixel标签，即按照上一步的图形，按照矩形进行标号
        init_label_map = init_label_map.repeat(batchsize, 1, 1, 1)      # 实现batchsize维度

    init_label_map = init_label_map.reshape(batchsize, -1)      # batch中每张图都展平（这个是pixel到superpixel的映射，因此channel维度是1
    centroids = centroids.reshape(batchsize, channels, -1)      # batch中每张图的每个channel维度上展平

    return centroids, init_label_map