import torch


def poolfeat(input, prob, avg = True):
    """ A function to aggregate superpixel features from pixel features

    Args: 
      input (tensor): input feature tensor.
      prob (tensor): one-hot superpixel segmentation.
      avg (bool, optional): average or sum the pixel features to get superpixel features
    
    Returns:
      cluster_feat (tensor): the superpixel features
    
    Shape:
      input: (B, C, H, W)
      prob: (B, N, H, W)
      cluster_feat: (B, N, C)
    """ 
    B, C, H, W = input.shape
    B, N, H, W = prob.shape
    prob_flat = prob.view(B, N, -1)
    input_flat = input.view(B, C, -1)
    cluster_feat = torch.matmul(prob_flat, input_flat.permute(0, 2, 1))
    if avg:
        cluster_sum = torch.sum(prob_flat, dim = -1).view(B, N , 1)
        cluster_feat = cluster_feat / (cluster_sum + 1e-8)
    return cluster_feat