Repository: Project-MONAI/MONAI

Issue:
Fully differentiable and trainable Joint Bilateral Filter (TJBF) layer
**Is your feature request related to a problem? Please describe.**
The [joint bilateral filter (JBF)](https://doi.org/10.1145/1015706.1015777) is similar to the bilateral filter but calculates its intensity range kernel on a second input image -- a so-called guidance image. This allows to integrate more information into the filter process and has slightly different applications than the usual bilateral filter (because you need a second guidance image for the filtering). MONAI currently does not contain a  regular or differentiable JBF.

**Describe the solution you'd like**
I implemented a fully differentiable and trainable version of the JBF based on an analytical filter gradient (C++/CUDA). You can find the code [here](https://github.com/faebstn96/trainable-joint-bilateral-filter-source). We additionally published a [paper (Scientific Reports)](https://doi.org/10.1038/s41598-022-22530-4) presenting and using that trainable joint bilateral filter for CT denoising.
The structure of the code is very similar to #1840 but contains the JBF algorithm with gradients with respect to its input, its guidance image, and its filter parameters. My GitHub also contains a gradcheck script to compare the analytical gradient to a numerical gradient based on `torch.autograd.gradcheck`.

I am happy to create a pull request similar to #5757 and include the trainable JBF as an additional filter layer. Please let me know what you think and add any questions or comments!


Discussion:

Resolving patch:
diff --git a/monai/networks/layers/__init__.py b/monai/networks/layers/__init__.py
--- a/monai/networks/layers/__init__.py
+++ b/monai/networks/layers/__init__.py
@@ -14,7 +14,7 @@
 from .convutils import calculate_out_shape, gaussian_1d, polyval, same_padding, stride_minus_kernel_padding
 from .drop_path import DropPath
 from .factories import Act, Conv, Dropout, LayerFactory, Norm, Pad, Pool, split_args
-from .filtering import BilateralFilter, PHLFilter, TrainableBilateralFilter
+from .filtering import BilateralFilter, PHLFilter, TrainableBilateralFilter, TrainableJointBilateralFilter
 from .gmm import GaussianMixtureModel
 from .simplelayers import (
     LLTM,
diff --git a/monai/networks/layers/filtering.py b/monai/networks/layers/filtering.py
--- a/monai/networks/layers/filtering.py
+++ b/monai/networks/layers/filtering.py
@@ -17,7 +17,7 @@
 
 _C, _ = optional_import("monai._C")
 
-__all__ = ["BilateralFilter", "PHLFilter", "TrainableBilateralFilter"]
+__all__ = ["BilateralFilter", "PHLFilter", "TrainableBilateralFilter", "TrainableJointBilateralFilter"]
 
 
 class BilateralFilter(torch.autograd.Function):
@@ -33,10 +33,10 @@ class BilateralFilter(torch.autograd.Function):
     Args:
         input: input tensor.
 
-        spatial sigma: the standard deviation of the spatial blur. Higher values can
+        spatial_sigma: the standard deviation of the spatial blur. Higher values can
             hurt performance when not using the approximate method (see fast approx).
 
-        color sigma: the standard deviation of the color blur. Lower values preserve
+        color_sigma: the standard deviation of the color blur. Lower values preserve
             edges better whilst higher values tend to a simple gaussian spatial blur.
 
         fast approx: This flag chooses between two implementations. The approximate method may
@@ -200,45 +200,246 @@ class TrainableBilateralFilter(torch.nn.Module):
     Args:
         input: input tensor to be filtered.
 
+        spatial_sigma: tuple (sigma_x, sigma_y, sigma_z) initializing the trainable standard
+            deviations of the spatial filter kernels. Tuple length must equal the number of
+            spatial input dimensions.
+
+        color_sigma: trainable standard deviation of the intensity range kernel. This filter
+            parameter determines the degree of edge preservation.
+
+    Returns:
+        output (torch.Tensor): filtered tensor.
+    """
+
+    def __init__(self, spatial_sigma, color_sigma):
+        super().__init__()
+
+        if isinstance(spatial_sigma, float):
+            spatial_sigma = [spatial_sigma, spatial_sigma, spatial_sigma]
+            self.len_spatial_sigma = 3
+        elif len(spatial_sigma) == 1:
+            spatial_sigma = [spatial_sigma[0], 0.01, 0.01]
+            self.len_spatial_sigma = 1
+        elif len(spatial_sigma) == 2:
+            spatial_sigma = [spatial_sigma[0], spatial_sigma[1], 0.01]
+            self.len_spatial_sigma = 2
+        elif len(spatial_sigma) == 3:
+            spatial_sigma = [spatial_sigma[0], spatial_sigma[1], spatial_sigma[2]]
+            self.len_spatial_sigma = 3
+        else:
+            raise ValueError(
+                f"len(spatial_sigma) {spatial_sigma} must match number of spatial dims {self.ken_spatial_sigma}."
+            )
+
+        # Register sigmas as trainable parameters.
+        self.sigma_x = torch.nn.Parameter(torch.tensor(spatial_sigma[0]))
+        self.sigma_y = torch.nn.Parameter(torch.tensor(spatial_sigma[1]))
+        self.sigma_z = torch.nn.Parameter(torch.tensor(spatial_sigma[2]))
+        self.sigma_color = torch.nn.Parameter(torch.tensor(color_sigma))
+
+    def forward(self, input_tensor):
+        if input_tensor.shape[1] != 1:
+            raise ValueError(
+                f"Currently channel dimensions >1 ({input_tensor.shape[1]}) are not supported. "
+                "Please use multiple parallel filter layers if you want "
+                "to filter multiple channels."
+            )
+
+        len_input = len(input_tensor.shape)
+
+        # C++ extension so far only supports 5-dim inputs.
+        if len_input == 3:
+            input_tensor = input_tensor.unsqueeze(3).unsqueeze(4)
+        elif len_input == 4:
+            input_tensor = input_tensor.unsqueeze(4)
+
+        if self.len_spatial_sigma != len_input:
+            raise ValueError(f"Spatial dimension ({len_input}) must match initialized len(spatial_sigma).")
+
+        prediction = TrainableBilateralFilterFunction.apply(
+            input_tensor, self.sigma_x, self.sigma_y, self.sigma_z, self.sigma_color
+        )
+
+        # Make sure to return tensor of the same shape as the input.
+        if len_input == 3:
+            prediction = prediction.squeeze(4).squeeze(3)
+        elif len_input == 4:
+            prediction = prediction.squeeze(4)
+
+        return prediction
+
+
+class TrainableJointBilateralFilterFunction(torch.autograd.Function):
+    """
+    torch.autograd.Function for the TrainableJointBilateralFilter layer.
+
+    See:
+        F. Wagner, et al., Trainable joint bilateral filters for enhanced prediction stability in
+        low-dose CT, Scientific Reports (2022), https://doi.org/10.1038/s41598-022-22530-4
+
+    Args:
+        input: input tensor to be filtered.
+
+        guide: guidance image tensor to be used during filtering.
+
         sigma x: trainable standard deviation of the spatial filter kernel in x direction.
 
         sigma y: trainable standard deviation of the spatial filter kernel in y direction.
 
         sigma z: trainable standard deviation of the spatial filter kernel in z direction.
 
-        sigma color: trainable standard deviation of the intensity range kernel. This filter
+        color sigma: trainable standard deviation of the intensity range kernel. This filter
             parameter determines the degree of edge preservation.
 
     Returns:
         output (torch.Tensor): filtered tensor.
     """
 
-    def __init__(self, sigma_x, sigma_y, sigma_z, sigma_color):
-        super().__init__()
+    @staticmethod
+    def forward(ctx, input_img, guidance_img, sigma_x, sigma_y, sigma_z, color_sigma):
+        output_tensor, output_weights_tensor, do_dx_ki, do_dsig_r, do_dsig_x, do_dsig_y, do_dsig_z = _C.tjbf_forward(
+            input_img, guidance_img, sigma_x, sigma_y, sigma_z, color_sigma
+        )
 
-        # Register sigmas as trainable parameters.
-        self.sigma_x = torch.nn.Parameter(torch.tensor(sigma_x))
-        self.sigma_y = torch.nn.Parameter(torch.tensor(sigma_y))
-        self.sigma_z = torch.nn.Parameter(torch.tensor(sigma_z))
-        self.sigma_color = torch.nn.Parameter(torch.tensor(sigma_color))
+        ctx.save_for_backward(
+            input_img,
+            sigma_x,
+            sigma_y,
+            sigma_z,
+            color_sigma,
+            output_tensor,
+            output_weights_tensor,
+            do_dx_ki,
+            do_dsig_r,
+            do_dsig_x,
+            do_dsig_y,
+            do_dsig_z,
+            guidance_img,
+        )
 
-    def forward(self, input_tensor):
-        assert input_tensor.shape[1] == 1, (
-            "Currently channel dimensions >1 are not supported. "
-            "Please use multiple parallel filter layers if you want "
-            "to filter multiple channels."
+        return output_tensor
+
+    @staticmethod
+    def backward(ctx, grad_output):
+        input_img = ctx.saved_tensors[0]  # input image
+        sigma_x = ctx.saved_tensors[1]
+        sigma_y = ctx.saved_tensors[2]
+        sigma_z = ctx.saved_tensors[3]
+        color_sigma = ctx.saved_tensors[4]
+        output_tensor = ctx.saved_tensors[5]  # filtered image
+        output_weights_tensor = ctx.saved_tensors[6]  # weights
+        do_dx_ki = ctx.saved_tensors[7]  # derivative of output with respect to input, while k==i
+        do_dsig_r = ctx.saved_tensors[8]  # derivative of output with respect to range sigma
+        do_dsig_x = ctx.saved_tensors[9]  # derivative of output with respect to sigma x
+        do_dsig_y = ctx.saved_tensors[10]  # derivative of output with respect to sigma y
+        do_dsig_z = ctx.saved_tensors[11]  # derivative of output with respect to sigma z
+        guidance_img = ctx.saved_tensors[12]  # guidance image
+
+        # calculate gradient with respect to the sigmas
+        grad_color_sigma = torch.sum(grad_output * do_dsig_r)
+        grad_sig_x = torch.sum(grad_output * do_dsig_x)
+        grad_sig_y = torch.sum(grad_output * do_dsig_y)
+        grad_sig_z = torch.sum(grad_output * do_dsig_z)
+
+        grad_output_tensor, grad_guidance_tensor = _C.tjbf_backward(
+            grad_output,
+            input_img,
+            guidance_img,
+            output_tensor,
+            output_weights_tensor,
+            do_dx_ki,
+            sigma_x,
+            sigma_y,
+            sigma_z,
+            color_sigma,
         )
 
+        return grad_output_tensor, grad_guidance_tensor, grad_sig_x, grad_sig_y, grad_sig_z, grad_color_sigma
+
+
+class TrainableJointBilateralFilter(torch.nn.Module):
+    """
+    Implementation of a trainable joint bilateral filter layer as proposed in the corresponding publication.
+    The guidance image is used as additional (edge) information during filtering. All filter parameters and the
+    guidance image can be trained data-driven. The spatial filter kernels x, y, and z determine
+    image smoothing whereas the color parameter specifies the amount of edge preservation.
+    Can run on 1D, 2D, or 3D tensors (on top of Batch and Channel dimensions). Input tensor shape must match
+    guidance tensor shape.
+
+    See:
+        F. Wagner, et al., Trainable joint bilateral filters for enhanced prediction stability in
+        low-dose CT, Scientific Reports (2022), https://doi.org/10.1038/s41598-022-22530-4
+
+    Args:
+        input: input tensor to be filtered.
+
+        guide: guidance image tensor to be used during filtering.
+
+        spatial_sigma: tuple (sigma_x, sigma_y, sigma_z) initializing the trainable standard
+            deviations of the spatial filter kernels. Tuple length must equal the number of
+            spatial input dimensions.
+
+        color_sigma: trainable standard deviation of the intensity range kernel. This filter
+            parameter determines the degree of edge preservation.
+
+    Returns:
+        output (torch.Tensor): filtered tensor.
+    """
+
+    def __init__(self, spatial_sigma, color_sigma):
+        super().__init__()
+
+        if isinstance(spatial_sigma, float):
+            spatial_sigma = [spatial_sigma, spatial_sigma, spatial_sigma]
+            self.len_spatial_sigma = 3
+        elif len(spatial_sigma) == 1:
+            spatial_sigma = [spatial_sigma[0], 0.01, 0.01]
+            self.len_spatial_sigma = 1
+        elif len(spatial_sigma) == 2:
+            spatial_sigma = [spatial_sigma[0], spatial_sigma[1], 0.01]
+            self.len_spatial_sigma = 2
+        elif len(spatial_sigma) == 3:
+            spatial_sigma = [spatial_sigma[0], spatial_sigma[1], spatial_sigma[2]]
+            self.len_spatial_sigma = 3
+        else:
+            raise ValueError(
+                f"len(spatial_sigma) {spatial_sigma} must match number of spatial dims {self.ken_spatial_sigma}."
+            )
+
+        # Register sigmas as trainable parameters.
+        self.sigma_x = torch.nn.Parameter(torch.tensor(spatial_sigma[0]))
+        self.sigma_y = torch.nn.Parameter(torch.tensor(spatial_sigma[1]))
+        self.sigma_z = torch.nn.Parameter(torch.tensor(spatial_sigma[2]))
+        self.sigma_color = torch.nn.Parameter(torch.tensor(color_sigma))
+
+    def forward(self, input_tensor, guidance_tensor):
+        if input_tensor.shape[1] != 1:
+            raise ValueError(
+                f"Currently channel dimensions >1 ({input_tensor.shape[1]}) are not supported. "
+                "Please use multiple parallel filter layers if you want "
+                "to filter multiple channels."
+            )
+        if input_tensor.shape != guidance_tensor.shape:
+            raise ValueError(
+                "Shape of input image must equal shape of guidance image."
+                f"Got {input_tensor.shape} and {guidance_tensor.shape}."
+            )
+
         len_input = len(input_tensor.shape)
 
         # C++ extension so far only supports 5-dim inputs.
         if len_input == 3:
             input_tensor = input_tensor.unsqueeze(3).unsqueeze(4)
+            guidance_tensor = guidance_tensor.unsqueeze(3).unsqueeze(4)
         elif len_input == 4:
             input_tensor = input_tensor.unsqueeze(4)
+            guidance_tensor = guidance_tensor.unsqueeze(4)
 
-        prediction = TrainableBilateralFilterFunction.apply(
-            input_tensor, self.sigma_x, self.sigma_y, self.sigma_z, self.sigma_color
+        if self.len_spatial_sigma != len_input:
+            raise ValueError(f"Spatial dimension ({len_input}) must match initialized len(spatial_sigma).")
+
+        prediction = TrainableJointBilateralFilterFunction.apply(
+            input_tensor, guidance_tensor, self.sigma_x, self.sigma_y, self.sigma_z, self.sigma_color
         )
 
         # Make sure to return tensor of the same shape as the input.


Test patch:
diff --git a/tests/test_trainable_joint_bilateral.py b/tests/test_trainable_joint_bilateral.py
new file mode 100644
--- /dev/null
+++ b/tests/test_trainable_joint_bilateral.py
@@ -0,0 +1,612 @@
+# Copyright (c) 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.
+
+from __future__ import annotations
+
+import unittest
+
+import numpy as np
+import torch
+from parameterized import parameterized
+from torch.autograd import gradcheck
+
+from monai.networks.layers.filtering import TrainableJointBilateralFilterFunction
+from tests.utils import skip_if_no_cpp_extension, skip_if_no_cuda
+
+TEST_CASES = [
+    [
+        # Case Description
+        "1 dimension, 1 channel, low spatial sigmas, low color sigma",
+        # (sigma_x, sigma_y, sigma_z, color_sigma)
+        (1.0, 1.0, 1.0, 0.2),
+        # Input
+        [
+            # Batch 0
+            [
+                # Channel 0
+                [1, 0, 0, 0, 1]
+            ],
+            # Batch 1
+            [
+                # Channel 0
+                [0, 0, 1, 0, 0]
+            ],
+        ],
+        # Guide
+        [
+            # Batch 0
+            [
+                # Channel 0
+                [1, 1, 0, 0, 1]
+            ],
+            # Batch 1
+            [
+                # Channel 0
+                [0, 0, 1, 0, 1]
+            ],
+        ],
+        # Expected
+        [
+            # Batch 0
+            [
+                # Channel 0
+                [0.622459, 0.377540, 0.000001, 0.000001, 0.999997]
+            ],
+            # Batch 1
+            [
+                # Channel 0
+                [0.000000, 0.000001, 0.880793, 0.000002, 0.119203]
+            ],
+        ],
+    ],
+    [
+        # Case Description
+        "1 dimension, 1 channel, low spatial sigmas, high color sigma",
+        # (sigma_x, sigma_y, sigma_z, color_sigma)
+        (1, 1, 1, 0.9),
+        # Input
+        [
+            # Batch 0
+            [
+                # Channel 0
+                [1, 0, 0, 0, 1]
+            ],
+            # Batch 1
+            [
+                # Channel 0
+                [0, 0, 1, 0, 0]
+            ],
+        ],
+        # Guide
+        [
+            # Batch 0
+            [
+                # Channel 0
+                [1, 1, 0, 0, 1]
+            ],
+            # Batch 1
+            [
+                # Channel 0
+                [0, 0, 1, 0, 1]
+            ],
+        ],
+        # Expected
+        [
+            # Batch 0
+            [
+                # Channel 0
+                [0.595404, 0.302253, 0.070203, 0.163038, 0.714200]
+            ],
+            # Batch 1
+            [
+                # Channel 0
+                [0.043465, 0.158126, 0.536864, 0.182809, 0.092537]
+            ],
+        ],
+    ],
+    [
+        # Case Description
+        "1 dimension, 1 channel, high spatial sigmas, low color sigma",
+        # (sigma_x, sigma_y, sigma_z, color_sigma)
+        (4, 4, 4, 0.2),
+        # Input
+        [
+            # Batch 0
+            [
+                # Channel 0
+                [1, 0, 0, 0, 1]
+            ],
+            # Batch 1
+            [
+                # Channel 0
+                [0, 0, 1, 0, 0]
+            ],
+        ],
+        # Guide
+        [
+            # Batch 0
+            [
+                # Channel 0
+                [1, 1, 0, 0, 1]
+            ],
+            # Batch 1
+            [
+                # Channel 0
+                [0, 0, 1, 0, 1]
+            ],
+        ],
+        # Expected
+        [
+            # Batch 0
+            [
+                # Channel 0
+                [0.623709, 0.632901, 0.000003, 0.000003, 0.680336]
+            ],
+            # Batch 1
+            [
+                # Channel 0
+                [0.000001, 0.000001, 0.531206, 0.000001, 0.468788]
+            ],
+        ],
+    ],
+    [
+        # Case Description
+        "1 dimension, 1 channel, high spatial sigmas, high color sigma",
+        # (sigma_x, sigma_y, sigma_z, color_sigma)
+        (4, 4, 4, 0.9),
+        # Input
+        [
+            # Batch 0
+            [
+                # Channel 0
+                [1, 0, 0, 0, 1]
+            ],
+            # Batch 1
+            [
+                # Channel 0
+                [0, 0, 1, 0, 0]
+            ],
+        ],
+        # Guide
+        [
+            # Batch 0
+            [
+                # Channel 0
+                [1, 1, 0, 0, 1]
+            ],
+            # Batch 1
+            [
+                # Channel 0
+                [0, 0, 1, 0, 1]
+            ],
+        ],
+        # Expected
+        [
+            # Batch 0
+            [
+                # Channel 0
+                [0.464455, 0.463098, 0.276430, 0.275530, 0.478105]
+            ],
+            # Batch 1
+            [
+                # Channel 0
+                [0.134956, 0.138247, 0.293759, 0.141954, 0.281082]
+            ],
+        ],
+    ],
+    [
+        # Case Description
+        "2 dimensions, 1 channel, high spatial sigmas, high color sigma",
+        # (sigma_x, sigma_y, sigma_z, color_sigma)
+        (4, 4, 4, 0.9),
+        # Input
+        [
+            # Batch 0
+            [
+                # Channel 0
+                [[1, 0, 0, 0, 1], [0, 0, 0, 0, 0], [0, 0, 0, 0, 0], [0, 0, 0, 0, 0], [1, 0, 0, 0, 1]]
+            ],
+            # Batch 1
+            [
+                # Channel 0
+                [[0, 0, 0, 0, 0], [0, 0, 0, 0, 0], [0, 0, 1, 0, 0], [0, 0, 0, 0, 0], [0, 0, 0, 0, 0]]
+            ],
+        ],
+        # Guide
+        [
+            # Batch 0
+            [
+                # Channel 0
+                [[1, 1, 0, 0, 1], [0, 0, 0, 1, 0], [1, 1, 0, 0, 0], [0, 0, 0, 0, 0], [0, 1, 0, 0, 1]]
+            ],
+            # Batch 1
+            [
+                # Channel 0
+                [[0, 0, 0, 1, 0], [0, 0, 0, 1, 1], [0, 0, 0, 0, 0], [0, 0, 0, 0, 0], [1, 0, 0, 0, 0]]
+            ],
+        ],
+        # Expected
+        [
+            # Batch 0
+            [
+                # Channel 0
+                [
+                    [0.186535, 0.187357, 0.105377, 0.103652, 0.198665],
+                    [0.112617, 0.108847, 0.105970, 0.189602, 0.102954],
+                    [0.178338, 0.179829, 0.107473, 0.105256, 0.103963],
+                    [0.117651, 0.113304, 0.109876, 0.107392, 0.105853],
+                    [0.121557, 0.177689, 0.113150, 0.110388, 0.192877],
+                ]
+            ],
+            # Batch 1
+            [
+                # Channel 0
+                [
+                    [0.047156, 0.047865, 0.048233, 0.038611, 0.047911],
+                    [0.047607, 0.048292, 0.048633, 0.039251, 0.038611],
+                    [0.047715, 0.048369, 0.048678, 0.048633, 0.048233],
+                    [0.047477, 0.048094, 0.048369, 0.048292, 0.047865],
+                    [0.039190, 0.047477, 0.047715, 0.047607, 0.047156],
+                ]
+            ],
+        ],
+    ],
+    [
+        # Case Description
+        "3 dimensions, 1 channel, high spatial sigmas, high color sigma",
+        # (sigma_x, sigma_y, sigma_z, color_sigma)
+        (4, 4, 4, 0.9),
+        # Input
+        [
+            # Batch 0
+            [
+                # Channel 0
+                [
+                    # Frame 0
+                    [[1, 0, 0, 0, 1], [0, 0, 0, 0, 0], [0, 0, 0, 0, 0], [0, 0, 0, 0, 0], [1, 0, 0, 0, 1]],
+                    # Frame 1
+                    [[0, 0, 0, 0, 0], [0, 0, 0, 0, 0], [0, 0, 0, 0, 0], [0, 0, 0, 0, 0], [0, 0, 0, 0, 0]],
+                    # Frame 2
+                    [[0, 0, 0, 0, 0], [0, 0, 0, 0, 0], [0, 0, 0, 0, 0], [0, 0, 0, 0, 0], [0, 0, 0, 0, 0]],
+                    # Frame 3
+                    [[0, 0, 0, 0, 0], [0, 0, 0, 0, 0], [0, 0, 0, 0, 0], [0, 0, 0, 0, 0], [0, 0, 0, 0, 0]],
+                    # Frame 4
+                    [[1, 0, 0, 0, 1], [0, 0, 0, 0, 0], [0, 0, 0, 0, 0], [0, 0, 0, 0, 0], [1, 0, 0, 0, 1]],
+                ]
+            ]
+        ],
+        # Guide
+        [
+            # Batch 0
+            [
+                # Channel 0
+                [
+                    # Frame 0
+                    [[1, 1, 0, 0, 1], [0, 0, 0, 0, 0], [0, 0, 0, 0, 0], [0, 0, 0, 0, 0], [0, 1, 0, 0, 1]],
+                    # Frame 1
+                    [[0, 0, 0, 0, 0], [0, 0, 0, 0, 0], [0, 0, 0, 0, 0], [0, 0, 0, 0, 0], [0, 0, 0, 0, 0]],
+                    # Frame 2
+                    [[0, 0, 0, 0, 0], [0, 0, 0, 0, 0], [0, 0, 0, 0, 0], [0, 0, 0, 0, 0], [0, 0, 0, 0, 0]],
+                    # Frame 3
+                    [[0, 0, 0, 0, 0], [0, 0, 0, 0, 0], [0, 0, 0, 0, 0], [0, 0, 0, 0, 0], [0, 0, 0, 0, 0]],
+                    # Frame 4
+                    [[1, 1, 0, 1, 1], [0, 0, 0, 0, 0], [0, 0, 0, 0, 0], [0, 0, 0, 0, 0], [1, 0, 0, 1, 1]],
+                ]
+            ]
+        ],
+        # Expected
+        [
+            # Batch 0
+            [
+                # Channel 0
+                [
+                    # Frame 0
+                    [
+                        [0.089316, 0.088903, 0.033707, 0.033461, 0.091881],
+                        [0.035173, 0.034324, 0.033747, 0.033448, 0.033427],
+                        [0.035619, 0.034710, 0.034074, 0.033720, 0.033646],
+                        [0.036364, 0.035387, 0.034688, 0.034275, 0.034148],
+                        [0.037401, 0.085687, 0.035583, 0.035109, 0.089741],
+                    ],
+                    # Frame 1
+                    [
+                        [0.034248, 0.033502, 0.033023, 0.032816, 0.032881],
+                        [0.034339, 0.033546, 0.033020, 0.032767, 0.032785],
+                        [0.034721, 0.033876, 0.033298, 0.032994, 0.032965],
+                        [0.035397, 0.034490, 0.033856, 0.033501, 0.033424],
+                        [0.036357, 0.035383, 0.034688, 0.034279, 0.034155],
+                    ],
+                    # Frame 2
+                    [
+                        [0.033748, 0.033047, 0.032609, 0.032441, 0.032541],
+                        [0.033782, 0.033041, 0.032562, 0.032353, 0.032410],
+                        [0.034104, 0.033316, 0.032792, 0.032538, 0.032554],
+                        [0.034714, 0.033872, 0.033298, 0.032998, 0.032972],
+                        [0.035604, 0.034702, 0.034074, 0.033727, 0.033660],
+                    ],
+                    # Frame 3
+                    [
+                        [0.033533, 0.032871, 0.032471, 0.032340, 0.032476],
+                        [0.033511, 0.032815, 0.032380, 0.032212, 0.032310],
+                        [0.033775, 0.033037, 0.032562, 0.032356, 0.032417],
+                        [0.034324, 0.033539, 0.033020, 0.032774, 0.032799],
+                        [0.035151, 0.034313, 0.033747, 0.033459, 0.033449],
+                    ],
+                    # Frame 4
+                    [
+                        [0.091383, 0.090681, 0.032608, 0.091418, 0.092851],
+                        [0.033525, 0.032867, 0.032471, 0.032344, 0.032483],
+                        [0.033733, 0.033039, 0.032609, 0.032448, 0.032555],
+                        [0.034226, 0.033491, 0.033023, 0.032827, 0.032903],
+                        [0.089445, 0.034216, 0.033707, 0.090126, 0.091748],
+                    ],
+                ]
+            ]
+        ],
+    ],
+]
+
+
+@skip_if_no_cpp_extension
+class JointBilateralFilterTestCaseCpuPrecise(unittest.TestCase):
+    @parameterized.expand(TEST_CASES)
+    def test_cpu_precise(self, test_case_description, sigmas, input, guide, expected):
+
+        # Params to determine the implementation to test
+        device = torch.device("cpu")
+
+        # Create input tensor and apply filter
+        input_tensor = torch.from_numpy(np.array(input)).to(dtype=torch.double, device=device)
+        guide_tensor = torch.from_numpy(np.array(guide)).to(dtype=torch.double, device=device)
+
+        len_input = len(input_tensor.shape)
+        # C++ extension so far only supports 5-dim inputs.
+        if len_input == 3:
+            input_tensor = input_tensor.unsqueeze(3).unsqueeze(4)
+            guide_tensor = guide_tensor.unsqueeze(3).unsqueeze(4)
+        elif len_input == 4:
+            input_tensor = input_tensor.unsqueeze(4)
+            guide_tensor = guide_tensor.unsqueeze(4)
+
+        output = TrainableJointBilateralFilterFunction.apply(input_tensor, guide_tensor, *sigmas).cpu().numpy()
+
+        # Make sure to return tensor of the same shape as the input.
+        if len_input == 3:
+            output = output.squeeze(4).squeeze(3)
+        elif len_input == 4:
+            output = output.squeeze(4)
+
+        # Ensure result are as expected.
+        np.testing.assert_allclose(output, expected, atol=1e-5)
+
+    @parameterized.expand(TEST_CASES)
+    def test_cpu_precise_backwards(self, test_case_description, sigmas, input, guide, expected):
+
+        # Params to determine the implementation to test
+        device = torch.device("cpu")
+
+        # Prepare input tensor
+        input_tensor = torch.from_numpy(np.array(input)).to(dtype=torch.double, device=device)
+        input_tensor.requires_grad = True
+        guide_tensor = torch.from_numpy(np.array(guide)).to(dtype=torch.double, device=device)
+
+        # C++ extension so far only supports 5-dim inputs.
+        len_input = len(input_tensor.shape)
+        if len_input == 3:
+            input_tensor = input_tensor.unsqueeze(3).unsqueeze(4)
+            guide_tensor = guide_tensor.unsqueeze(3).unsqueeze(4)
+        elif len_input == 4:
+            input_tensor = input_tensor.unsqueeze(4)
+            guide_tensor = guide_tensor.unsqueeze(4)
+
+        # Check gradient toward input.
+        args = (
+            input_tensor,
+            guide_tensor,
+            torch.tensor(sigmas[0]),
+            torch.tensor(sigmas[1]),
+            torch.tensor(sigmas[2]),
+            torch.tensor(sigmas[3]),
+        )
+        gradcheck(TrainableJointBilateralFilterFunction.apply, args, eps=1e-6, atol=1e-5, raise_exception=False)
+        input_tensor = input_tensor.detach()
+        input_tensor.requires_grad = False
+
+        # Check gradient toward guide.
+        guide_tensor.requires_grad = True
+        args = (
+            input_tensor,
+            guide_tensor,
+            torch.tensor(sigmas[0]),
+            torch.tensor(sigmas[1]),
+            torch.tensor(sigmas[2]),
+            torch.tensor(sigmas[3]),
+        )
+        gradcheck(TrainableJointBilateralFilterFunction.apply, args, eps=1e-6, atol=1e-5, raise_exception=False)
+        guide_tensor = guide_tensor.detach()
+        guide_tensor.guide_tensor = False
+
+        # Check gradient toward sigma_x.
+        args = (
+            input_tensor,
+            guide_tensor,
+            torch.tensor(sigmas[0], dtype=torch.double, requires_grad=True),
+            torch.tensor(sigmas[1]),
+            torch.tensor(sigmas[2]),
+            torch.tensor(sigmas[3]),
+        )
+        gradcheck(TrainableJointBilateralFilterFunction.apply, args, eps=1e-2, atol=1e-3, raise_exception=False)
+
+        # Check gradient toward sigma_y.
+        args = (
+            input_tensor,
+            guide_tensor,
+            torch.tensor(sigmas[0]),
+            torch.tensor(sigmas[1], dtype=torch.double, requires_grad=True),
+            torch.tensor(sigmas[2]),
+            torch.tensor(sigmas[3]),
+        )
+        gradcheck(TrainableJointBilateralFilterFunction.apply, args, eps=1e-2, atol=1e-3, raise_exception=False)
+
+        # Check gradient toward sigma_z.
+        args = (
+            input_tensor,
+            guide_tensor,
+            torch.tensor(sigmas[0]),
+            torch.tensor(sigmas[1]),
+            torch.tensor(sigmas[2], dtype=torch.double, requires_grad=True),
+            torch.tensor(sigmas[3]),
+        )
+        gradcheck(TrainableJointBilateralFilterFunction.apply, args, eps=1e-2, atol=1e-3, raise_exception=False)
+
+        # Check gradient toward sigma_color.
+        args = (
+            input_tensor,
+            guide_tensor,
+            torch.tensor(sigmas[0]),
+            torch.tensor(sigmas[1]),
+            torch.tensor(sigmas[2]),
+            torch.tensor(sigmas[3], dtype=torch.double, requires_grad=True),
+        )
+        gradcheck(TrainableJointBilateralFilterFunction.apply, args, eps=1e-3, atol=1e-3, raise_exception=False)
+
+
+@skip_if_no_cuda
+@skip_if_no_cpp_extension
+class JointBilateralFilterTestCaseCudaPrecise(unittest.TestCase):
+    @parameterized.expand(TEST_CASES)
+    def test_cuda_precise(self, test_case_description, sigmas, input, guide, expected):
+
+        # Skip this test
+        if not torch.cuda.is_available():
+            return
+
+        # Params to determine the implementation to test
+        device = torch.device("cuda")
+
+        # Create input tensor and apply filter
+        input_tensor = torch.from_numpy(np.array(input)).to(dtype=torch.double, device=device)
+        guide_tensor = torch.from_numpy(np.array(guide)).to(dtype=torch.double, device=device)
+
+        len_input = len(input_tensor.shape)
+        # C++ extension so far only supports 5-dim inputs.
+        if len_input == 3:
+            input_tensor = input_tensor.unsqueeze(3).unsqueeze(4)
+            guide_tensor = guide_tensor.unsqueeze(3).unsqueeze(4)
+        elif len_input == 4:
+            input_tensor = input_tensor.unsqueeze(4)
+            guide_tensor = guide_tensor.unsqueeze(4)
+
+        output = TrainableJointBilateralFilterFunction.apply(input_tensor, guide_tensor, *sigmas).cpu().numpy()
+
+        # Make sure to return tensor of the same shape as the input.
+        if len_input == 3:
+            output = output.squeeze(4).squeeze(3)
+        elif len_input == 4:
+            output = output.squeeze(4)
+
+        # Ensure result are as expected.
+        np.testing.assert_allclose(output, expected, atol=1e-5)
+
+    @parameterized.expand(TEST_CASES)
+    def test_cuda_precise_backwards(self, test_case_description, sigmas, input, guide, expected):
+
+        # Params to determine the implementation to test
+        device = torch.device("cuda")
+
+        # Prepare input tensor
+        input_tensor = torch.from_numpy(np.array(input)).to(dtype=torch.double, device=device)
+        input_tensor.requires_grad = True
+        guide_tensor = torch.from_numpy(np.array(guide)).to(dtype=torch.double, device=device)
+
+        # C++ extension so far only supports 5-dim inputs.
+        len_input = len(input_tensor.shape)
+        if len_input == 3:
+            input_tensor = input_tensor.unsqueeze(3).unsqueeze(4)
+            guide_tensor = guide_tensor.unsqueeze(3).unsqueeze(4)
+        elif len_input == 4:
+            input_tensor = input_tensor.unsqueeze(4)
+            guide_tensor = guide_tensor.unsqueeze(4)
+
+        # Check gradient toward input.
+        args = (
+            input_tensor,
+            guide_tensor,
+            torch.tensor(sigmas[0]),
+            torch.tensor(sigmas[1]),
+            torch.tensor(sigmas[2]),
+            torch.tensor(sigmas[3]),
+        )
+        gradcheck(TrainableJointBilateralFilterFunction.apply, args, eps=1e-6, atol=1e-5, raise_exception=False)
+        input_tensor = input_tensor.detach()
+        input_tensor.requires_grad = False
+
+        # Check gradient toward guide.
+        guide_tensor.requires_grad = True
+        args = (
+            input_tensor,
+            guide_tensor,
+            torch.tensor(sigmas[0]),
+            torch.tensor(sigmas[1]),
+            torch.tensor(sigmas[2]),
+            torch.tensor(sigmas[3]),
+        )
+        gradcheck(TrainableJointBilateralFilterFunction.apply, args, eps=1e-6, atol=1e-5, raise_exception=False)
+        guide_tensor = guide_tensor.detach()
+        guide_tensor.guide_tensor = False
+
+        # Check gradient toward sigma_x.
+        args = (
+            input_tensor,
+            guide_tensor,
+            torch.tensor(sigmas[0], dtype=torch.double, requires_grad=True),
+            torch.tensor(sigmas[1]),
+            torch.tensor(sigmas[2]),
+            torch.tensor(sigmas[3]),
+        )
+        gradcheck(TrainableJointBilateralFilterFunction.apply, args, eps=1e-2, atol=1e-3, raise_exception=False)
+
+        # Check gradient toward sigma_y.
+        args = (
+            input_tensor,
+            guide_tensor,
+            torch.tensor(sigmas[0]),
+            torch.tensor(sigmas[1], dtype=torch.double, requires_grad=True),
+            torch.tensor(sigmas[2]),
+            torch.tensor(sigmas[3]),
+        )
+        gradcheck(TrainableJointBilateralFilterFunction.apply, args, eps=1e-2, atol=1e-3, raise_exception=False)
+
+        # Check gradient toward sigma_z.
+        args = (
+            input_tensor,
+            guide_tensor,
+            torch.tensor(sigmas[0]),
+            torch.tensor(sigmas[1]),
+            torch.tensor(sigmas[2], dtype=torch.double, requires_grad=True),
+            torch.tensor(sigmas[3]),
+        )
+        gradcheck(TrainableJointBilateralFilterFunction.apply, args, eps=1e-2, atol=1e-3, raise_exception=False)
+
+        # Check gradient toward sigma_color.
+        args = (
+            input_tensor,
+            guide_tensor,
+            torch.tensor(sigmas[0]),
+            torch.tensor(sigmas[1]),
+            torch.tensor(sigmas[2]),
+            torch.tensor(sigmas[3], dtype=torch.double, requires_grad=True),
+        )
+        gradcheck(TrainableJointBilateralFilterFunction.apply, args, eps=1e-3, atol=1e-3, raise_exception=False)
+
+
+if __name__ == "__main__":
+    unittest.main()