Python中基于Numpy高效拼接4D数组图像块的方法

问题描述 #

给定一个形状为(N,W,H,3)的4D NumPy数组，其中：

• N是图像块总数
• W、H是单块图像的宽和高
• 3是RGB颜色通道数

这些图像块是按行优先顺序从原始图像分割得到的正方形块（比如3×3分割后，数组顺序是[patch11,patch12,patch13,patch21,patch22,patch23,patch31,patch32,patch33]）。需要仅用Python基础函数和NumPy（不使用PIL/OpenCV）实现最高效的拼接方法，替代现有的双重循环实现。

现有循环实现 #

import numpy as np

def reconstruct_image(patches, num_rows, num_cols):
    # num_rows和num_cols分别是图像块的行数和列数
    patch_height, patch_width, channels = patches.shape[1], patches.shape[2], patches.shape[3]

    # 初始化空的完整图像数组
    full_image = np.zeros((num_rows * patch_height, num_cols * patch_width, channels), dtype=patches.dtype)

    # 遍历图像块的行和列
    for i in range(num_rows):
        for j in range(num_cols):
            # 获取当前图像块在4D数组中的索引
            patch_index = i * num_cols + j
            # 将图像块放置到完整图像的对应位置
            full_image[i*patch_height:(i+1)*patch_height, j*patch_width:(j+1)*patch_width, :] = patches[patch_index]

    return full_image

N = 9  # 图像块数量
W, H, C = 10, 10, 3  # 图像块尺寸（宽×高×通道数）
num_rows, num_cols = 3, 3  # 图像块的行列数（3×3）
patches = np.random.rand(N, W, H, C)  # 示例图像块数据

reconstructed_image = reconstruct_image(patches, num_rows, num_cols)

高效无循环实现 #

利用NumPy的**形状重塑(reshape)和轴交换(transpose)**可以完全避免循环，大幅提升性能，核心思路是：

1. 先将图像块数组重塑为(num_rows, num_cols, W, H, 3)，还原块的行列结构
2. 交换轴的顺序，把块内的宽高维度和块的行列维度合并
3. 最终拼接成完整图像

实现代码：

import numpy as np

def reconstruct_image_fast(patches, num_rows, num_cols):
    patch_h, patch_w, channels = patches.shape[1], patches.shape[2], patches.shape[3]
    # 重塑为块的行列结构，再交换轴合并维度
    return patches.reshape(num_rows, num_cols, patch_h, patch_w, channels) \
                  .transpose(0, 2, 1, 3, 4) \
                  .reshape(num_rows * patch_h, num_cols * patch_w, channels)

# 测试示例
N = 9
W, H, C = 10, 10, 3
num_rows, num_cols = 3, 3
patches = np.random.rand(N, W, H, C)

# 验证两种方法结果一致
recon1 = reconstruct_image(patches, num_rows, num_cols)
recon2 = reconstruct_image_fast(patches, num_rows, num_cols)
assert np.array_equal(recon1, recon2)

性能优势说明 #

• 循环实现依赖Python层面的迭代，每次赋值都涉及数组切片操作，在块数量较多时性能会显著下降
• 无循环实现完全基于NumPy的底层C操作，避免了Python循环的开销，对于大规模图像块数据，速度提升可达几十甚至上百倍

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