import torch


def get_centroid_indices(masks):
    """
    Params:
        masks: Tensor[num_objects, height, width]
    Returns:
        centroids: Tensor[num_objects, (x, y)]
    """
    _, height, width = masks.shape
    dtype = masks.dtype
    device = masks.device

    location_x = torch.arange(0, width, 1, dtype=dtype, device=device) # Tensor[width]
    location_y = torch.arange(0, height, 1, dtype=dtype, device=device) # Tensor[height]

    total_area = masks.sum(dim=(1,2)) + 1e-9
    centroids_x = torch.sum(masks.sum(dim=1) * location_x[None,:], dim=1) / total_area # Tensor[num_objects]
    centroids_y = torch.sum(masks.sum(dim=2) * location_y[None,:], dim=1) / total_area # Tensor[num_objects]

    centroids = torch.stack((centroids_x, centroids_y), dim=1) # Tensor[num_objects, (x, y)]
    centroids = centroids.to(torch.int64)
    return centroids