Implementation:Kornia Kornia Distance Transform

Knowledge Sources	Kornia
Domains	Vision, Image_Processing, Distance_Fields
Last Updated	2026-02-09 15:00 GMT

Overview

Approximates the Euclidean distance transform of 2D images and 3D volumes using cascaded convolution operations.

Description

The distance_transform module in the Kornia contrib package provides both a functional interface (distance_transform) and an nn.Module class (DistanceTransform) to approximate the Euclidean distance transform. The value at each pixel/voxel in the output represents the distance to the nearest non-zero element in the input. It uses the method described in Pham et al. (2021) based on cascaded convolutions with exponential distance kernels. The implementation supports both 2D images with shape (B, C, H, W) and 3D volumes with shape (B, C, D, H, W), and processes each channel independently.

Usage

Import this module when you need to compute distance transforms on binary or thresholded images/volumes for tasks such as shape analysis, level set methods, loss functions in segmentation, or generating signed distance fields.

Code Reference

Source Location

• Repository: Kornia
• File: kornia/contrib/distance_transform.py
• Lines: 1-165

Signature

def distance_transform(
    image: torch.Tensor,
    kernel_size: int = 3,
    h: float = 0.35,
) -> torch.Tensor: ...

class DistanceTransform(nn.Module):
    def __init__(self, kernel_size: int = 3, h: float = 0.35) -> None: ...
    def forward(self, image: torch.Tensor) -> torch.Tensor: ...

Import

from kornia.contrib import distance_transform, DistanceTransform

I/O Contract

Inputs

Name	Type	Required	Description
image	torch.Tensor	Yes	Binary image or volume, shape (B, C, H, W) for 2D or (B, C, D, H, W) for 3D; must be floating point
kernel_size	int	No	Size of the convolution kernel, must be odd and >= 3 (default: 3)
h	float	No	Positive float controlling the approximation of the min function (default: 0.35)

Outputs

Name	Type	Description
distance_map	torch.Tensor	Distance transform result with same shape as input; values represent distance to nearest non-zero element

Usage Examples

import torch
from kornia.contrib import distance_transform, DistanceTransform

# 2D distance transform
image_2d = torch.zeros(1, 1, 5, 5)
image_2d[:, :, 1, 2] = 1.0  # single non-zero pixel
dt_2d = distance_transform(image_2d)
print(dt_2d.shape)  # torch.Size([1, 1, 5, 5])

# 3D distance transform on a volume
volume = torch.zeros(1, 1, 5, 5, 5)
volume[:, :, 2, 2, 2] = 1.0
dt_3d = distance_transform(volume)
print(dt_3d.shape)  # torch.Size([1, 1, 5, 5, 5])

# Using the nn.Module interface
dt_module = DistanceTransform(kernel_size=5, h=0.35)
result = dt_module(image_2d)

# Multi-channel input (channels processed independently)
multi_ch = torch.zeros(1, 3, 32, 32)
multi_ch[:, 0, 10, 10] = 1.0
multi_ch[:, 1, 20, 20] = 1.0
multi_ch[:, 2, 5, 5] = 1.0
dt_multi = DistanceTransform()(multi_ch)

Related Pages

• Principle:Kornia_Kornia_Distance_Transform_Algorithm