Repository: kornia/kornia

Issue:
Add Joint Bilateral Filter
## 🚀 Feature
<!-- A clear and concise description of the feature proposal -->

Joint Bilateral Filter as described in the paper http://vis.berkeley.edu/papers/fnf/fnfSmall.pdf

Compared to Bilateral Filter, the only difference is using another image (aka guidance) to calculate the color/intensity kernel.

OpenCV: https://docs.opencv.org/4.x/da/d17/group__ximgproc__filters.html#ga80b9b58fb85dd069691b709285ab985c

Once #2242 is merged, extending the current code for Joint Bilateral Filter should be straight-forward.

## Motivation

<!-- A clear and concise description of what the problem is. Ex. I'm always frustrated when [...]. If this is related to another GitHub issue, please link here too -->

## Pitch

<!-- A clear and concise description of what you want to happen. -->

## Alternatives

<!-- A clear and concise description of any alternative solutions or features you've considered, if any. -->

## Additional context

<!-- Add any other context or screenshots about the feature request here. -->

______________________________________________________________________

#### Consider also to contribute to Kornia universe projects :)

<sub>

- [**Tutorials**](https://github.com/kornia/tutorials): our repository containing the tutorials.

</sub>



Discussion:

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
@@ -363,6 +363,7 @@ def main():
     kernel = torch.tensor([[0, 1, 0], [1, 1, 1], [0, 1, 0]])
     transforms: dict = {
         "bilateral_blur": (((11, 11), 0.1, (3, 3)), 1),
+        "joint_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),
@@ -378,7 +379,11 @@ def main():
     # ITERATE OVER THE TRANSFORMS
     for fn_name, (args, num_samples) in transforms.items():
         img_in = img5.repeat(num_samples, 1, 1, 1)
-        args_in = (img_in, *args)
+        if fn_name == "joint_bilateral_blur":
+            guide = K.geometry.resize(img2.repeat(num_samples, 1, 1, 1), img_in.shape[-2:])
+            args_in = (img_in, guide, *args)
+        else:
+            args_in = (img_in, *args)
         # import function and apply
         fn = getattr(mod, fn_name)
         out = fn(*args_in)
@@ -392,6 +397,8 @@ def main():
             out = K.enhance.normalize_min_max(out)
         if fn_name == "spatial_gradient":
             out = out.permute(2, 1, 0, 3, 4).squeeze()
+        if fn_name == "joint_bilateral_blur":
+            out = torch.cat([args_in[1], out], dim=-1)
         # save the output image
         out = torch.cat([img_in[0], *(out[i] for i in range(out.size(0)))], dim=-1)
         out_np = K.utils.tensor_to_image((out * 255.0).byte())
diff --git a/kornia/filters/__init__.py b/kornia/filters/__init__.py
--- a/kornia/filters/__init__.py
+++ b/kornia/filters/__init__.py
@@ -1,6 +1,6 @@
 from __future__ import annotations
 
-from .bilateral import BilateralBlur, bilateral_blur
+from .bilateral import BilateralBlur, JointBilateralBlur, bilateral_blur, joint_bilateral_blur
 from .blur import BoxBlur, box_blur
 from .blur_pool import (
     BlurPool2D,
@@ -99,5 +99,7 @@
     "get_gaussian_kernel2d_t",
     "get_gaussian_kernel3d_t",
     "bilateral_blur",
+    "joint_bilateral_blur",
     "BilateralBlur",
+    "JointBilateralBlur",
 ]
diff --git a/kornia/filters/bilateral.py b/kornia/filters/bilateral.py
--- a/kornia/filters/bilateral.py
+++ b/kornia/filters/bilateral.py
@@ -1,20 +1,75 @@
 from __future__ import annotations
 
 from kornia.core import Module, Tensor, pad
-from kornia.core.check import KORNIA_CHECK_IS_TENSOR, KORNIA_CHECK_SHAPE
+from kornia.core.check import KORNIA_CHECK, KORNIA_CHECK_IS_TENSOR, KORNIA_CHECK_SHAPE
 
-from .gaussian import get_gaussian_kernel2d
-from .kernels import _unpack_2d_ks
+from .kernels import _unpack_2d_ks, get_gaussian_kernel2d
 from .median import _compute_zero_padding
 
 
+def _bilateral_blur(
+    input: Tensor,
+    guidance: Tensor | None,
+    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:
+    "Single implementation for both Bilateral Filter and Joint Bilateral Filter"
+
+    KORNIA_CHECK_IS_TENSOR(input)
+    KORNIA_CHECK_SHAPE(input, ['B', 'C', 'H', 'W'])
+    if guidance is not None:
+        # NOTE: allow guidance and input having different number of channels
+        KORNIA_CHECK_IS_TENSOR(guidance)
+        KORNIA_CHECK_SHAPE(guidance, ['B', 'C', 'H', 'W'])
+        KORNIA_CHECK(
+            (guidance.shape[0] == input.shape[0]) and (guidance.shape[-2:] == input.shape[-2:]),
+            "guidance and input should have the same batch size and spatial dimensions",
+        )
+
+    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_input = pad(input, (pad_x, pad_x, pad_y, pad_y), mode=border_type)
+    unfolded_input = padded_input.unfold(2, ky, 1).unfold(3, kx, 1).flatten(-2)  # (B, C, H, W, K x K)
+
+    if guidance is None:
+        guidance = input
+        unfolded_guidance = unfolded_input
+    else:
+        padded_guidance = pad(guidance, (pad_x, pad_x, pad_y, pad_y), mode=border_type)
+        unfolded_guidance = padded_guidance.unfold(2, ky, 1).unfold(3, kx, 1).flatten(-2)  # (B, C, H, W, K x K)
+
+    diff = unfolded_guidance - guidance.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_input * kernel).sum(-1) / kernel.sum(-1)
+    return out
+
+
 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",
+    color_distance_type: str = 'l1',
 ) -> Tensor:
     r"""Blur a tensor using a Bilateral filter.
 
@@ -48,71 +103,56 @@ def bilateral_blur(
         >>> 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)
+    return _bilateral_blur(input, None, kernel_size, sigma_color, sigma_space, border_type, color_distance_type)
 
-    kernel = space_kernel * color_kernel
-    out = (unfolded * kernel).sum(-1) / kernel.sum(-1)
-    return out
 
+def joint_bilateral_blur(
+    input: Tensor,
+    guidance: 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 Joint Bilateral filter.
 
-class BilateralBlur(Module):
-    r"""Blur a tensor using a Bilateral filter.
+    .. image:: _static/img/joint_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.
+    This operator is almost identical to a Bilateral filter. The only difference
+    is that the color Gaussian kernel is computed based on another image called
+    a guidance image. See :func:`bilateral_blur()` for more information.
 
     Arguments:
+        input: the input tensor with shape :math:`(B,C,H,W)`.
+        guidance: the guidance 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 gaussian_blur2d.
+          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'``.
+          match OpenCV implementation.
 
     Returns:
-        the blurred input tensor.
-
-    Shape:
-        - Input: :math:`(B, C, H, W)`
-        - Output: :math:`(B, C, H, W)`
+        the blurred tensor with shape :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)
+        >>> guidance = torch.rand(2, 4, 5, 5)
+        >>> output = joint_bilateral_blur(input, guidance, (3, 3), 0.1, (1.5, 1.5))
         >>> output.shape
         torch.Size([2, 4, 5, 5])
     """
+    return _bilateral_blur(input, guidance, kernel_size, sigma_color, sigma_space, border_type, color_distance_type)
+
 
+# trick to make mypy not throw errors about difference in .forward() signatures of subclass and superclass
+class _BilateralBlur(Module):
     def __init__(
         self,
         kernel_size: tuple[int, int] | int,
@@ -138,7 +178,92 @@ def __repr__(self) -> str:
             f"color_distance_type={self.color_distance_type})"
         )
 
+
+class BilateralBlur(_BilateralBlur):
+    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 :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 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 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
         )
+
+
+class JointBilateralBlur(_BilateralBlur):
+    r"""Blur a tensor using a Joint Bilateral filter.
+
+    This operator is almost identical to a Bilateral filter. The only difference
+    is that the color Gaussian kernel is computed based on another image called
+    a guidance image. See :class:`BilateralBlur` for more information.
+
+    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 :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.
+
+    Returns:
+        the blurred input tensor.
+
+    Shape:
+        - Input: :math:`(B, C, H, W)`, :math:`(B, C, H, W)`
+        - Output: :math:`(B, C, H, W)`
+
+    Examples:
+        >>> input = torch.rand(2, 4, 5, 5)
+        >>> guidance = torch.rand(2, 4, 5, 5)
+        >>> blur = JointBilateralBlur((3, 3), 0.1, (1.5, 1.5))
+        >>> output = blur(input, guidance)
+        >>> output.shape
+        torch.Size([2, 4, 5, 5])
+    """
+
+    def forward(self, input: Tensor, guidance: Tensor) -> Tensor:
+        return joint_bilateral_blur(
+            input,
+            guidance,
+            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
--- a/test/filters/test_bilateral.py
+++ b/test/filters/test_bilateral.py
@@ -1,7 +1,7 @@
 import pytest
 import torch
 
-from kornia.filters import BilateralBlur, bilateral_blur
+from kornia.filters import BilateralBlur, JointBilateralBlur, bilateral_blur, joint_bilateral_blur
 from kornia.testing import BaseTester, tensor_to_gradcheck_var
 
 
@@ -148,3 +148,179 @@ def test_opencv_rgb(self, device, dtype):
 
         out = bilateral_blur(img, kernel_size, sigma_color, sigma_distance)
         self.assert_close(out, expected, rtol=1e-2, atol=1e-2)
+
+
+class TestJointBilateralBlur(BaseTester):
+    @pytest.mark.parametrize("input_depth", [1, 3])
+    @pytest.mark.parametrize("guidance_depth", [1, 3])
+    def test_smoke(self, input_depth, guidance_depth, device, dtype):
+        b, h, w = 2, 8, 15
+        kernel_size = 5
+        sigma_color = 0.1
+        sigma_space = (2, 2)
+        inp = torch.rand(b, input_depth, h, w, device=device, dtype=dtype)
+        guide = torch.rand(b, guidance_depth, h, w, device=device, dtype=dtype)
+
+        out = joint_bilateral_blur(inp, guide, kernel_size, sigma_color, sigma_space)
+        assert isinstance(out, torch.Tensor)
+        assert out.shape == (b, input_depth, h, w)
+
+    def test_same_input(self, device, dtype):
+        shape = (2, 3, 8, 15)
+        kernel_size = 5
+        sigma_color = 0.1
+        sigma_space = (2, 2)
+        inp = torch.rand(shape, device=device, dtype=dtype)
+
+        out1 = joint_bilateral_blur(inp, inp, kernel_size, sigma_color, sigma_space)
+        out2 = bilateral_blur(inp, kernel_size, sigma_color, sigma_space)
+        self.assert_close(out1, out2)
+
+    @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)
+        guide = torch.zeros(shape, device=device, dtype=dtype)
+        actual = joint_bilateral_blur(inp, guide, kernel_size, 0.1, (1, 1))
+        assert actual.shape == shape
+
+    def test_exception(self):
+        inp = torch.rand(1, 1, 5, 5)
+        guide = torch.rand(1, 1, 5, 5)
+
+        with pytest.raises(Exception) as errinfo:
+            joint_bilateral_blur(inp, guide, 3, 1, 1)
+        assert 'Not a Tensor type. Go' in str(errinfo)
+
+        with pytest.raises(Exception) as errinfo:
+            joint_bilateral_blur(inp, torch.randn(1, 1, 2, 4), 3, 1, (1, 1))
+        assert 'guidance and input should have the same' in str(errinfo)
+
+        with pytest.raises(Exception) as errinfo:
+            joint_bilateral_blur(inp, torch.randn(2, 1, 5, 5), 3, 1, (1, 1))
+        assert 'guidance and input should have the same' in str(errinfo)
+
+        with pytest.raises(ValueError) as errinfo:
+            joint_bilateral_blur(inp, guide, 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)
+        guide = torch.rand(3, 5, 5, device=device, dtype=dtype).expand(batch_size, -1, -1, -1)
+
+        actual = joint_bilateral_blur(inp, guide, 3, 1, (1, 1))
+        assert actual.is_contiguous()
+
+    def test_gradcheck(self, device):
+        img = torch.rand(1, 2, 5, 4, device=device)
+        guide = torch.rand(1, 2, 5, 4, device=device)
+        img = tensor_to_gradcheck_var(img)  # to var
+        guide = tensor_to_gradcheck_var(guide)
+        self.gradcheck(joint_bilateral_blur, (img, guide, 3, 1, (1, 1)))
+
+    def test_module(self, device, dtype):
+        shape = (2, 3, 11, 7)
+        kernel_size = 5
+        sigma_color = 0.1
+        sigma_space = (2, 2)
+        img = torch.rand(shape, device=device, dtype=dtype)
+        guide = torch.rand(shape, device=device, dtype=dtype)
+        params = (kernel_size, sigma_color, sigma_space)
+
+        op = joint_bilateral_blur
+        op_module = JointBilateralBlur(*params)
+        self.assert_close(op_module(img, guide), op(img, guide, *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.rand(2, 3, 8, 8, device=device, dtype=dtype)
+        guide = torch.rand(2, 3, 8, 8, device=device, dtype=dtype)
+        op = JointBilateralBlur(kernel_size, 1, (1, 1), color_distance_type=color_distance_type)
+        op_optimized = torch_optimizer(op)
+
+        self.assert_close(op(inpt, guide), op_optimized(inpt, guide))
+
+        sigma_color = torch.rand(inpt.shape[0], device=device, dtype=dtype)
+        sigma_space = torch.rand(inpt.shape[0], 2, device=device, dtype=dtype)
+        op = JointBilateralBlur(kernel_size, sigma_color, sigma_space, color_distance_type=color_distance_type)
+        op_optimized = torch_optimizer(op)
+
+        self.assert_close(op(inpt, guide), op_optimized(inpt, guide))
+
+    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
+
+        guide = [[161, 87, 93, 6], [91, 182, 97, 154], [70, 123, 109, 70], [119, 28, 60, 109]]
+        guide = torch.tensor(guide, 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.ximgproc.jointBilateralFilter(img[0, 0].numpy(), guide[0, 0].numpy(), 5, 0.1, 0.5)
+        expected = [
+            [0.38221005, 0.5027215, 0.49131155, 0.8937083],
+            [0.3976327, 0.55548316, 0.69680846, 0.65291953],
+            [0.44903287, 0.65470666, 0.7295845, 0.5840189],
+            [0.5867507, 0.3472942, 0.66494286, 0.81431836],
+        ]
+        expected = torch.tensor(expected, device=device, dtype=dtype).view(1, 1, 4, 4)
+
+        out = joint_bilateral_blur(img, guide, 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
+
+        guide = [
+            [[136, 196, 198, 21], [149, 185, 196, 141], [115, 110, 87, 155], [126, 82, 109, 207]],
+            [[188, 42, 48, 0], [73, 200, 58, 173], [53, 140, 130, 34], [129, 3, 41, 73]],
+            [[85, 36, 51, 0], [33, 82, 37, 87], [35, 70, 57, 36], [50, 10, 30, 43]],
+        ]
+        guide = torch.tensor(guide, 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.ximgproc.jointBilateralFilter(
+        #     img[0].permute(1, 2, 0).numpy(),
+        #     guide[0].permute(1, 2, 0).numpy(),
+        #     5, 0.1, 0.5
+        # )
+        expected = [
+            [
+                [0.6671455, 0.74172455, 0.7328562, 1.0],
+                [0.66403687, 0.81948805, 0.8357967, 0.8431371],
+                [0.7673652, 0.836736, 0.8925936, 0.7494889],
+                [0.8120746, 0.49431974, 0.8778994, 0.9764218],
+            ],
+            [
+                [0.23984201, 0.40574068, 0.35012922, 0.88235295],
+                [0.2555589, 0.43905944, 0.65692204, 0.58039117],
+                [0.30529743, 0.5791127, 0.68724936, 0.47124338],
+                [0.48745763, 0.23940866, 0.6072882, 0.82737976],
+            ],
+            [
+                [0.2868149, 0.43398118, 0.39570162, 0.6862745],
+                [0.30228183, 0.45868847, 0.6153678, 0.50980365],
+                [0.3252297, 0.5474381, 0.6367775, 0.47098273],
+                [0.51799667, 0.32952034, 0.53696644, 0.6078228],
+            ],
+        ]
+        expected = torch.tensor(expected, device=device, dtype=dtype).view(1, 3, 4, 4)
+
+        out = joint_bilateral_blur(img, guide, kernel_size, sigma_color, sigma_distance)
+        self.assert_close(out, expected, rtol=1e-2, atol=1e-2)