Repository: kornia/kornia

Issue:
[Feat] Bilateral filter
## 🚀 Feature

Would be great to have a differentiable version of Bilater Filtering `kornia.filter.bilateral_blur`

We can follow opencv signature:
https://docs.opencv.org/3.4/d4/d86/group__imgproc__filter.html#ga9d7064d478c95d60003cf839430737ed

Resource: http://jamesgregson.ca/bilateral-filtering-in-python.html



Discussion:
Such a feature would be great. Especially if it also works in 3D and is implemented in CUDA (or rather running efficiently on the GPU). See e.g. for a starting point in CUDA:
https://github.com/opencv/opencv_contrib/blob/master/modules/cudaimgproc/src/cuda/bilateral_filter.cu
http://developer.download.nvidia.com/compute/DevZone/C/html_x64/Image_Processing.html#bilateralFilter
@tvercaut not sure we are going to support custom cuda kernels for now. I would go first for a naive solution using plane pytorch
OK, I played a bit with it. For a simple 449 x 359 image and a spatial kernel diameter of 9, I ended up with these timings (the pytorch based version is buggy but I would think the timing is roughly OK):
```
OpenCV CPU
100 loops, best of 3: 8.71 ms per loop

OpenCV UMat (OpenCL)
1000 loops, best of 3: 289 µs per loop

"Naive" numpy (Gregson based)
1 loop, best of 3: 793 ms per loop

"Naive" pytorch (Gregson based) - debugging required
10 loops, best of 3: 79.9 ms per loop
```
The CPU-GPU IO is not counted.
@tvercaut great ! consider submitting a PR. We can assist you to complete
Sorry, I don't have the bandwidth to do any real software development these days. If helpful, my quick and dirty notebook is here:
https://colab.research.google.com/drive/1hUsuHsCD1_QOhx7yrTfSsWsylmaKRu8b?usp=sharing
@tvercaut this is already an amazing contribution :)) thanks a lot 
Hello! I would like to work in the Bilateral filter functionality.
Hi @AnimeshMaheshwari22 

> Hello! I would like to work in the Bilateral filter functionality.

perfect!!!! Let's see how much work requires this Feature and then we can determine how to proceed. :)

Luis




Agree, let’s continue this discussion in slack 
Any updates on this?
@AnimeshMaheshwari22 is going to give a try on that, @AhmadM-DL feel free to join our slack in case you want to be involved too in that feature.
Don't seem to find any `kornia.filter.bilateral_blur` variable. Did its name change in new version?

@avani17101 bilateral file was not finished since we couldn't fins a proper API vs "fast" implementation. This particular algorithm needs to almost iterate through the different pixels to compute weighting in a local neighborhood.  This is extremely slow with python -- we should find another good vectored way to implement decently efficient. /cc @ducha-aiki 
Hello. This is my apprehension as well. The functionality I found close to this was uniform spatial bilateral filtering that is relatively faster, but that is not the exact requirement.
@edgarriba how about this implementation? https://github.com/sunny2109/bilateral_filter_Pytorch/blob/main/bilateral_filter.py
It seems one can just use unfold2d, no?
might fit !! @lferraz suggested something similar with unfolding 
Since I had submitted a PR, I can continue that and give it a try!

> Since I had submitted a PR, I can continue that and give it a try!

would be great to follow this new implementation
I was looking for bilateral filter and saw this issue and the unfinished PR. I took a try, following ideas from sunny2109 implementations.

```python
import torch
import torch.nn.functional as F
from kornia.filters.gaussian import get_gaussian_kernel2d
import numpy as np
import cv2
import requests
import io
from PIL import Image
import matplotlib.pyplot as plt
from skimage.restoration import denoise_bilateral


def bilateral_filter(imgs, kernel_size, sigma_intensity, sigma_space):
    pad = (kernel_size - 1) // 2
    padded = F.pad(imgs, (pad, pad, pad, pad), mode="reflect")
    unfolded = padded.unfold(2, kernel_size, 1).unfold(3, kernel_size, 1).flatten(-2)  # (B, C, H, W, K x K)
    intensity_kernel = torch.exp(-0.5 * (unfolded - imgs[..., None]).square() / sigma_intensity ** 2)
    space_kernel = get_gaussian_kernel2d(
        kernel_size,
        (sigma_space, sigma_space),
        device=imgs.device,
        dtype=imgs.dtype
    ).view(1, 1, 1, 1, -1)
    kernel = space_kernel * intensity_kernel
    out = (unfolded * kernel).sum(-1) / kernel.sum(-1)
    return out


img = io.BytesIO(requests.get("https://upload.wikimedia.org/wikipedia/en/7/7d/Lenna_%28test_image%29.png").content)
x_np = np.array(Image.open(img))
x_np = x_np.astype(np.float32) / 255
x = torch.from_numpy(x_np).permute(2, 0, 1).unsqueeze(0)

out1 = bilateral_filter(x, 11, 0.5, 2)
out2 = cv2.bilateralFilter(x_np, 11, 0.5, 2)
out3 = denoise_bilateral(x_np, 11, 0.5, 2, bins=100_000, mode="reflect", channel_axis=-1)

diff = np.abs(out1.squeeze(0).permute(1, 2, 0).numpy() - out2) / out2
print(diff.mean(), diff.max(), diff.min())  # 0.015516709 2.511976 0.0

diff = np.abs(out1.squeeze(0).permute(1, 2, 0).numpy() - out3) / out3
print(diff.mean(), diff.max(), diff.min())  # 0.047820166 2.9121413 0.0

plt.subplot(1, 3, 1)
plt.imshow(out1.squeeze(0).permute(1, 2, 0).mul(255).numpy().astype(np.uint8))
plt.title("pytorch")

plt.subplot(1, 3, 2)
plt.imshow((out2 * 255).astype(np.uint8))
plt.title("opencv")

plt.subplot(1, 3, 3)
plt.imshow((out3 * 255).astype(np.uint8))
plt.title("skimage")

plt.show()
```

However, comparing with OpenCV and skimage implementations, the results are significantly different. Even OpenCV and skimage are not the same. OpenCV is the sharpest, followed by my PyTorch implementation, and then skimage. skimage has very strong ghosting artifacts.

![Figure_1](https://user-images.githubusercontent.com/26946864/221741288-ff3a3fcb-fe89-43cb-9829-3acc5cc5eda2.png)

@ducha-aiki @edgarriba If you can find out any problems with my implementation, I'm happy to try and fix.

OpenCV code: https://github.com/opencv/opencv/blob/4.x/modules/imgproc/src/bilateral_filter.dispatch.cpp and https://github.com/opencv/opencv/blob/4.x/modules/imgproc/src/bilateral_filter.dispatch.cpp. It's not very readable so I can't really understand what would lead to different results. I noticed that their spatial Gaussian kernel is limited by the radius, not by a square window. (`if (sqrt(i*i + j*j) > radius) continue`). This would make the values in the 4 corners of the filter a bit different.

https://github.com/opencv/opencv/blob/b3c9842a98a263cef9daaded8dfa20a82f30e06b/modules/imgproc/src/bilateral_filter.dispatch.cpp#L206-L207

Quality-wise, reducing intensity sigma (preserve more edges) in my PyTorch implementation can make the result closer to OpenCV's.

I haven't looked in skimage code in-depth yet.
> It's not very readable so I can't really understand what would lead to different results. I noticed that their spatial Gaussian kernel is limited by the radius, not by a square window. (`if (sqrt(i*i + j*j) > radius) continue`).

wondering whether we can solve that by clipping or similar. Have you benchmarked times ? It's fine to included this implementation, but might be worth to explore prototyping a triton example too. See: https://github.com/openai/triton/blob/main/python/tutorials/01-vector-add.py

not sure though how we solve then the auto differentiation; or we just have different modes. /cc @ducha-aiki @tvercaut @t-vi 
@gau-nernst that is very clever implementation, I love it!
Input is 512x512x3. CPU is Intel W-2295, GPU is RTX3090. Time measured to process 1 image. For batch 32, time is divided by 32.

Implementation | kernel_size=3 | kernel_size=5 | kernel_size=7
--|--|--|--
This PyTorch (CPU - batch 1) | 11.0294ms | 53.6026ms | 99.6424ms
This PyTorch (CPU - batch 1 - compile) | 2.8805ms | 14.0371ms | 24.8569ms
This PyTorch (CPU - batch 32) | 19.2862ms | 46.0436ms | 83.7909ms
This PyTorch (CPU - batch 32 - compile) | 4.5717ms | 13.0927ms | 24.5206ms
This PyTorch (GPU - batch 1) | 0.9488ms | 1.9329ms | 3.4727ms
This PyTorch (GPU - batch 1 - compile) | 0.1224ms | 0.1212ms | 0.1226ms
This PyTorch (GPU - batch 32) | 0.6444ms | 1.6317ms | OOM
This PyTorch (GPU - batch 32 - compile) | 0.0038ms | 0.0038ms | 0.0039ms
OpenCV | 0.2250ms | 0.4486ms | 1.4840ms
Scikit-image | 42.2278ms | 112.9555ms | 214.5870ms

It's not surprising that this is much slower than OpenCV. Doing calculations on unfolded tensors is expensive. In terms of speed, I can try optimize some operations e.g. use `torch.bmm()` at some places. Scikit-image being slower than PyTorch is quite surprising, considering that they use custom Cython code. Or perhaps I'm doing something wrong with scikit-image.

@edgarriba regarding the `if (sqrt(i*i + j*j) > radius) continue`, I think it's possible to create a boolean mask to mask out the Gaussian kernel before unfold it. I will try and let you know.

UPDATE: add `torch.compile()` results with pytorch==2.0.0.dev20230228. The results are insane!
I would benchmark also against torchdynamo, we are seeing some speed improvements compared to opencv. /cc @johnnv1 https://github.com/kornia/kornia-benchmark/tree/master/dynamo
@edgarriba @gau-nernst I believe that better strategy would be to merge correct, but slower implementation first and do the speed-ups in the consequent PRs
I just added speed test with `torch.compile()`. It's unbelievable!

@ducha-aiki Yes I agree, correctness first before speed optimization.
> I just added speed test with `torch.compile()`. It's unbelievable!

tell me more :)
I think I have figured out the difference. For RGB images, OpenCV uses sum of absolute differences across three channels to calculate the intensity difference $|r_1-r_2| + |g_1-g_2| + |b_1-b_2|$

https://github.com/opencv/opencv/blob/4.x/modules/imgproc/src/bilateral_filter.simd.hpp#L880

The code above should be modified like this (take absolute value of the difference, then sum over channel dimension)

```python
    intensity_diff = (unfolded - imgs.unsqueeze(-1)).abs().sum(1, keepdim=True)  # (B, 1, H, W, K x K)
    intensity_kernel = torch.exp(-0.5 / sigma_intensity ** 2 * intensity_diff.square())
```

Numerically, the absolute difference is 0.12% on average, and 12.8% maximum. Visually, I think they look identical. Using the same settings as above (kernel size = 11, sigma intensity / color = 0.5, sigma space = 2)

![Figure_1](https://user-images.githubusercontent.com/26946864/221893553-5b6bdbb9-3c44-41e4-af99-592e6131751e.png)

The radial mask `if (sqrt(i*i + j*j) > radius) continue` can be implemented like the following

```python
    mask = torch.arange(-pad, pad+1, device=imgs.device).square()
    mask = (mask.view(-1, 1) + mask.view(1, -1)) <= kernel_size * kernel_size
    space_kernel = space_kernel * mask
```

I don't observe any difference when including the radial mask.

Also, regarding bilateral filter for RGB images, this paper (https://ieeexplore.ieee.org/document/710815) suggests converting to Lab color space and applying Euclidean distance (square root of sum of squared differences) $\sqrt{(L_1-L_2)^2 + (a_1-a_2)^2 + (b_1-b_2)^2}$. Matlab also suggests doing the same (https://www.mathworks.com/help/images/ref/imbilatfilt.html).
These may also be of interest:
https://github.com/faebstn96/trainable-bilateral-filter-source
https://github.com/faebstn96/trainable-joint-bilateral-filter-source
https://docs.monai.io/en/stable/networks.html#bilateralfilter
> These may also be of interest: https://github.com/faebstn96/trainable-bilateral-filter-source https://github.com/faebstn96/trainable-joint-bilateral-filter-source https://docs.monai.io/en/stable/networks.html#bilateralfilter

I saw those custom C++/CUDA kernel code. I'm not familiar with writing such code so I probably won't be able to write it (unless we just directly copy it?) Compiling and packaging will be more difficult, and harder to maintain. For example, I tried compiling Monai (they have no pre-compiled kernels on PyPI yet) for my macbook. There were no errors in the compilation/installation phase but I couldn't call the bilateral filter function (I got some missing symbol errors).

Nevertheless, looking at these implementations, it seems like they use L2 distance for color difference (like Matlab). 

https://github.com/faebstn96/trainable-bilateral-filter-source/blob/main/csrc/bf_layer_gpu_forward.cu#L114

https://github.com/Project-MONAI/MONAI/blob/dev/monai/csrc/filtering/bilateral/bilateralfilter_cuda.cu#L111

Resolving patch:
diff --git a/docs/generate_example_images.py b/docs/generate_example_images.py
--- a/docs/generate_example_images.py
+++ b/docs/generate_example_images.py
@@ -361,6 +361,7 @@ def main():
     mod = importlib.import_module("kornia.filters")
     kernel = torch.tensor([[0, 1, 0], [1, 1, 1], [0, 1, 0]])
     transforms: dict = {
+        "bilateral_blur": (((11, 11), 0.1, (3, 3)), 1),
         "box_blur": (((5, 5),), 1),
         "median_blur": (((5, 5),), 1),
         "gaussian_blur2d": (((5, 5), (1.5, 1.5)), 1),
diff --git a/kornia/filters/__init__.py b/kornia/filters/__init__.py
--- a/kornia/filters/__init__.py
+++ b/kornia/filters/__init__.py
@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+from .bilateral import BilateralBlur, bilateral_blur
 from .blur import BoxBlur, box_blur
 from .blur_pool import (
     BlurPool2D,
@@ -97,4 +98,6 @@
     "get_gaussian_kernel1d_t",
     "get_gaussian_kernel2d_t",
     "get_gaussian_kernel3d_t",
+    "bilateral_blur",
+    "BilateralBlur",
 ]
diff --git a/kornia/filters/bilateral.py b/kornia/filters/bilateral.py
new file mode 100644
--- /dev/null
+++ b/kornia/filters/bilateral.py
@@ -0,0 +1,144 @@
+from __future__ import annotations
+
+from kornia.core import Module, Tensor, pad
+from kornia.core.check import KORNIA_CHECK_IS_TENSOR, KORNIA_CHECK_SHAPE
+
+from .gaussian import get_gaussian_kernel2d
+from .kernels import _unpack_2d_ks
+from .median import _compute_zero_padding
+
+
+def bilateral_blur(
+    input: Tensor,
+    kernel_size: tuple[int, int] | int,
+    sigma_color: float | Tensor,
+    sigma_space: tuple[float, float] | Tensor,
+    border_type: str = 'reflect',
+    color_distance_type: str = "l1",
+) -> Tensor:
+    r"""Blur a tensor using a Bilateral filter.
+
+    .. image:: _static/img/bilateral_blur.png
+
+    The operator is an edge-preserving image smoothing filter. The weight
+    for each pixel in a neighborhood is determined not only by its distance
+    to the center pixel, but also the difference in intensity or color.
+
+    Arguments:
+        input: the input tensor with shape :math:`(B,C,H,W)`.
+        kernel_size: the size of the kernel.
+        sigma_color: the standard deviation for intensity/color Gaussian kernel.
+          Smaller values preserve more edges.
+        sigma_space: the standard deviation for spatial Gaussian kernel.
+          This is similar to ``sigma`` in :func:`gaussian_blur2d()`.
+        border_type: the padding mode to be applied before convolving.
+          The expected modes are: ``'constant'``, ``'reflect'``,
+          ``'replicate'`` or ``'circular'``. Default: ``'reflect'``.
+        color_distance_type: the type of distance to calculate intensity/color
+          difference. Only ``'l1'`` or ``'l2'`` is allowed. Use ``'l1'`` to
+          match OpenCV implementation. Use ``'l2'`` to match Matlab implementation.
+          Default: ``'l1'``.
+
+    Returns:
+        the blurred tensor with shape :math:`(B, C, H, W)`.
+
+    Examples:
+        >>> input = torch.rand(2, 4, 5, 5)
+        >>> output = bilateral_blur(input, (3, 3), 0.1, (1.5, 1.5))
+        >>> output.shape
+        torch.Size([2, 4, 5, 5])
+    """
+    KORNIA_CHECK_IS_TENSOR(input)
+    KORNIA_CHECK_SHAPE(input, ['B', 'C', 'H', 'W'])
+    if isinstance(sigma_color, Tensor):
+        KORNIA_CHECK_SHAPE(sigma_color, ['B'])
+        sigma_color = sigma_color.to(device=input.device, dtype=input.dtype).view(-1, 1, 1, 1, 1)
+
+    kx, ky = _unpack_2d_ks(kernel_size)
+    pad_x, pad_y = _compute_zero_padding(kernel_size)
+
+    padded = pad(input, (pad_x, pad_x, pad_y, pad_y), mode=border_type)
+    unfolded = padded.unfold(2, ky, 1).unfold(3, kx, 1).flatten(-2)  # (B, C, H, W, K x K)
+
+    diff = unfolded - input.unsqueeze(-1)
+    if color_distance_type == "l1":
+        color_distance_sq = diff.abs().sum(1, keepdim=True).square()
+    elif color_distance_type == "l2":
+        color_distance_sq = diff.square().sum(1, keepdim=True)
+    else:
+        raise ValueError("color_distance_type only acceps l1 or l2")
+    color_kernel = (-0.5 / sigma_color**2 * color_distance_sq).exp()  # (B, 1, H, W, K x K)
+
+    space_kernel = get_gaussian_kernel2d(kernel_size, sigma_space, device=input.device, dtype=input.dtype)
+    space_kernel = space_kernel.view(-1, 1, 1, 1, kx * ky)
+
+    kernel = space_kernel * color_kernel
+    out = (unfolded * kernel).sum(-1) / kernel.sum(-1)
+    return out
+
+
+class BilateralBlur(Module):
+    r"""Blur a tensor using a Bilateral filter.
+
+    The operator is an edge-preserving image smoothing filter. The weight
+    for each pixel in a neighborhood is determined not only by its distance
+    to the center pixel, but also the difference in intensity or color.
+
+    Arguments:
+        kernel_size: the size of the kernel.
+        sigma_color: the standard deviation for intensity/color Gaussian kernel.
+          Smaller values preserve more edges.
+        sigma_space: the standard deviation for spatial Gaussian kernel.
+          This is similar to ``sigma`` in gaussian_blur2d.
+        border_type: the padding mode to be applied before convolving.
+          The expected modes are: ``'constant'``, ``'reflect'``,
+          ``'replicate'`` or ``'circular'``. Default: ``'reflect'``.
+        color_distance_type: the type of distance to calculate intensity/color
+          difference. Only ``'l1'`` or ``'l2'`` is allowed. Use ``'l1'`` to
+          match OpenCV implementation. Use ``'l2'`` to match Matlab implementation.
+          Default: ``'l1'``.
+
+    Returns:
+        the blurred input tensor.
+
+    Shape:
+        - Input: :math:`(B, C, H, W)`
+        - Output: :math:`(B, C, H, W)`
+
+    Examples:
+        >>> input = torch.rand(2, 4, 5, 5)
+        >>> blur = BilateralBlur((3, 3), 0.1, (1.5, 1.5))
+        >>> output = blur(input)
+        >>> output.shape
+        torch.Size([2, 4, 5, 5])
+    """
+
+    def __init__(
+        self,
+        kernel_size: tuple[int, int] | int,
+        sigma_color: float,
+        sigma_space: tuple[float, float],
+        border_type: str = 'reflect',
+        color_distance_type: str = "l1",
+    ) -> None:
+        super().__init__()
+        self.kernel_size = kernel_size
+        self.sigma_color = sigma_color
+        self.sigma_space = sigma_space
+        self.border_type = border_type
+        self.color_distance_type = color_distance_type
+
+    def __repr__(self) -> str:
+        return (
+            f"{self.__class__.__name__}"
+            f"(kernel_size={self.kernel_size}, "
+            f"sigma_color={self.sigma_color}, "
+            f"sigma_space={self.sigma_space}, "
+            f"border_type={self.border_type}, "
+            f"color_distance_type={self.color_distance_type})"
+        )
+
+    def forward(self, input: Tensor) -> Tensor:
+        return bilateral_blur(
+            input, self.kernel_size, self.sigma_color, self.sigma_space, self.border_type, self.color_distance_type
+        )


Test patch:
diff --git a/test/filters/test_bilateral.py b/test/filters/test_bilateral.py
new file mode 100644
--- /dev/null
+++ b/test/filters/test_bilateral.py
@@ -0,0 +1,150 @@
+import pytest
+import torch
+
+from kornia.filters import BilateralBlur, bilateral_blur
+from kornia.testing import BaseTester, tensor_to_gradcheck_var
+
+
+class TestBilateralBlur(BaseTester):
+    @pytest.mark.parametrize("shape", [(1, 1, 8, 15), (2, 3, 11, 7)])
+    @pytest.mark.parametrize("kernel_size", [5, (3, 5)])
+    @pytest.mark.parametrize("sigma_color", [1, 0.1])
+    @pytest.mark.parametrize("sigma_space", [(1, 1), (1.5, 1)])
+    @pytest.mark.parametrize("color_distance_type", ["l1", "l2"])
+    def test_smoke(self, shape, kernel_size, sigma_color, sigma_space, color_distance_type, device, dtype):
+        inp = torch.zeros(shape, device=device, dtype=dtype)
+
+        # tensor sigmas -> with batch dim
+        sigma_color = torch.rand(shape[0], device=device, dtype=dtype)
+        sigma_space = torch.rand(shape[0], 2, device=device, dtype=dtype)
+        actual_A = bilateral_blur(inp, kernel_size, sigma_color, sigma_space, "reflect", color_distance_type)
+        assert isinstance(actual_A, torch.Tensor)
+        assert actual_A.shape == shape
+
+        # float and tuple sigmas -> same sigmas across batch
+        sigma_color_ = sigma_color[0].item()
+        sigma_space_ = tuple(sigma_space[0].cpu().numpy())
+        actual_B = bilateral_blur(inp, kernel_size, sigma_color_, sigma_space_, "reflect", color_distance_type)
+        assert isinstance(actual_B, torch.Tensor)
+        assert actual_B.shape == shape
+
+        self.assert_close(actual_A[0], actual_B[0])
+
+    @pytest.mark.parametrize("shape", [(1, 1, 8, 15), (2, 3, 11, 7)])
+    @pytest.mark.parametrize("kernel_size", [5, (3, 5)])
+    def test_cardinality(self, shape, kernel_size, device, dtype):
+        inp = torch.zeros(shape, device=device, dtype=dtype)
+        actual = bilateral_blur(inp, kernel_size, 0.1, (1, 1))
+        assert actual.shape == shape
+
+    def test_exception(self):
+        with pytest.raises(Exception) as errinfo:
+            bilateral_blur(torch.rand(1, 1, 5, 5), 3, 1, 1)
+        assert 'Not a Tensor type. Go' in str(errinfo)
+
+        with pytest.raises(ValueError) as errinfo:
+            bilateral_blur(torch.rand(1, 1, 5, 5), 3, 0.1, (1, 1), color_distance_type="l3")
+        assert 'color_distance_type only acceps l1 or l2' in str(errinfo)
+
+    def test_noncontiguous(self, device, dtype):
+        batch_size = 3
+        inp = torch.rand(3, 5, 5, device=device, dtype=dtype).expand(batch_size, -1, -1, -1)
+
+        actual = bilateral_blur(inp, 3, 1, (1, 1))
+        assert actual.is_contiguous()
+
+    def test_gradcheck(self, device):
+        img = torch.rand(1, 2, 5, 4, device=device)
+        img = tensor_to_gradcheck_var(img)  # to var
+        self.gradcheck(bilateral_blur, (img, 3, 1, (1, 1)))
+
+    @pytest.mark.parametrize("shape", [(1, 1, 8, 15), (2, 3, 11, 7)])
+    @pytest.mark.parametrize("kernel_size", [5, (3, 5)])
+    @pytest.mark.parametrize("sigma_color", [1, 0.1])
+    @pytest.mark.parametrize("sigma_space", [(1, 1), (1.5, 1)])
+    @pytest.mark.parametrize("color_distance_type", ["l1", "l2"])
+    def test_module(self, shape, kernel_size, sigma_color, sigma_space, color_distance_type, device, dtype):
+        img = torch.rand(shape, device=device, dtype=dtype)
+        params = (kernel_size, sigma_color, sigma_space, "reflect", color_distance_type)
+
+        op = bilateral_blur
+        op_module = BilateralBlur(*params)
+        self.assert_close(op_module(img), op(img, *params))
+
+    @pytest.mark.parametrize('kernel_size', [5, (5, 7)])
+    @pytest.mark.parametrize('color_distance_type', ["l1", "l2"])
+    def test_dynamo(self, kernel_size, color_distance_type, device, dtype, torch_optimizer):
+        inpt = torch.ones(2, 3, 8, 8, device=device, dtype=dtype)
+        op = BilateralBlur(kernel_size, 1, (1, 1), color_distance_type=color_distance_type)
+        op_optimized = torch_optimizer(op)
+
+        self.assert_close(op(inpt), op_optimized(inpt))
+
+        sigma_color = torch.rand(inpt.shape[0], device=device, dtype=dtype)
+        sigma_space = torch.rand(inpt.shape[0], 2, device=device, dtype=dtype)
+        op = BilateralBlur(kernel_size, sigma_color, sigma_space, color_distance_type=color_distance_type)
+        op_optimized = torch_optimizer(op)
+
+        self.assert_close(op(inpt), op_optimized(inpt))
+
+    def test_opencv_grayscale(self, device, dtype):
+        img = [[95, 130, 108, 228], [98, 142, 187, 166], [114, 166, 190, 141], [150, 83, 174, 216]]
+        img = torch.tensor(img, device=device, dtype=dtype).view(1, 1, 4, 4) / 255
+
+        kernel_size = 5
+        sigma_color = 0.1
+        sigma_distance = (0.5, 0.5)
+
+        # Expected output generated with OpenCV:
+        # import cv2
+        # expected = cv2.bilateralFilter(img[0, 0].numpy(), 5, 0.1, 0.5)
+        expected = [
+            [0.38708255, 0.5060622, 0.43372786, 0.8876763],
+            [0.39813757, 0.55695623, 0.72320986, 0.6593296],
+            [0.4527661, 0.6484203, 0.7295754, 0.5705908],
+            [0.5774919, 0.32919288, 0.6949335, 0.83184093],
+        ]
+        expected = torch.tensor(expected, device=device, dtype=dtype).view(1, 1, 4, 4)
+
+        out = bilateral_blur(img, kernel_size, sigma_color, sigma_distance)
+        self.assert_close(out, expected, rtol=1e-2, atol=1e-2)
+
+    def test_opencv_rgb(self, device, dtype):
+        img = [
+            [[170, 189, 182, 255], [169, 209, 216, 215], [196, 213, 228, 191], [207, 126, 224, 249]],
+            [[61, 104, 74, 225], [65, 112, 176, 148], [78, 147, 176, 120], [124, 61, 155, 211]],
+            [[73, 111, 90, 175], [77, 117, 163, 130], [83, 139, 163, 120], [132, 84, 137, 155]],
+        ]
+        img = torch.tensor(img, device=device, dtype=dtype).view(1, 3, 4, 4) / 255
+
+        kernel_size = 5
+        sigma_color = 0.1
+        sigma_distance = (0.5, 0.5)
+
+        # Expected output generated with OpenCV:
+        # import cv2
+        # expected = cv2.bilateralFilter(img[0].permute(1, 2, 0).numpy(), 5, 0.1, 0.5)
+        expected = [
+            [
+                [0.6658919, 0.7486991, 0.7140039, 0.9999949],
+                [0.6656203, 0.815614, 0.852062, 0.84256846],
+                [0.7658699, 0.83580506, 0.88873357, 0.7496973],
+                [0.8123873, 0.49414372, 0.87789816, 0.97619873],
+            ],
+            [
+                [0.24242306, 0.40987095, 0.29138556, 0.8823465],
+                [0.2543548, 0.43856043, 0.68934506, 0.58119816],
+                [0.3045888, 0.5758538, 0.6885629, 0.47137713],
+                [0.48865014, 0.23922202, 0.6099074, 0.82698977],
+            ],
+            [
+                [0.28948042, 0.43686634, 0.35377124, 0.686273],
+                [0.30078027, 0.4582056, 0.63826954, 0.5113827],
+                [0.32491508, 0.5446742, 0.63721484, 0.47087318],
+                [0.51836246, 0.32941142, 0.5409612, 0.60793126],
+            ],
+        ]
+        expected = torch.tensor(expected, device=device, dtype=dtype).view(1, 3, 4, 4)
+
+        out = bilateral_blur(img, kernel_size, sigma_color, sigma_distance)
+        self.assert_close(out, expected, rtol=1e-2, atol=1e-2)