针对8x8子矩阵处理的双层循环代码性能优化求助

你的代码核心是对图像的每个重叠8x8窗口计算DCT，统计低幅值系数数量，再更新对应区域的投票矩阵和零计数矩阵。双层循环的方式在Python中效率极低，图像尺寸较大时耗时会呈指数级增长。下面是针对性的优化思路和实现：

核心优化点 #

• 避免Python层的双层循环，用向量化操作批量处理所有窗口
• 一次性生成所有重叠8x8窗口，批量计算DCT和零计数
• 用矩阵广播和累积比较替代逐窗口的区域更新

优化后的代码 #

import cv2
import numpy as np
from skimage.util import view_as_windows

def compute_grid_votes(path):
    # 加载图像并转换为float32
    L = cv2.imread(path, cv2.IMREAD_GRAYSCALE).astype(np.float32)
    h, w = L.shape

    # 初始化votes和zeros
    votes = np.full((h, w), -1, dtype=np.int16)
    zeros = np.zeros((h, w), dtype=np.float32)

    # 生成所有重叠的8x8窗口（形状：(h-7, w-7, 8, 8)）
    windows = view_as_windows(L, window_shape=(8, 8), step=1)

    # 批量计算所有窗口的DCT
    # 先将窗口转为适合cv2.dct的形状：(N, 8, 8)，N为窗口总数
    dct_blocks = np.array([cv2.dct(window) for window in windows.reshape(-1, 8, 8)])
    # 统计每个窗口的低幅值系数数量（绝对值<0.5）
    z_values = np.sum(np.abs(dct_blocks) < 0.5, axis=(1, 2)).reshape(windows.shape[0], windows.shape[1])

    # 生成每个窗口对应的投票索引值：8*(i%8)+(j%8)
    i_indices = np.arange(h-7)[:, np.newaxis]
    j_indices = np.arange(w-7)
    vote_indices = 8 * (i_indices % 8) + (j_indices % 8)

    # 遍历每个窗口的位置，用向量化操作更新区域
    for i in range(h-7):
        for j in range(w-7):
            z = z_values[i, j]
            idx = vote_indices[i, j]
            # 获取当前窗口对应的区域
            roi_zeros = zeros[i:i+8, j:j+8]
            roi_votes = votes[i:i+8, j:j+8]

            # 按原逻辑更新votes和zeros
            mask_greater = z > roi_zeros
            mask_equal = z == roi_zeros

            roi_votes[mask_greater] = idx
            roi_votes[mask_equal] = -1
            roi_zeros[mask_greater] = z

    return votes

进一步加速方案 #

如果上述代码仍达不到性能要求，可以尝试：

• 用Numba装饰器（@numba.jit(nopython=True)）修饰循环部分，将Python代码编译为机器码
• 替换view_as_windows为numpy原生的as_strided（需注意内存对齐），减少第三方库依赖
• 全向量化更新：将所有窗口的z值和索引按像素位置聚合，直接取每个像素的最大z值对应的索引，完全消除循环

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