merge mater

CONTRIBUTING.md

@@ -32,6 +32,8 @@ _Pull request early_
 
 We encourage you to create pull requests early. It helps us track the contributions under development, whether they are ready to be merged or not. Change your pull request's title to begin with `[WIP]` until it is ready for formal review.
 
+Please note that, as per PyTorch, MONAI uses American English spelling. This means classes and variables should be: normali**z**e, visuali**z**e, colo~~u~~r, etc.
+
 ### Preparing pull requests
 To ensure the code quality, MONAI relies on several linting tools ([flake8 and its plugins](https://gitlab.com/pycqa/flake8), [black](https://github.com/psf/black), [isort](https://github.com/timothycrosley/isort)),
 static type analysis tools ([mypy](https://github.com/python/mypy), [pytype](https://github.com/google/pytype)), as well as a set of unit/integration tests.

docs/source/networks.rst

@@ -183,6 +183,11 @@ Layers
 ~~~~~~~~~~~~~~~~
 .. autoclass:: GaussianFilter
     :members:
+
+`BilateralFilter`
+~~~~~~~~~~~~~~~~~
+.. autoclass:: BilateralFilter
+    :members:
 
 `HilbertTransform`
 ~~~~~~~~~~~~~~~~~~

monai/csrc/ext.cpp

@@ -12,11 +12,16 @@ limitations under the License.
 */
 
 #include <torch/extension.h>
+
+#include "filtering/filtering.h"
 #include "lltm/lltm.h"
 #include "resample/pushpull.h"
 #include "utils/resample_utils.h"
 
 PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
+  // filtering
+  m.def("bilateral_filter", &BilateralFilter, "Bilateral Filter");
+
   // lltm
   m.def("lltm_forward", &lltm_forward, "LLTM forward");
   m.def("lltm_backward", &lltm_backward, "LLTM backward");

monai/csrc/filtering/bilateral/bilateral.h

@@ -0,0 +1,42 @@
+/*
+Copyright 2020 MONAI Consortium
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+    http://www.apache.org/licenses/LICENSE-2.0
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+*/
+
+#pragma once
+
+#include <torch/extension.h>
+#include "utils/common_utils.h"
+
+torch::Tensor BilateralFilterCpu(torch::Tensor input, float spatial_sigma, float color_sigma);
+torch::Tensor BilateralFilterPHLCpu(torch::Tensor input, float spatial_sigma, float color_sigma);
+
+#ifdef WITH_CUDA
+torch::Tensor BilateralFilterCuda(torch::Tensor input, float spatial_sigma, float color_sigma);
+torch::Tensor BilateralFilterPHLCuda(torch::Tensor input, float spatial_sigma, float color_sigma);
+#endif
+
+torch::Tensor BilateralFilter(torch::Tensor input, float spatial_sigma, float color_sigma, bool usePHL) {
+  torch::Tensor (*filterFunction)(torch::Tensor, float, float);
+
+#ifdef WITH_CUDA
+  if (torch::cuda::is_available() && input.is_cuda()) {
+    CHECK_CONTIGUOUS_CUDA(input);
+    filterFunction = usePHL ? &BilateralFilterPHLCuda : &BilateralFilterCuda;
+  } else {
+    filterFunction = usePHL ? &BilateralFilterPHLCpu : &BilateralFilterCpu;
+  }
+#else
+  filterFunction = usePHL ? &BilateralFilterPHLCpu : &BilateralFilterCpu;
+#endif
+
+  return filterFunction(input, spatial_sigma, color_sigma);
+}

monai/csrc/filtering/bilateral/bilateralfilter_cpu.cpp

@@ -0,0 +1,167 @@
+/*
+Copyright 2020 MONAI Consortium
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+    http://www.apache.org/licenses/LICENSE-2.0
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+*/
+
+#include <math.h>
+#include <torch/extension.h>
+
+#include "utils/tensor_description.h"
+
+struct Indexer {
+ public:
+  Indexer(int dimensions, int* sizes) {
+    m_dimensions = dimensions;
+    m_sizes = sizes;
+    m_index = new int[dimensions]{0};
+  }
+
+  bool operator++(int) {
+    for (int i = 0; i < m_dimensions; i++) {
+      m_index[i] += 1;
+
+      if (m_index[i] < m_sizes[i]) {
+        return true;
+      } else {
+        m_index[i] = 0;
+      }
+    }
+
+    return false;
+  }
+
+  int& operator[](int dimensionIndex) {
+    return m_index[dimensionIndex];
+  }
+
+ private:
+  int m_dimensions;
+  int* m_sizes;
+  int* m_index;
+};
+
+template <typename scalar_t>
+void BilateralFilterCpu(torch::Tensor inputTensor, torch::Tensor outputTensor, float spatialSigma, float colorSigma) {
+  // Getting tensor description.
+  TensorDescription desc = TensorDescription(inputTensor);
+
+  // Raw tensor data pointers.
+  scalar_t* inputTensorData = inputTensor.data_ptr<scalar_t>();
+  scalar_t* outputTensorData = outputTensor.data_ptr<scalar_t>();
+
+  // Pre-calculate common values
+  int windowSize = (int)ceil(5.0f * spatialSigma) | 1; // ORing last bit to ensure odd window size
+  int halfWindowSize = floor(0.5f * windowSize);
+  scalar_t spatialExpConstant = -1.0f / (2 * spatialSigma * spatialSigma);
+  scalar_t colorExpConstant = -1.0f / (2 * colorSigma * colorSigma);
+
+  // Kernel sizes.
+  int* kernelSizes = new int[desc.dimensions];
+
+  for (int i = 0; i < desc.dimensions; i++) {
+    kernelSizes[i] = windowSize;
+  }
+
+  // Pre-calculate gaussian kernel in 1D.
+  scalar_t* gaussianKernel = new scalar_t[windowSize];
+
+  for (int i = 0; i < windowSize; i++) {
+    int distance = i - halfWindowSize;
+    gaussianKernel[i] = exp(distance * distance * spatialExpConstant);
+  }
+
+  // Kernel aggregates used to calculate
+  // the output value.
+  scalar_t* valueSum = new scalar_t[desc.channelCount];
+  scalar_t weightSum = 0;
+
+  // Looping over the batches
+  for (int b = 0; b < desc.batchCount; b++) {
+    int batchOffset = b * desc.batchStride;
+
+    // Looping over all dimensions for the home element
+    Indexer homeIndex = Indexer(desc.dimensions, desc.sizes);
+    do // while(homeIndex++)
+    {
+      // Calculating indexing offset for the home element
+      int homeOffset = batchOffset;
+
+      for (int i = 0; i < desc.dimensions; i++) {
+        homeOffset += homeIndex[i] * desc.strides[i];
+      }
+
+      // Zero kernel aggregates.
+      for (int i = 0; i < desc.channelCount; i++) {
+        valueSum[i] = 0;
+      }
+
+      weightSum = 0.0f;
+
+      // Looping over all dimensions for the neighbour element
+      Indexer kernelIndex = Indexer(desc.dimensions, kernelSizes);
+      do // while(kernelIndex++)
+      {
+        // Calculating buffer offset for the neighbour element
+        // Index is clamped to the border in each dimension.
+        int neighbourOffset = batchOffset;
+
+        for (int i = 0; i < desc.dimensions; i++) {
+          int neighbourIndex = homeIndex[i] + kernelIndex[i] - halfWindowSize;
+          int neighbourIndexClamped = std::min(desc.sizes[i] - 1, std::max(0, neighbourIndex));
+          neighbourOffset += neighbourIndexClamped * desc.strides[i];
+        }
+
+        // Euclidean color distance.
+        scalar_t colorDistanceSquared = 0;
+
+        for (int i = 0; i < desc.channelCount; i++) {
+          scalar_t diff = inputTensorData[homeOffset + i * desc.channelStride] -
+              inputTensorData[neighbourOffset + i * desc.channelStride];
+          colorDistanceSquared += diff * diff;
+        }
+
+        // Calculating and combining the spatial
+        // and color weights.
+        scalar_t spatialWeight = 1;
+
+        for (int i = 0; i < desc.dimensions; i++) {
+          spatialWeight *= gaussianKernel[kernelIndex[i]];
+        }
+
+        scalar_t colorWeight = exp(colorDistanceSquared * colorExpConstant);
+        scalar_t totalWeight = spatialWeight * colorWeight;
+
+        // Aggregating values.
+        for (int i = 0; i < desc.channelCount; i++) {
+          valueSum[i] += inputTensorData[neighbourOffset + i * desc.channelStride] * totalWeight;
+        }
+
+        weightSum += totalWeight;
+      } while (kernelIndex++);
+
+      for (int i = 0; i < desc.channelCount; i++) {
+        outputTensorData[homeOffset + i * desc.channelStride] = valueSum[i] / weightSum;
+      }
+    } while (homeIndex++);
+  }
+}
+
+torch::Tensor BilateralFilterCpu(torch::Tensor inputTensor, float spatialSigma, float colorSigma) {
+  // Preparing output tensor.
+  torch::Tensor outputTensor = torch::zeros_like(inputTensor);
+
+  AT_DISPATCH_FLOATING_TYPES_AND_HALF(inputTensor.type(), "BilateralFilterCpu", ([&] {
+                                        BilateralFilterCpu<scalar_t>(
+                                            inputTensor, outputTensor, spatialSigma, colorSigma);
+                                      }));
+
+  return outputTensor;
+}

monai/csrc/filtering/bilateral/bilateralfilter_cpu_phl.cpp

@@ -0,0 +1,89 @@
+/*
+Copyright 2020 MONAI Consortium
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+    http://www.apache.org/licenses/LICENSE-2.0
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+*/
+
+#include <torch/extension.h>
+
+#include "filtering/permutohedral/permutohedral.h"
+#include "utils/tensor_description.h"
+
+template <typename scalar_t>
+void BilateralFilterPHLCpu(
+    torch::Tensor inputTensor,
+    torch::Tensor outputTensor,
+    float spatialSigma,
+    float colorSigma) {
+  // Getting tensor description.
+  TensorDescription desc = TensorDescription(inputTensor);
+
+  int featureChannels = desc.channelCount + desc.dimensions;
+
+  // Preparing memory
+  scalar_t* inputTensorData = inputTensor.data_ptr<scalar_t>();
+  scalar_t* outputTensorData = outputTensor.data_ptr<scalar_t>();
+  scalar_t* data = new scalar_t[desc.channelStride * desc.channelCount];
+  scalar_t* features = new scalar_t[desc.channelStride * featureChannels];
+
+  // Precalculating inverse sigmas
+  float invSpatialSigma = 1.0f / spatialSigma;
+  float invColorSigma = 1.0f / colorSigma;
+
+  // Looping over batches
+  for (int b = 0; b < desc.batchCount; b++) {
+    int batchOffset = b * desc.batchStride;
+
+    // Creating features (also permuting input data to be channel last. Permutohedral
+    // implementation should be changed to channel first to avoid this)
+    for (int i = 0; i < desc.channelStride; i++) {
+      // Color features (and permutation)
+      for (int c = 0; c < desc.channelCount; c++) {
+        features[i * featureChannels + c] = invColorSigma * inputTensorData[batchOffset + i + c * desc.channelStride];
+        data[i * desc.channelCount + c] = inputTensorData[batchOffset + i + c * desc.channelStride];
+      }
+
+      // Spatial features
+      int offsetRemanider = i;
+
+      for (int d = 0; d < desc.dimensions; d++) {
+        int coord = offsetRemanider / desc.strides[d];
+        offsetRemanider -= coord * desc.strides[d];
+
+        features[i * featureChannels + desc.channelCount + d] = invSpatialSigma * coord;
+      }
+    }
+
+    // Filtering data with respect to the features.
+    scalar_t* output =
+        PermutohedralCPU<scalar_t>(data, features, desc.channelCount, featureChannels, desc.channelStride);
+
+    // Writing output tensor.
+    for (int i = 0; i < desc.channelStride; i++) {
+      for (int c = 0; c < desc.channelCount; c++) {
+        outputTensorData[batchOffset + i + c * desc.channelStride] = output[i * desc.channelCount + c];
+      }
+    }
+  }
+
+  delete[] data;
+  delete[] features;
+}
+
+// Function to choose template implementation based on dynamic, channels and dimensions
+torch::Tensor BilateralFilterPHLCpu(torch::Tensor inputTensor, float spatialSigma, float colorSigma) {
+  torch::Tensor outputTensor = torch::zeros_like(inputTensor);
+
+  AT_DISPATCH_FLOATING_TYPES(inputTensor.type(), "BilateralFilterPhlCpu", ([&] {
+                               BilateralFilterPHLCpu<scalar_t>(inputTensor, outputTensor, spatialSigma, colorSigma);
+                             }));
+
+  return outputTensor;
+}