Package: nlm-torch
torch_nlm- memory efficient non-local means on PyTorchContextNon-local means takes a long time compute as it is quadratic for the number of pixels in an image. For simple images this is OK. For larger or three-dimensional (common in medical imaging) images this is impeditive for its application. Here, I introduce a PyTorch-based solution which uses convolutions to extract neighbours (non-local means here does not use the complete image but rather a neighbourhood with $nn$ pixels or $nn*n$ voxels) and calculates the non-local means average. By porting this to PyTorch we can make easy use of the very efficient GPU parallelization and speed up what oftentimes is a very time consuming algorithm.When should you use it?You want to run NLM for small images:just usescikit-imageYou want to run NLM for bigger images AND you have a GPU:use thisYou want to run NLM for relatively big images and you DO NOT have a GPU:good luckBenchmarkI only benchmarkedtorch_nlmagainstscikit-imagein 2d because the latter is prohibitively slow in 3d. Results below.Brief explanation of non-local means (NLM)For an image $I \in \mathbb{R}^{h \times w}$ consider pixel $I_{i,j}$ with coordinates $i,j$.To obtain the non-local mean of this pixel:$$\frac{1}{W}\sum{}^{h,w}{a,b=1} w(I{i,j},I_{a,b}) * I_{a,b}$$where $w(I_{i,j},I_{a,b})$ is the weight of pixel $I_{a,b}$ given $I_{i,j}$ and $W=\sum{}^{h,w}{a,b=1} w(I{i,j},I_{a,b})$.In other words, the non-local means of a given pixel is the weighted average of all pixels. Weights, here, are calculated as the $w(I_{i,j},I_{a,b}) = \exp(- \frac{(I_{i,j} - I_{a,b})^2}{h^2})$, where $h$ is a constant (akin to a standard deviation). To make computation tractable, rather than looping over all pixels, a simple solution is to restrict the neighbourhood to a small square surrounding the image - this is the solution used here.UsageTo use this package all you have to do is clone and install this (apyproject.tomlis provided so that you can easily install this withpoetry). Alternatively, userequirements.txtwithpip(i.e.pip install -r requirements.txt).Two main functions are exported:nlm2dandnlm3d, which are aliases for the most efficienttorch-based NLMM versions (apply_nonlocal_means_2d_mem_efficientandapply_nonlocal_means_3d_mem_efficient), respectively. So if you want to apply it to your favourite image and have a CUDA compatible GPU:importtorch# necessary for obvious reasonsfromtorch_nlmimportnlm2dimage=...# here you define your imageimage_torch=torch.as_tensor(image).to("cuda")image_nlm=nlm2d(image_torch,# the imagekernel_size=11,# neighbourhood sizestd=1.0,# the sigmakernel_size_mean=3,# the kernel used to compute the average pixel intensitysub_filter_size=32# how many neighbourhoods are computed per iteration)sub_filter_sizeis what allows large neighbourhoods - given that users may have relatively small GPU cards, they may opt for smallersub_filter_sizeswhich will enable them to load much smaller sets of neighbourhoods for distance/weight calculations. You may want to run a few tests to figure out the bestsub_filter_sizebefore deploying thisen masse.Implementation detailsThis code was optimized for speed. Three main functions are provided here -apply_nonlocal_means_2d,apply_windowed_nonlocal_means_2dandapply_nonlocal_means_2d_mem_efficient. The first two are development versions, the latter is the one you should use (exposed asnlm_2d).apply_nonlocal_means_2dRetrieves all neighbours as a large tensor and calculates the NLM of the image.Problems:large neighbourhoods will lead to OOMapply_windowed_nonlocal_means_2dDoes the same asapply_nonlocal_means_2dbut uses strided patches to do this, thus reducing memory requirements.Problems:leads to visible striding patch artifactsapply_nonlocal_means_2d_mem_efficientDoes the same asapply_nonlocal_means_2dbut loops over sets of neighbourhoods to calculate weights. $W$ is updated as an accummulator. This version requires defining a "batch size" analogue so that only a few neighbours are calculated at a time. This allows for massive neighbourhoods to be computed in a very parallel fashion.Problems:none for now! But time is a teacher to us all.Generalising to 3DThe good aspect of this is that it requires very little effort to generalise these functions to 3D. These are made available with the same names as above but replacing2dwith3d. The version you want to use isnlm_3d.Expected runtimeFor a large image such asassets/threegorges-1024x1127.jpg(source; size: 1024x1127),apply_nonlocal_means_2d_mem_efficienttakes ~3-5 seconds with a neighbourhood with 51x51 pixels when running on GPU.Example below (obtained by runningpython test.py assets/threegorges-1024x1127.jpg):First panel: original imageSecond panel: original image + noiseThird panel: original image + noise + NLMFourth panel: difference between original image and original image + noise + NLMNote on benchmarking:while 2d benchmarks are reasonable, 3d benchmarks will take a lot of time because ofscikit-image's implementation. Expect times of ~4,000 seconds for a $256 \times 256 \times 256$ images with a neighbourhood size of 17 (torch_nlmran in ~70-80 seconds 😊). You will needscikit-imagefor benchmarking.