更快的numpy数组索引

Question

我想索引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

Answer 1

有没有更快，更“ 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。

Answer 2

这是您要找的东西吗？ 您可以使用numpy轻松选择中心和步长，然后遍历数组或直接在任何操作中使用它。

import numpy as np
d = np.arange(100).reshape((10,10))

center, step = 3

print(d[center::step, center::step])

array([[33, 36, 39],
      [63, 66, 69],
      [93, 96, 99]])