更快的numpy数组索引

时间:2018-10-22 18:52:29

标签: 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”的方法可以给我相同的结果?

P.S。 img是一个numpy ndarray

2 个答案:

答案 0 :(得分:0)

  

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

答案 1 :(得分:0)

这是您要找的吗?您可以使用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]])