CODEWITHSUNDEEP
pythonnumpy
Liondancer
Liondancer

Reputation: 16469

Optimize NumPy iterating through image and changing values

Is there a way to optimize this snippet of code? With my current im value, it is taking me ~28 seconds. Wondering if I could reduce this time.

im = output_image[-min_xy[1]:-min_xy[1] + image_2.shape[0], -min_xy[0]:-min_xy[0] + image_2.shape[1]]
    for idx, rk in enumerate(im):
        for ix, k in enumerate(rk):
            image_2[idx][ix] = avg(im[idx][ix], image_2[idx][ix])

type(image_2) and type(im) is <type 'numpy.ndarray'>

im.shape and image_2.shape is (2386, 3200, 3)

what my avg() does is 

def avg(a1, a2):
    if [0., 0., 0.] in a1:
        return a2
    else:
        return (a1 + a2) / 2

NOTE: a1 is an array of size 3 ex: array([ 0.68627451, 0.5372549 , 0.4745098])

Upvotes: 0

Views: 50

Answers (1)

Divakar
Divakar

Reputation: 221614

The only obstacle to vectorization seemed to be that IF conditional in avg. To get past it, simply use the choosing capability of np.where and thus have our solution like so -

avgs = (im + image_2)/2.0
image_2_out = np.where((im == 0).any(-1,keepdims=1), image_2, avgs)

Note that this assumes with if [0., 0., 0.] in a1, you meant to check for ANY one match. If you meant to check for ALL zeros, simply use .all instead of .any.

Alternatively, to do in-situ edits in image_2, use a mask for boolean-indexing -

mask = ~(im == 0).any(-1)
image_2[mask] = avgs[mask]

Upvotes: 1

Related Questions