使用np.pad()为RGB图像填充时出现广播错误的原因及解决方法

编写color_image_padding函数实现RGB图像边界零填充功能，输入图像维度为(Width, Height, 3)，最后一维对应3个颜色通道。

最初实现代码：

import numpy as np

def color_image_padding(image: np.ndarray) -> np.ndarray:
    return np.pad(image, pad_width=1)

运行时报错：

operands could not be broadcast together with shapes (4,4,3) (4,4,5)

最初怀疑报错和颜色通道有关，不确定np.pad是否会自动拆分3个通道分别完成填充。

排查过程 #

后续定位到通用填充函数image_padding()报错的根因：传入的图像中混有2维结构的灰度图，最小复现代码如下：

bar = np.ones((1, 3))
print(bar.ndim) # 输出2，对应2维灰度图

def image_padding(image: np.ndarray, amt: int) -> np.ndarray:
    return np.pad(image, pad_width=((amt, amt), (amt, amt), (0, 0)))

image_padding(bar, 2)

完整报错栈：

ValueError                                Traceback (most recent call last)
~\AppData\Local\Temp/ipykernel_8116/4065018867.py in <module>
----> 1 img(bar, 3)

~\AppData\Local\Temp/ipykernel_8116/1455868751.py in img(image, amt)
      1 def img(image, amt):
----> 2     return np.pad(image, pad_width=((amt, amt), (amt, amt), (0, 0)))

<__array_function__ internals> in pad(*args, **kwargs)

~\anaconda3\lib\site-packages\numpy\lib\arraypad.py in pad(array, pad_width, mode, **kwargs)
    741 
    742     # Broadcast to shape (array.ndim, 2)
--> 743     pad_width = _as_pairs(pad_width, array.ndim, as_index=True)
    744 
    745     if callable(mode):

~\anaconda3\lib\site-packages\numpy\lib\arraypad.py in _as_pairs(x, ndim, as_index)
    516     # Converting the array with `tolist` seems to improve performance
    517     # when iterating and indexing the result (see usage in `pad`)
--> 518     return np.broadcast_to(x, (ndim, 2)).tolist()
    519 
    520 

<__array_function__ internals> in broadcast_to(*args, **kwargs)

~\anaconda3\lib\site-packages\numpy\lib\stride_tricks.py in broadcast_to(array, shape, subok)
    409            [1, 2, 3]])
    410     """
--> 411     return _broadcast_to(array, shape, subok=subok, readonly=True)
    412 
    413 

~\anaconda3\lib\site-packages\numpy\lib\stride_tricks.py in _broadcast_to(array, shape, subok, readonly)
    346                          'negative')
    347     extras = []
--> 348     it = np.nditer(
    349         (array,), flags=['multi_index', 'refs_ok', 'zerosize_ok'] + extras,
    350         op_flags=['readonly'], itershape=shape, order='C')

ValueError: operands could not be broadcast together with remapped shapes [original->remapped]: (3,2)  and requested shape (2,2)

np.pad对pad_width参数的维度匹配有严格要求：

• 传入单个标量作为pad_width时，会对输入数组的所有维度都对称添加该宽度的填充。最初的实现给3维RGB图传pad_width=1，会给通道维也前后各加1个单位填充，通道数从3变为5，和后续逻辑预期的3通道形状不匹配，触发第一个形状广播报错。
• 传入按维度划分的填充序列时，序列长度必须和输入数组的维度数完全一致。给2维灰度图传入3组填充值（对应宽、高、通道3个维度）时，参数长度和数组维度不匹配，就会触发第二个广播报错。

填充前先判断图像维度，针对灰度图、彩色图分别传入维度匹配的pad_width参数即可，修正后的代码如下：

def image_padding(image: np.ndarray, amt: int) -> np.ndarray:
    if image.ndim == 2:
        # 2维灰度图，仅对宽高两个维度做填充
        return np.pad(image, pad_width=(amt, amt))
    elif image.ndim == 3:
        # 3维彩色图，仅对宽高维度填充，通道维不做填充
        return np.pad(image, pad_width=((amt, amt), (amt, amt), (0, 0)))

原排查给出的修正代码存在笔误：elif分支误写为img.ndim == 3，变量名和入参不一致会触发未定义错误，上述代码已修正该问题。

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