Faster numpy array indexing

Question

I want to index some specific pixels inside a RGB image. I'm rather new to Python so I've implemented the indexing like I would do that in Java/C#.

Here's the code:

# 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

Is there a faster, more "pythonic" approach that would give me the same result?

PS img is a numpy ndarray

Answer 1

Is there a faster, more "pythonic" approach that would give me the same result?

Maybe, it depends if your operations can be vectorized and done in the c-layer (ie it's important to know what the # Do stuff... block actually is for your particular case). You can at least get the data with a slice and fancy indexing instead of a loop (and in one line):

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

Let's keep it simple and pretend you just want to square each pixel value and store the result (I doubt you actually want to do this, it's just to show that if your operation is vectorized this will be faster):

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

I put this in a file called test.py, and will use the IPython REPL for timings (just because it's convenient, you could use cProfile or something too):

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

Is this what you're looking for? You can select a center and a step size easily with numpy and then iterate over the array or directly use it in any operation.

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]])