import numpy as np
channels_last = np.random.random((150, 150, 40))
channels_first = np.random.random((40, 150, 150))
def gradient_channels_first(image_data):
g = np.empty((image_data.shape[0] * 2,) + image_data.shape[1:])
for i in xrange(image_data.shape[0]):
gx, gy = np.gradient(image_data[i])
g[2 * i] = gx
g[2 * i + 1] = gy
return g
def gradient_channels_last(image_data):
g = np.empty(image_data.shape[:2] + (image_data.shape[2] * 2,))
for i in xrange(image_data.shape[2]):
gx, gy = np.gradient(image_data[..., i])
g[..., 2 * i] = gx
g[..., 2 * i + 1] = gy
return g
%timeit gradient_channels_first(channels_first)
100 loops, best of 3: 13 ms per loop
%timeit gradient_channels_last(channels_last)
10 loops, best of 3: 22.7 ms per loop