Abstract
TorchMorph is a PyTorch extension providing GPU-accelerated morphological and distance-transform operators across up to eight dimensions with a SciPy-compatible API.
Morphological transforms are long-standing tools for shape and mask processing, but the de facto reference implementation in the Python ecosystem, i.e. scipy.ndimage, is CPU-only, single-array, and therefore unusable inside a GPU training loop without an expensive device-to-host round trip. GPU vision libraries built on PyTorch cover a narrow subset of these operators, typically restricted to two spatial dimensions and flat structuring elements. We present TorchMorph, a lightweight PyTorch extension that closes this gap. TorchMorph exposes 22 public operators covering binary morphology, greyscale morphology, exact and approximate distance transforms, and entropy-regularised optimal transport, all implemented as fused CUDA kernels that operate directly on (B, C, Spatial...) CUDA tensors with up to eight spatial dimensions. The API deliberately mirrors scipy.ndimage argument-for-argument, including border modes, structuring-element origins and pre-allocated outputs, so that existing pipelines port with a change of import. We describe the layered architecture and the kernel designs behind each operator family. Against single-threaded CPU references, batched execution reaches up to 1.1e3 times the throughput of scipy.ndimage on greyscale morphology and up to 350x on exact Euclidean distance transforms, while the Sinkhorn solver runs up to 42x faster than POT. Binary and chamfer operators reproduce their SciPy counterparts exactly, and every float-valued operator agrees with the CPU reference to within 1.8e-6 absolute error. TorchMorph is released under the MIT licence at https://intcomp.github.io/tm.
Community
This paper presents TorchMorph, a CUDA-accelerated morphological transforms library for PyTorch. It exposes 22 operators (binary and greyscale morphology, exact/chamfer/brute-force distance transforms, entropic optimal transport) as fused CUDA kernels running directly on (B, C, Spatial...) tensors with up to 8 spatial dimensions, behind an API that mirrors scipy.ndimage argument-for-argument, so porting is a change of import. Existing GPU alternatives don't cover this: Kornia is 2-D only, cuCIM runs on CuPy and caps its exact EDT at 3-D, MONAI forwards to SciPy channel-wise.
Against single-threaded references it reaches up to 1.1×10³× scipy.ndimage throughput on greyscale morphology, 350× on exact EDT, and 42× over POT on Sinkhorn. Binary and chamfer operators match SciPy bit-exactly; float-valued operators agree to within 1.8×10⁻⁶. MIT licensed.
Get this paper in your agent:
hf papers read 2608.24738 Don't have the latest CLI?
curl -LsSf https://hf.co/cli/install.sh | bash Models citing this paper 0
No model linking this paper
Datasets citing this paper 0
No dataset linking this paper
Spaces citing this paper 0
No Space linking this paper
Collections including this paper 0
No Collection including this paper