import torch


def upfeat(input, prob):
    """ A function to compute pixel features from superpixel features

    Args: 
      input (tensor): superpixel feature tensor.
      prob (tensor): one-hot superpixel segmentation.
    
    Returns:
      reconstr_feat (tensor): the pixel features.
    
    Shape:
      input: (B, N, C)
      prob: (B, N, H, W)
      reconstr_feat: (B, C, H, W)
    """
    B, N, H, W = prob.shape
    prob_flat = prob.view(B, N, -1)
    reconstr_feat = torch.matmul(prob_flat.permute(0, 2, 1), input)
    reconstr_feat = reconstr_feat.view(B, H, W, -1).permute(0, 3, 1, 2)

    return reconstr_feat