添加 numpy 数组时避免溢出

Avoid overflow when adding numpy arrays

我想添加数据类型为 uint8 的 numpy 数组。我知道这些数组中的值可能大到足以发生溢出。所以我得到类似的东西:

a = np.array([100, 200, 250], dtype=np.uint8)
b = np.array([50, 50, 50], dtype=np.uint8)
a += b

现在,a 是 [150 250 44]。但是,我想要的不是溢出,而是对于 uint8 来说太大的值是 uint8 允许的最大值。所以我想要的结果是 [150 250 255].

我可以用下面的代码得到这个结果:

a = np.array([100, 200, 250], dtype=np.uint8)
b = np.array([50, 50, 50], dtype=np.uint8)
c = np.zeros((1,3), dtype=np.uint16)
c += a
c += b
c[c>255] = 255
a = np.array(c, dtype=np.uint8)

问题是,我的数组非常大,因此创建具有更大数据类型的第三个数组可能是内存问题。是否有一种快速且内存效率更高的方法来实现所描述的结果?

怎么样

>>> a + np.minimum(255 - a, b)
array([150, 250, 255], dtype=uint8)

通常使用

获取数据类型的最大值
np.iinfo(np.uint8).max

您可以通过创建第三个 dtype uint8 数组和一个 bool 数组 (它们一起比一个 uint16 数组的内存效率更高)来实现这一点.

np.putmask 对于避免临时数组很有用。

a = np.array([100, 200, 250], dtype=np.uint8)
b = np.array([50, 50, 50], dtype=np.uint8)
c = 255 - b  # a temp uint8 array here
np.putmask(a, c < a, c)  # a temp bool array here
a += b

但是,正如@moarningsun 正确指出的那样,bool 数组占用的内存量与 uint8 数组相同,因此这不一定有用。可以通过在任何给定时间避免拥有多个临时数组 :

来解决此问题
a = np.array([100, 200, 250], dtype=np.uint8)
b = np.array([50, 50, 50], dtype=np.uint8)
b = 255 - b  # old b is gone shortly after new array is created
np.putmask(a, b < a, b)  # a temp bool array here, then it's gone
a += 255 - b  # a temp array here, then it's gone

此方法用内存消耗换取 CPU。


另一种方法是预先计算所有可能的结果,这是 O(1) 额外内存(即独立于数组的大小):

c = np.clip(np.arange(256) + np.arange(256)[..., np.newaxis], 0, 255).astype(np.uint8)
c
=> array([[  0,   1,   2, ..., 253, 254, 255],
          [  1,   2,   3, ..., 254, 255, 255],
          [  2,   3,   4, ..., 255, 255, 255],
          ..., 
          [253, 254, 255, ..., 255, 255, 255],
          [254, 255, 255, ..., 255, 255, 255],
          [255, 255, 255, ..., 255, 255, 255]], dtype=uint8)

c[a,b]
=> array([150, 250, 255], dtype=uint8)

如果您的数组非常大,这种方法是最节省内存的方法。同样,它的处理时间很昂贵,因为它用较慢的 2dim 数组索引代替了超快的​​整数加法。

工作原理说明

上面 c 数组的构造使用了 numpy 广播技巧。添加一个形状为 (N,) 的数组和形状为 (1,N) 的数组广播为 (N,N)-like,因此结果是一个包含所有可能总和的 NxN 数组。然后,我们剪辑它。我们得到一个 2dim 数组满足: c[i,j]=min(i+j,255) for each i,j.

然后剩下的就是使用花哨的索引来获取正确的值。使用您提供的输入,我们访问:

c[( [100, 200, 250] , [50, 50, 50] )]

第一个索引数组引用第一个暗淡,第二个索引数组引用第二个暗淡。因此,结果是一个与索引数组 ((N,)) 具有相同形状的数组,由值 [ c[100,50] , c[200,50] , c[250,50] ].

组成

这里有一个方法:

>>> a = np.array([100, 200, 250], dtype=np.uint8)
>>> b = np.array([50, 50, 50], dtype=np.uint8)
>>> a+=b; a[a<b]=255
>>> a
array([150, 250, 255], dtype=uint8)

您可以使用 Numba 真正做到这一点,例如:

import numba

@numba.jit('void(u1[:],u1[:])', locals={'temp': numba.uint16})
def add_uint8_inplace_clip(a, b):
    for i in range(a.shape[0]):
        temp = a[i] + b[i]
        a[i] = temp if temp<256 else 255

add_uint8_inplace_clip(a, b)

或者用 Numexpr,例如:

import numexpr

numexpr.evaluate('where((a+b)>255, 255, a+b)', out=a, casting='unsafe')

Numexpr upcasts uint8int32 内部,然后将其放回 uint8 数组。

a function in numpy 个:

numpy.nan_to_num(x)[source]

Replace nan with zero and inf with finite numbers.

Returns an array or scalar replacing Not a Number (NaN) with zero, (positive) infinity with a very large number and negative infinity with a very small (or negative) number.

New Array with the same shape as x and dtype of the element in x with the greatest precision.

If x is inexact, then NaN is replaced by zero, and infinity (-infinity) is replaced by the largest (smallest or most negative) floating point value that fits in the output dtype. If x is not inexact, then a copy of x is returned.

我不确定它是否适用于 uint8,因为在输出中提到了浮点数,但对于其他读者来说,它可能有用

def non_overflowing_sum(a, b)
    c = np.uint16(a)+b
    c[np.where(c>255)] = 255
    return np.uint8( c )

它也交换内存,但我发现更优雅并且临时 uint16 在 return

转换后被释放

OpenCV有这样一个函数:cv2.addWeighted