[Feature] 3d implementation of the sobel filter

🚀 Feature

Hello, I’m using kornia to work with MRI images for the auto-differentiation feature but out of the deep-learning context. It came super usefull to not re-implement every thing.

So I found that the sobel filter had not been implemented for the 3d images case. As I need it I propose here an implementation, more for the users that like me may need it, waiting for the devs to publish one. What I have done is taking much inspiration from the diff implementation of the spatial gradient. To my understanding it should works with the auto-differentiation form backpropagation. However I am not sure that is the best and faster implementation and feel free to comment or improve.

My proposal :

def get_sobel_kernel_3d():
     """ Return a [3,1,3,3,3] sobel kernel"""
    return torch.tensor(
    [
        [[[-1,0,1],
          [-2,0,2],
          [-1,0,1]],

         [[-2,0,2],
          [-4,0,4],
          [-2,0,2]],

         [[-1,0,1],
          [-2,0,2],
          [-1,0,1]]],

        [[[-1,-2,-1],
          [0,0,0],
          [1,2,1]],

         [[-2,-4,-2],
          [0,0,0],
          [2,4,2]],

         [[-1,-2,-1],
          [0,0,0],
          [1,2,1]]],

        [[[-1,-2,-1],
          [-2,-4,-2],
          [-1,-2,-1]],

         [[0,0,0],
          [0,0,0],
          [0,0,0]],

         [[1,2,1],
          [2,4,2],
          [1,2,1]]]
    ]).unsqueeze(1)

def spacialGradient_3d(image):
    """ Implementation of a sobel 3d spatial gradient inspired by the kornia library.

    :param image: Tensor [B,1,H,W,D]
    :return: Tensor [B,3,H,W,D]

    :Example:
    H,W,D = (50,75,100)
    image = torch.zeros((H,W,D))
    mX,mY,mZ = torch.meshgrid(torch.arange(H),
                              torch.arange(W),
                              torch.arange(D))

    mask_rond = ((mX - H//2)**2 + (mY - W//2)**2).sqrt() < H//4
    mask_carre = (mX > H//4) & (mX < 3*H//4) & (mZ > D//4) & (mZ < 3*D//4)
    mask_diamand = ((mY - W//2).abs() + (mZ - D//2).abs()) < W//4
    mask = mask_rond & mask_carre & mask_diamand
    image[mask] = 1


    grad_image = spacialGradient_3d(image[None,None])
    """

    # sobel kernel is not implemented for 3D images yet in kornia
    # grad_image = SpatialGradient3d(mode='sobel')(image)
    kernel = get_sobel_kernel_3d().to(image.device).to(image.dtype)
    spatial_pad = [1,1,1,1,1,1]
    image_padded = F.pad(image,spatial_pad,'replicate').repeat(1,3,1,1,1)
    grad_image =  F.conv3d(image_padded,kernel,padding=0,groups=3,stride=1)

    return grad_image

[shijianjian]

@antonfrancois Please go ahead for the PR!

Nice repo you got! I particularly like the visualization of the gradient arrows you did. We were actually thinking of making tutorials to demonstrate the differentiability of the low-level computer vision operations. So I think it could be a great match to your examples!

\cc @edgarriba @ducha-aiki Take a look at this.

[antonfrancois]

As you can see here https://github.com/kornia/kornia/blob/43a9469b93f61aa91e0a779d3abfe2eb68a9eb22/kornia/filters/kernels.py#L308 at line 326, the filter is not implemented.

For examples you can check my repository : https://github.com/antonfrancois/Demeter_metamorphosis There are some example in jupyter notebooks, where I use kornia for basic operations. However it is a bit hidden. What kind of tutorials do you think about ?