weighted_median.py

import numpy as np
import numba


@numba.jit(nopython=True, nogil=True, fastmath=True)
def local_median(data, weights, kernel):
    return median(data * weights * kernel)


@numba.jit(nopython=True, nogil=True, fastmath=True, parallel=True)
def weighted_median_filter(data, weights, kernel, out):

    # half length of kernel (rounded down)
    kh = [int(s / 2.) for s in kernel.shape]

    for i in range(kh[0], data.shape[0] - kh[0]):
        for j in range(kh[1], data.shape[1] - kh[1]):
            for k in range(kh[2], data.shape[2] - kh[2]):

                i_low, i_high = i - kh[0], i + kh[0]
                j_low, j_high = j - kh[1], j + kh[1]
                k_low, k_high = k - kh[2], k + kh[2]

                local_data = data[i_low:i_high, j_low:j_high, k_low:k_high]
                local_weights = weights[i_low:i_high, j_low:j_high, k_low:k_high]
                local_kernel = kernel[i_low:i_high, j_low:j_high, k_low:k_high]

                out[i, j, k] = median(local_data * local_weights)



data = rand(100, 100, 100)
kernel = ones((5, 5, 5))
weights = ones_like(data)
out = zeros_like(data)

weighted_median_filter(data, weights, kernel, out)