Hello, I am wondering if there are ways to compute the centroid of a binary blob in an image such that they are differentiable? Essentially, I am looking for a differentiable twin of the loss function shown here:
class CentroidLoss(nn.Module):
    def __init__(self):
        super(CentroidLoss, self).__init__()
        
    def forward(self, rendered_silhouette, original_silhouette):
        rendered_sil_indices = torch.nonzero(rendered_silhouette)
        rendered_sil_centroid = torch.Tensor.float(rendered_sil_indices).mean(dim=0)

        original_sil_indices = torch.nonzero(original_silhouette)
        original_sil_centroid = torch.Tensor.float(original_sil_indices).mean(dim=0)

        centroid_dist = ((rendered_sil_centroid[0] - original_sil_centroid[0]) ** 2 \
                        + (rendered_sil_centroid[0] - original_sil_centroid[0]) ** 2 ) \
                        / (rendered_silhouette.shape[0]*rendered_silhouette.shape[1])

        return centroid_dist

“rendered_silhouette” and “original_silhouette” are binary images, each containing a blob.