时间戳

首页 > 解决方案 > 更快的 numpy 数组索引

python - 更快的 numpy 数组索引

问题描述

我想索引 RGB 图像中的一些特定像素。我对 Python 比较陌生，所以我已经像在 Java/C# 中那样实现了索引。

这是代码：

# Pane width (height)
pane_step, center = 20, 10

for i in range(0, field_size * pane_step, pane_step):
    for j in range(0, field_size * pane_step, pane_step):
        r, g, b, = img[i + center, center + j, :]
        if (r, g, b) == (255, 0, 0):
            grid[int(i / pane_step)][int(j / pane_step)] = 2
        elif (r, g, b) == (0, 128, 0):
            grid[int(i / pane_step)][int(j / pane_step)] = 1
        elif (r, g, b) == (0, 0, 0):
            grid[int(i / pane_step)][int(j / pane_step)] = -1

有没有更快、更“pythonic”的方法可以给我同样的结果？

PS img 是一个 numpy ndarray

标签： pythonnumpy

解决方案

有没有更快、更“pythonic”的方法可以给我同样的结果？

也许，这取决于您的操作是否可以在 c 层中进行矢量化和完成（即，# Do stuff...对于您的特定情况，了解块实际上是什么很重要）。您至少可以使用切片和精美的索引而不是循环（并且在一行中）获取数据：

import numpy as np
img = np.arange(30000).reshape(100, 100, 3)

pane_step, center = 20, 10
field_size = 5

rs = []
gs = []
bs = []

# Original code
for i in range(0, field_size * pane_step, pane_step):
    for j in range(0, field_size * pane_step, pane_step):
        r, g, b, = img[i + center, center + j, :]
        rs.append(r)
        gs.append(g)
        bs.append(b)

# You want center (i = 0) to field_size * pane_step + center (i=field_size * pane_step) going by pane_step for the first dim
# Same for second dim
# Use fancy indexing and slice notation instead of doing one ind at a time
r_fancy, g_fancy, b_fancy = img[center:center+pane_step*field_size:pane_step, center:center+pane_step*field_size:pane_step, :].T

r_fancy_list = [x for x in r_fancy.T.reshape(-1)]
g_fancy_list = [x for x in g_fancy.T.reshape(-1)]
b_fancy_list = [x for x in b_fancy.T.reshape(-1)]

# Same data, just a different "shape" and orientation (you could transpose the result right away)
assert r_fancy_list == rs
assert g_fancy_list == gs
assert b_fancy_list == bs

让我们保持简单，假设您只想对每个像素值求平方并存储结果（我怀疑您是否真的想要这样做，这只是为了表明如果您的操作是矢量化的，这会更快）：

import numpy as np

# Original code with squaring
def square_em():
    img = np.arange(30000).reshape(100, 100, 3)
    pane_step, center = 20, 10
    field_size = 5
    rs = []
    gs = []
    bs = []
    for i in range(0, field_size * pane_step, pane_step):
        for j in range(0, field_size * pane_step, pane_step):
            r, g, b, = img[i + center, center + j, :]
            # Doing stuff...?
            rs.append(r**2)
            gs.append(g**2)
            bs.append(b**2)
    return rs, gs, bs

# Vectorized squaring
def square_em_vec():
    img = np.arange(30000).reshape(100, 100, 3)
    pane_step, center = 20, 10
    field_size = 5
    # Scroll over, tacked on a **2 to the end...
    r, g, b = img[center:center+pane_step*field_size:pane_step, center:center+pane_step*field_size:pane_step, :].T ** 2
    return r.T, g.T, b.T

我把它放在一个名为 test.py 的文件中，并将使用 IPython REPL 进行计时（只是因为它很方便，你也可以使用cProfile或其他东西）：

In [1]: from test import square_em, square_em_vec

In [2]: %timeit square_em()
10000 loops, best of 3: 83.9 µs per loop

In [3]: %timeit square_em_vec()
The slowest run took 5.00 times longer than the fastest. This could mean that an intermediate result is being cached.
10000 loops, best of 3: 32.8 µs per loop

HTH。

推荐阅读