fast_binned_median_numpy.py

import numpy as np
import time

kernel_cols = 2
kernel_rows = 3
rows = 180
cols = 160
data = np.array(np.arange(0, cols * rows).reshape((rows, cols)))
print('DATA')
print(data)
print('')

t0 = time.time()
medians = np.zeros((rows // kernel_rows, cols // kernel_cols))
for y in range(rows // kernel_rows):
    for x in range(cols // kernel_cols):
        array = []
        for kx in range(kernel_cols):
            for ky in range(kernel_rows):
                array.append(data[y * kernel_rows + ky, x * kernel_cols + kx])
        medians[y, x] = np.median(array)
t1 = time.time()
print('MEDIAN REFERENCE')
print(medians)
print('')

print('SPLIT')
t2 = time.time()
array1 = np.hsplit(data, [(i + 1) * kernel_cols for i in range(cols // kernel_cols - 1)])
array2 = np.concatenate(array1).reshape(rows * cols // kernel_cols // kernel_rows, kernel_cols * kernel_rows)
array3 = np.median(array2, axis=1).reshape(cols // kernel_cols, rows // kernel_rows)
fast_binned_medians = array3.transpose()
t3 = time.time()
print(fast_binned_medians)

print(t1 - t0)
print(t3 - t2)

assert np.min(fast_binned_medians == medians)