import torch


def process_per_level(dense_features, mask, level, stride):
    """
    extract sparse features and indices, compute locations in image of them as well

    Arguments:
        dense_features: N*C*H*W
        mask: N*H*W*1

    Returns:
        sparse_features: m*3
        indices: m*C
        locations: m*2
    """

    # swap channels
    dense_features = dense_features.permute(0, 2, 3, 1).contiguous()

    mask = mask > 0

    indices = torch.nonzero(mask.squeeze(3))  # [[batch_idx, y_idx, x_idx]], m*3
    sparse_features = torch.masked_select(dense_features, mask).view(-1, dense_features.size(-1))
    locations = indices[:, [2, 1]].float() * stride + stride // 2

    return sparse_features, locations, indices