Track papers across labs, topics, and trends.
Search papers, labs, and topics across Lattice.
TorchMorph is a PyTorch extension that implements 22 morphological transform operators as fused CUDA kernels, enabling efficient processing of shape and mask data directly on GPU tensors. This tool addresses the limitations of existing CPU-only implementations, significantly improving throughput—up to 1,100 times faster for greyscale morphology and 350 times for exact Euclidean distance transforms compared to scipy.ndimage. The API is designed to be compatible with scipy.ndimage, allowing seamless integration into existing workflows with minimal changes.
TorchMorph achieves up to 1,100x faster morphological processing on GPUs, revolutionizing shape and mask operations in deep learning pipelines.
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.