Add Median Filter 2D

[santiago-imelio]

I was writing this notebook to learn Nx.Defn and thought that with some improvements it could potentially belong here :)

Would love to get feedback on how to improve my implementation. Once I have some feedback I can continue with documentation and tests.

Here is the result from applying to a noisy image:

Original

Median filter applied (kernel_size = 5)

[polvalente]

Suggested change

	{_, _, _, _, _, result} =
	while {i = 0, t, kernel_tensor, height, width, output}, i < height do
	{result, _} =
	while {output, {i = 0, t, kernel_tensor, height, width}}, i < height do

Refactor the loops to follow this kind of pattern as it's easier to match on and brings the output to the front

[polvalente]

Actually, here's an idea for a refactor using vectorization.

Let's use hardcoded values for simplicity:

window_shape = {3, 3}
tensor = Nx.iota({5, 5})

# first, let's build out a base indices tensor. these can be used with gather to access each element in the tensor
idx = Nx.stack([Nx.iota(tensor.shape, axis: 0), Nx.iota(tensor.shape, axis: 1)], axis: -1) |> Nx.reshape({:auto, 2}) |> Nx.vectorize(:elements)

# then, we can build an offset mask for each window:

padding_x = div(elem(window_shape, 0), 2)
padding_y = div(elem(window_shape, 1), 2)
window_offsets = Nx.stack([
  Nx.iota(window_shape, axis: 0) |> Nx.subtract(padding_x),
  Nx.iota(window_shape, axis: 1)  |> Nx.subtract(padding_y)
], axis: -1) |> Nx.reshape({:auto, 2}) |> Nx.vectorize(:offsets)

# because we're vectorizing with different names, when we operate on the tensors, we'll get
# a cross-operation between each element of one with all of the others. This can be used
# to build each median window.

idx = Nx.add(idx, window_offsets) |> Nx.add(Nx.tensor([padding_x, padding_y]))

tensor
|> Nx.pad(0, [{padding_x, padding_x, 0}, {padding_y, padding_y, 0}])
|> Nx.gather(idx) # the tensor here is now vectorized {elements: 25, offsets: 9} with scalar elements
|> Nx.devectorize(keep_names: false)
|> Nx.median(axis: 1)
|> Nx.reshape(tensor.shape)

I haven't cross-checked results, this might be wrong

[polvalente]

@santiago-imelio did you get around to trying the vectorized version out?

[polvalente]

Instead of Nx.pad the boundaries can be handled using different functions (e.g.Nx.reflect or custom implementations) to handle mode. I wouldn't worry about this in this PR, though.

[santiago-imelio]

I haven't yet, first I would like to understand the concept of vectorization and play with it, need some time for that. Also I'm somewhat of a noob to scientific programming so bare with me haha.

[polvalente]

No worries. I think there is a guide on vectorization in the main nx repo

[santiago-imelio]

@polvalente I was able to use vectorization, although it's somwhat different from your suggestion. Seems to be working fine with the same examples (with minor discrepancies). Let me know if the approach is correct.

[santiago-imelio]

I think the approach I'm using for getting the window using Nx.slice will raise some problems if we want to implement different strategies for handling boundaries like scipy does

https://docs.scipy.org/doc/scipy/reference/generated/scipy.ndimage.median_filter.html#scipy.ndimage.median_filter

[polvalente]

Nice! This is even better than what I originally had in mind!

[polvalente]

I think this approach should be easily extendable to n dimensions, actually!