Python 使用可变二维遮罩在numpy中高效过滤三维矩阵

我有一个[10000x3000x128]维度的3D numpy数组点，其中第一个维度是帧数，第二个维度是每个帧中的点数，第三个维度是与每个点关联的128个元素的特征向量。我想做的是通过使用[10000x3000]尺寸的布尔2D掩码有效地过滤每一帧中的点，并且对于每个选定的点也采用相关的128维特征向量。此外，在输出中，我仍然需要一个3D向量，而不是一个合并的2D向量，并且可能避免任何for循环。 实际上我正在做的是：

# example of points
points = np.array([10000, 3000, 128])

# fg, bg = 2D dimensional boolean np.array

# init empty lists
fg_points, bg_points = [], []

for i in range(points.shape[0]):
     fg_mask_tmp, bg_mask_tmp = fg[i], bg[i]
     fg_points.append(points[i,fg_mask_tmp,:])
     bg_points.append(points[i,bg_mask_tmp,:])

fg_features, bg_features = np.array(fg_points), np.array(bg_points)

但这是一个非常幼稚的解决方案，可以用一种更像numpy的方式来改进。 此外，我还尝试了其他解决方案，如：

fg_features = points[fg,:]

但此解决方案不会保留合并前两个维度的数组的维度，因为每个帧的过滤点数可能会有所不同。 我尝试过的另一个解决方案是通过在最后一个维度上附加[128]真值来放大2D遮罩，但没有成功的结果。 有人知道可能的有效解决方案吗

提前感谢您的帮助