向量化循环操作

Question

我有一个使用 for 循环的操作。谁能建议一种使用 numpy 对操作进行矢量化的方法？

# rgb is a 3 channel image  
# points are computed using vector mult op (same size as rgb image)
# dtypes - rgb is uint8 and points is float
buffer = []
for v in range(rgb.shape[1]):
        for u in range(rgb.shape[0]):
            X,Y,Z = points[u,v,:]
            r,g,b = rgb[u,v,:] 
            buffer.append(struct.pack('ffffBBBBIII', X,Y,Z,0.0,
                                                     b,g,r,
                                                     255,0,0,0))

我想压缩上面的操作并获得缓冲区。任何指针都会有所帮助

Answer 1

我可以这样想：

def pack(points, rgb):
    X, Y, Z = points.transpose(2, 0, 1)
    r, g, b = rgb.transpose(2, 0, 1)
    rec = np.rec.fromarrays([X,Y,Z,r,g,b])
    def fn(point):
        X, Y, Z, r, g, b = point
        return struct.pack('ffffBBBBIII', X, Y, Z, 0.0, b, g, r, 255, 0, 0, 0)
    return np.vectorize(fn)(rec)

Answer 2

您可以使用 structured arrray 和

这样的 dtype

np.dtype([('point', 'f8', (4,)), ('rgb', 'uint8', (4,)), ('something', 'uint8', (4,))])

之后可以调用数组的方法ndarray.tobytes()获取buffer

初始化见here or here。

Answer 3

如果你有结构类型（C 类型）和 numpy 数值类型之间的对应关系，这应该相当简单。结构的文档是 here, while numpy's is here。相关的转换是：

'f' -> np.single（Python 没有等效类型）
'B' -> np.ubyte
'I' -> np.uintc

您可以通过 custom dtype 将输出创建为值数组，就像 struct 允许您做的那样：

dt = np.dtype([(c, np.single) for c in 'XYZW'] +
              [(c, np.ubyte) for c in 'RGBA'] +
              [('', np.intc, 3)])

为每个通道创建单独的字段（例如 [('X', np.single), ('Y', np.single), ...）而不是为所有通道创建一个字段（例如 [('XYZW', np.single, 4), ...）的原因是您希望能够访问数组统一的步伐。您不会分配给的空白部分可以是每个元素中的单个块：('zeros', np.intc, 3).

您可以使用其他数据类型来提供您想要的结果。例如，您可以命名您的字段，或将它们拆分为单独的通道。我建议您在视图中编写输出数组后执行此操作以简化处理。

现在你有了一个 dtype，用它创建一个数组：

output = np.zeros(rgb.shape[:2], dtype=dt)

现在您可以使用 dt.fields attribute combined with output.setfield:

存储字段

for name, plane in zip('XYZ', np.moveaxis(points, -1, 0)):
    tp, offset, *_ = dt.fields[name]
    output.setfield(plane, tp, offset)

for name, plane in zip('RGB', np.moveaxis(rgb, -1, 0)):
    tp, offset, *_ = dt.fields[name]
    output.setfield(plane, tp, offset)

tp, offset, *_ = dt.fields['A']
output.setfield(255, tp, offset)

您可以使用 itertoools.chain 缩短为单个循环：

from itertools import chain

for name, plane in zip('XYZRGBA', chain(np.moveaxis(points, -1, 0),
                                        np.moveaxis(rgb, -1, 0),
                                        [255])):
    tp, offset, *_ = dt.fields[name]
    output.setfield(plane, tp, offset)

请注意，这里的循环并不是非常昂贵。它只经历了七次迭代。结果的每个元素都是一个缓冲区，其形式与您的结构调用所创建的形式完全相同。您可以通过分解 output 数组来丢弃形状信息。

结果是一个带有自定义数据类型的数组，相当于 'ffffBBBBIII' 的结构格式规范。每个元素都是一个标量，由字段名称索引：

>>> output[0, 0]['W']
0.0

如果需要，您可以在数组中创建替代视图，例如，将值分组到类别或类似的东西中：

>>> dt2 = np.dtype([('p', np.single, 4), ('i', np.ubyte, 4), ('z', np.intc, 3)]
>>> output2 = output.view(dtype=dt2)
>>> output2[0, 0]['p']
array([0.501182 , 0.7935149, 0.9981835, 0.       ], dtype=float32)  # Random example data

此视图不复制数据，只是以不同的方式解释现有缓冲区。在内部，表示仍然是您试图通过 struct.

实现的压缩版本

向量化循环操作

Vectorize loop operation

python

numpy

vectorization

numpy-ndarray