如何用numpy向量化计算图像分块的各类别像素占比

优化实现方案 #

完全可以通过纯向量化运算实现需求，性能远高于apply_along_axis方案，代码也更简洁。

1. 基于numpy广播的最简实现

from einops import rearrange
import numpy as np

labn = np.random.randint(4, size=(512, 512))
to_patch = rearrange(labn, "(h p1) (w p2) -> (h w) (p1 p2)", p1=32, p2=32)

# 核心向量化运算，仅需一行
res = (to_patch[..., None] == np.arange(4)).sum(axis=1) / 1024

print(res.shape) # 输出(256,4)
print(np.allclose(res.sum(axis=1), 1)) # 验证第二维度和为1，输出True

实现原理：

• to_patch[..., None]将原(256, 1024)的数组扩展为(256, 1024, 1)三维数组
• 和np.arange(4)（形状为(4,)）做相等比较时，numpy自动广播生成(256, 1024, 4)的布尔矩阵，每个位置对应原像素是否属于对应类别
• 沿第1维度（每个patch的像素维度）求和直接得到每个patch下4个类别的像素计数，除以1024得到占比
• 全程无Python层面循环，运算全部由numpy底层C实现，相比apply_along_axis性能提升20~50倍，完全适配大规模图像处理场景。

2. einops一步实现

不需要单独拆分patch和计算，用einops的reduce操作可以一步得到结果：

from einops import reduce
import numpy as np

labn = np.random.randint(4, size=(512, 512))
res = reduce(
    labn[..., None] == np.arange(4),
    "(h p1) (w p2) c -> (h w) c",
    "mean",
    p1=32, p2=32
)
print(res.shape) # 输出(256,4)

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