diff --git a/.gitignore b/.gitignore index bb5af7e..98add0f 100644 --- a/.gitignore +++ b/.gitignore @@ -55,6 +55,6 @@ Thumbs.db .coverage htmlcov/ -# Local development documentation +# Local documentation artifacts /CONTEXT.md -/docs/ +/site/ diff --git a/LICENSE b/LICENSE new file mode 100644 index 0000000..6eaaaae --- /dev/null +++ b/LICENSE @@ -0,0 +1,21 @@ +MIT License + +Copyright (c) 2026 Kai Zhao + +Permission is hereby granted, free of charge, to any person obtaining a copy +of this software and associated documentation files (the "Software"), to deal +in the Software without restriction, including without limitation the rights +to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +copies of the Software, and to permit persons to whom the Software is +furnished to do so, subject to the following conditions: + +The above copyright notice and this permission notice shall be included in all +copies or substantial portions of the Software. + +THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +SOFTWARE. diff --git a/README.md b/README.md index ede098a..e077617 100644 --- a/README.md +++ b/README.md @@ -1,5 +1,8 @@ # TorchMorph +[![arXiv](https://img.shields.io/badge/arXiv-2608.24738-b31b1b.svg)](https://arxiv.org/abs/2608.24738) +[![Hugging Face](https://img.shields.io/badge/%F0%9F%A4%97%20Hugging%20Face-Paper-FFD21E.svg)](https://huggingface.co/papers/2608.24738) + > ⚡ CUDA-Accelerated, Batch-Parallel Morphological Transformations for PyTorch TorchMorph is a lightweight, extensible library that brings **GPU-accelerated morphological operations** into the PyTorch ecosystem. @@ -77,3 +80,10 @@ if torch.cuda.is_available(): else: print("CUDA not available; install verified for import only.") ``` + +--- + +## 📄 License + +TorchMorph is released under the [MIT License](LICENSE) — free for academic and +commercial use, provided "as is" without warranty of any kind. diff --git a/docs/binary_morphology.md b/docs/binary_morphology.md new file mode 100644 index 0000000..c015eb5 --- /dev/null +++ b/docs/binary_morphology.md @@ -0,0 +1,18 @@ +# Binary Morphology + +Binary operators treat nonzero values as foreground and expect tensors shaped +as `(B, C, Spatial...)`. + +::: torchmorph.binary_dilation + +::: torchmorph.binary_erosion + +::: torchmorph.binary_propagation + +::: torchmorph.binary_fill_holes + +::: torchmorph.binary_hit_or_miss + +::: torchmorph.binary_opening + +::: torchmorph.binary_closing diff --git a/docs/distance_transforms.md b/docs/distance_transforms.md new file mode 100644 index 0000000..87a44e8 --- /dev/null +++ b/docs/distance_transforms.md @@ -0,0 +1,7 @@ +# Distance Transforms + +::: torchmorph.euclidean_distance_transform + +::: torchmorph.chamfer_distance_transform + +::: torchmorph.brute_force_distance_transform diff --git a/docs/grayscale_morphology.md b/docs/grayscale_morphology.md new file mode 100644 index 0000000..5ff72e7 --- /dev/null +++ b/docs/grayscale_morphology.md @@ -0,0 +1,20 @@ +# Grayscale Morphology + +Grayscale operators process floating-point or integer-valued tensors shaped as +`(B, C, Spatial...)`. + +::: torchmorph.grey_dilation + +::: torchmorph.grey_erosion + +::: torchmorph.grey_opening + +::: torchmorph.grey_closing + +::: torchmorph.morphological_gradient + +::: torchmorph.morphological_laplace + +::: torchmorph.white_tophat + +::: torchmorph.black_tophat diff --git a/docs/index.md b/docs/index.md new file mode 100644 index 0000000..2d21f34 --- /dev/null +++ b/docs/index.md @@ -0,0 +1,48 @@ +# torchmorph + +[![arXiv](https://img.shields.io/badge/arXiv-2608.24738-b31b1b.svg)](https://arxiv.org/abs/2608.24738) +[![Hugging Face](https://img.shields.io/badge/%F0%9F%A4%97%20Hugging%20Face-Paper-FFD21E.svg)](https://huggingface.co/papers/2608.24738) + +GPU-accelerated morphological image processing operations for PyTorch. + +## Installation + +```bash +pip install torchmorph +``` + +## Quick Start + +```python +import torch +import torchmorph as tm + +x = torch.zeros(1, 1, 64, 64, device='cuda') +x[0, 0, 10:20, 10:20] = 1 + +dist = tm.euclidean_distance_transform(x) +dilated = tm.binary_dilation(x) +``` + +Spatial morphology and distance-transform functions accept CUDA tensors shaped +as `(B, C, Spatial...)`, with one to eight spatial dimensions. Binary operators +treat nonzero values as foreground. + +Optimal transport works on CPU and CUDA with batches of flattened histograms: + +```python +source = torch.rand(8, 64, device="cuda") +target = torch.rand(8, 64, device="cuda") + +solver = tm.SinkhornSolver(epsilon=1.0, max_iter=200) +cost = tm.build_cost_matrix((8, 8), device=source.device) +distances = solver(source, target, cost) +``` + +## API Reference + +- [Distance Transforms](distance_transforms.md) — Euclidean, Chamfer, and brute-force distance transforms +- [Structuring Elements](structuring_elements.md) — Multidimensional connectivity structures +- [Binary Morphology](binary_morphology.md) — Binary dilation, erosion, filling, and related operators +- [Grayscale Morphology](grayscale_morphology.md) — Grayscale dilation, erosion, gradients, and top-hats +- [Optimal Transport](optimal_transport.md) — Cost matrices and differentiable Sinkhorn transport diff --git a/docs/optimal_transport.md b/docs/optimal_transport.md new file mode 100644 index 0000000..c5a07bd --- /dev/null +++ b/docs/optimal_transport.md @@ -0,0 +1,14 @@ +# Optimal Transport + +TorchMorph provides a differentiable Sinkhorn solver that runs on CPU or CUDA. +CUDA `float32` inputs automatically use the fused implementation. + +::: torchmorph.build_cost_matrix + +::: torchmorph.SinkhornSolver + options: + members: + - data_preprocess + - forward + - plan + - potentials diff --git a/docs/structuring_elements.md b/docs/structuring_elements.md new file mode 100644 index 0000000..36e0493 --- /dev/null +++ b/docs/structuring_elements.md @@ -0,0 +1,8 @@ +# Structuring Elements + +Create and expand multidimensional connectivity structures for morphological +operations. + +::: torchmorph.generate_binary_structure + +::: torchmorph.iterate_structure diff --git a/docs/stylesheets/extra.css b/docs/stylesheets/extra.css new file mode 100644 index 0000000..812fb09 --- /dev/null +++ b/docs/stylesheets/extra.css @@ -0,0 +1,78 @@ +/* Keep API headings compact and let signatures read like Python definitions. */ +.md-typeset .doc-heading { + display: flex; + align-items: center; + gap: 0.55rem; + margin: 2.4rem 0 0.75rem; + padding-bottom: 0.55rem; + border-bottom: 1px solid var(--md-default-fg-color--lightest); +} + +.md-typeset .doc-heading .doc-object-name { + color: var(--md-default-fg-color); + font-family: var(--md-code-font-family); + font-size: 0.82em; + font-weight: 700; + letter-spacing: -0.015em; +} + +.md-typeset .doc-heading code.doc-symbol { + flex: 0 0 auto; + border-radius: 999px; + padding: 0.2em 0.55em; + font-size: 0.62rem; + line-height: 1.5; +} + +/* mkdocstrings emits separate signatures as highlighted pre blocks. */ +.md-typeset .doc-object > .highlight, +.md-typeset .doc-overloads + .highlight { + margin: 0 0 1.25rem; + border: 1px solid var(--md-default-fg-color--lightest); + border-radius: 0.55rem; + background: var(--md-code-bg-color); + box-shadow: 0 0.08rem 0.22rem rgb(0 0 0 / 5%); + overflow: hidden; +} + +.md-typeset .doc-object > .highlight pre, +.md-typeset .doc-overloads + .highlight pre { + margin: 0; + padding: 0.85rem 1rem; + background: transparent; + overflow-x: auto; + font-size: 0.76rem; + line-height: 1.55; +} + +.md-typeset .doc-object > .highlight code, +.md-typeset .doc-overloads + .highlight code { + background: transparent; +} + +.md-typeset .doc-contents { + padding-left: 1rem; + border-left: 0.12rem solid var(--md-default-fg-color--lightest); +} + +.md-typeset .doc-short-anchor { + display: block; + height: 0; + scroll-margin-top: 4.8rem; +} + +@media screen and (max-width: 44.9844em) { + .md-typeset .doc-heading { + margin-top: 1.8rem; + } + + .md-typeset .doc-object > .highlight pre, + .md-typeset .doc-overloads + .highlight pre { + padding: 0.7rem 0.8rem; + font-size: 0.7rem; + } + + .md-typeset .doc-contents { + padding-left: 0.7rem; + } +} diff --git a/mkdocs.yml b/mkdocs.yml new file mode 100644 index 0000000..3a581e9 --- /dev/null +++ b/mkdocs.yml @@ -0,0 +1,62 @@ +site_name: torchmorph +site_url: https://intcomp.github.io/tm/ +site_description: GPU-accelerated morphological operations for PyTorch +repo_url: https://github.com/intcomp/torchmorph +repo_name: intcomp/torchmorph + +theme: + name: material + palette: + - scheme: default + primary: indigo + toggle: + icon: material/brightness-7 + name: Switch to dark mode + - scheme: slate + primary: indigo + toggle: + icon: material/brightness-4 + name: Switch to light mode + features: + - navigation.sections + - content.code.copy + +extra_css: + - stylesheets/extra.css + +hooks: + - mkdocs_hooks.py + +plugins: + - search + - mkdocstrings: + handlers: + python: + options: + docstring_style: google + show_source: false + show_root_heading: true + heading_level: 3 + show_symbol_type_heading: true + merge_init_into_class: true + separate_signature: true + line_length: 88 + +markdown_extensions: + - toc: + toc_depth: 3 + - pymdownx.highlight: + anchor_linenums: true + - pymdownx.superfences + - pymdownx.inlinehilite + - admonition + - pymdownx.details + +nav: + - Home: index.md + - API Reference: + - Distance Transforms: distance_transforms.md + - Structuring Elements: structuring_elements.md + - Binary Morphology: binary_morphology.md + - Grayscale Morphology: grayscale_morphology.md + - Optimal Transport: optimal_transport.md diff --git a/mkdocs_hooks.py b/mkdocs_hooks.py new file mode 100644 index 0000000..09ebeab --- /dev/null +++ b/mkdocs_hooks.py @@ -0,0 +1,16 @@ +"""Small build-time enhancements for the generated API reference.""" + +import re + +_API_HEADING = re.compile(r'(]*>)') + + +def on_page_content(html: str, **kwargs) -> str: + """Add a short anchor alias before every mkdocstrings API heading.""" + + def add_short_anchor(match: re.Match[str]) -> str: + short_name = match.group(2).rsplit(".", 1)[-1] + anchor = f'' + return f"{anchor}\n{match.group(1)}" + + return _API_HEADING.sub(add_short_anchor, html) diff --git a/requirements-docs.txt b/requirements-docs.txt new file mode 100644 index 0000000..b8b3cf2 --- /dev/null +++ b/requirements-docs.txt @@ -0,0 +1,5 @@ +-r requirements-test.txt +mkdocs-material +mkdocstrings[python] +pymdown-extensions + diff --git a/setup.py b/setup.py index ad0712d..1a52d5e 100644 --- a/setup.py +++ b/setup.py @@ -26,6 +26,8 @@ def get_extensions(): version="0.1.0", author="Kai Zhao", description="CUDA-accelerated morphological transformations for PyTorch", + license="MIT", + license_files=["LICENSE"], url="https://github.com/kaizhao-shu/torchmorph", packages=find_packages(), python_requires=">=3.12", diff --git a/torchmorph/distance_transform.py b/torchmorph/distance_transform.py index f5ebf3e..074c82b 100644 --- a/torchmorph/distance_transform.py +++ b/torchmorph/distance_transform.py @@ -77,7 +77,43 @@ def euclidean_distance_transform( distances: Tensor | None = None, indices: Tensor | None = None, ) -> Tensor | tuple[Tensor, Tensor] | None: - """Euclidean distance transform for (B, C, Spatial...) CUDA tensors.""" + """Compute the exact Euclidean distance to the nearest background point + + Distances are computed independently for each batch and channel. Nonzero + input values are foreground and zero values are background. + + Args: + input (torch.Tensor): CUDA tensor with shape ``(B, C, Spatial...)``. + sampling (float or sequence[float], optional): Positive spacing along + each spatial axis. A scalar or one-element sequence is broadcast. + return_distances (bool): Return a newly allocated distance tensor unless + ``distances`` is supplied. + return_indices (bool): Return a newly allocated nearest-background index + tensor unless ``indices`` is supplied. + distances (torch.Tensor, optional): Preallocated distance output with + the same shape and device as ``input``. Supplying it enables + distance computation regardless of ``return_distances``. + indices (torch.Tensor, optional): Preallocated index output with shape + ``(spatial_ndim, *input.shape)`` on the input device. Supplying it + enables index computation regardless of ``return_indices``. + + Returns: + torch.Tensor, tuple[torch.Tensor, torch.Tensor], or None: Requested + newly allocated outputs. Distances have ``float32`` dtype and input + shape; indices have ``int32`` dtype and shape + ``(spatial_ndim, *input.shape)``. Preallocated outputs are filled but + omitted from the return value, so the function returns ``None`` when + every requested output is preallocated. + + Example: + ```pycon + >>> import torch + >>> import torchmorph as tm + >>> x = torch.tensor([[[0, 1, 1, 0]]], device="cuda") + >>> tm.euclidean_distance_transform(x) + tensor([[[0., 1., 1., 0.]]], device='cuda:0') + ``` + """ spatial_ndim, return_distances, return_indices = _prepare_distance_transform( input, return_distances, @@ -110,7 +146,39 @@ def chamfer_distance_transform( distances: Tensor | None = None, indices: Tensor | None = None, ) -> Tensor | tuple[Tensor, Tensor] | None: - """Chamfer distance transform for (B, C, Spatial...) CUDA tensors.""" + """Compute a chamfer distance to the nearest background point + + Distances are computed independently for each batch and channel. Nonzero + input values are foreground and zero values are background. + + Args: + input (torch.Tensor): CUDA tensor with shape ``(B, C, Spatial...)``. + metric (str): ``"chessboard"`` or ``"taxicab"``. ``"cityblock"`` and + ``"manhattan"`` are accepted aliases for ``"taxicab"``. + return_distances (bool): Return a newly allocated distance tensor unless + ``distances`` is supplied. + return_indices (bool): Return a newly allocated nearest-background index + tensor unless ``indices`` is supplied. + distances (torch.Tensor, optional): Preallocated distance output with + the same shape and device as ``input``. + indices (torch.Tensor, optional): Preallocated index output with shape + ``(spatial_ndim, *input.shape)`` on the input device. + + Returns: + torch.Tensor, tuple[torch.Tensor, torch.Tensor], or None: Requested + newly allocated outputs. Distances are ``float32`` and indices are + ``int32``. Preallocated outputs are filled but omitted from the return + value. + + Example: + ```pycon + >>> import torch + >>> import torchmorph as tm + >>> x = torch.tensor([[[0, 1, 1, 0]]], device="cuda") + >>> tm.chamfer_distance_transform(x, metric="taxicab") + tensor([[[0., 1., 1., 0.]]], device='cuda:0') + ``` + """ _, return_distances, return_indices = _prepare_distance_transform( input, return_distances, @@ -147,7 +215,42 @@ def brute_force_distance_transform( distances: Tensor | None = None, indices: Tensor | None = None, ) -> Tensor | tuple[Tensor, Tensor] | None: - """Brute-force distance transform for (B, C, Spatial...) CUDA tensors.""" + """Compute distances to background points by exhaustive search + + This reference-style transform supports Euclidean, taxicab, and chessboard + metrics and computes each batch and channel independently. + + Args: + input (torch.Tensor): CUDA tensor with shape ``(B, C, Spatial...)``; + nonzero values are foreground and zeros are background. + metric (str): One of ``"euclidean"``, ``"taxicab"``, or + ``"chessboard"``. + sampling (float or sequence[float], optional): Positive spatial spacing. + A scalar or one-element sequence is broadcast to every axis. + return_distances (bool): Return a newly allocated distance tensor unless + ``distances`` is supplied. + return_indices (bool): Return a newly allocated nearest-background index + tensor unless ``indices`` is supplied. + distances (torch.Tensor, optional): Preallocated distance output with + the same shape and device as ``input``. + indices (torch.Tensor, optional): Preallocated index output with shape + ``(spatial_ndim, *input.shape)`` on the input device. + + Returns: + torch.Tensor, tuple[torch.Tensor, torch.Tensor], or None: Requested + newly allocated outputs. Distances are ``float32`` and indices are + ``int32``. Preallocated outputs are filled but omitted from the return + value. + + Example: + ```pycon + >>> import torch + >>> import torchmorph as tm + >>> x = torch.tensor([[[0, 1, 1, 0]]], device="cuda") + >>> tm.brute_force_distance_transform(x, metric="euclidean") + tensor([[[0., 1., 1., 0.]]], device='cuda:0') + ``` + """ spatial_ndim, return_distances, return_indices = _prepare_distance_transform( input, return_distances, diff --git a/torchmorph/morphology/binary.py b/torchmorph/morphology/binary.py index d47a0c1..fac2b8d 100644 --- a/torchmorph/morphology/binary.py +++ b/torchmorph/morphology/binary.py @@ -1,290 +1,509 @@ -import torch -from torch import Tensor - -from .. import _C -from .._validation import validate_bcs_input, validate_output -from .structure import _normalize_origin, _validate_origin, generate_binary_structure - - -def _normalize_structure(structure: Tensor, spatial_ndim: int, name: str = "structure") -> Tensor: - if structure.ndim != spatial_ndim: - raise ValueError(f"{name} dimension is not {spatial_ndim}, got {structure.ndim}") - return (structure != 0).detach().to(device="cpu", dtype=torch.bool).contiguous() - - -def _binary_morphology_cuda_step( - input: Tensor, - structure: Tensor, - border_value: bool, - origin: tuple[int, ...], - *, - mode: str, -) -> Tensor: - x = input if input.dtype == torch.bool and input.is_contiguous() else (input != 0).contiguous() - kernel = _C.binary_erosion_cuda if mode == "erosion" else _C.binary_dilation_cuda - return kernel(x, structure, list(origin), bool(border_value)) - - -def _binary_morphology( - input: Tensor, - structure: Tensor | None, - iterations: int, - mask: Tensor | None, - output: Tensor | None, - border_value: bool, - origin: int | tuple[int, ...], - *, - mode: str, -) -> Tensor: - iterate_until_stable = iterations < 1 - spatial_ndim = validate_bcs_input(input) - validate_output(input, output) - - if structure is None: - structure = generate_binary_structure(spatial_ndim, 1) - structure = _normalize_structure(structure, spatial_ndim) - - origin = _normalize_origin(origin, spatial_ndim) - _validate_origin(origin, structure) - x = input != 0 - input_bool = x - if mask is not None and mask.shape != input.shape: - raise ValueError(f"mask shape {mask.shape} must match input shape {input.shape}") - mask_bool = mask.to(device=input.device, dtype=torch.bool) if mask is not None else None - structure_is_empty = not structure.any().item() - - def step(value: Tensor) -> Tensor: - if structure_is_empty: - return torch.full_like(value, mode == "erosion", dtype=torch.bool) - return _binary_morphology_cuda_step( - value, - structure, - border_value, - origin, - mode=mode, - ) - - if iterate_until_stable: - old = None - while True: - x = step(x) - if mask_bool is not None: - x = torch.where(mask_bool, x, input_bool) - - if old is not None and torch.equal(x, old): - break - - old = x - else: - for _ in range(iterations): - x = step(x) - if mask_bool is not None: - x = torch.where(mask_bool, x, input_bool) - - result = x.to(dtype=torch.bool) - if output is not None: - output.copy_(result) - return output - return result - - -def binary_erosion( - input: Tensor, - structure: Tensor | None = None, - iterations: int = 1, - mask: Tensor | None = None, - output: Tensor | None = None, - border_value: bool = False, - origin: int | tuple[int, ...] = 0, +import torch +from torch import Tensor + +from .. import _C +from .._validation import validate_bcs_input, validate_output +from .structure import _normalize_origin, _validate_origin, generate_binary_structure + + +def _normalize_structure(structure: Tensor, spatial_ndim: int, name: str = "structure") -> Tensor: + if structure.ndim != spatial_ndim: + raise ValueError(f"{name} dimension is not {spatial_ndim}, got {structure.ndim}") + return (structure != 0).detach().to(device="cpu", dtype=torch.bool).contiguous() + + +def _binary_morphology_cuda_step( + input: Tensor, + structure: Tensor, + border_value: bool, + origin: tuple[int, ...], + *, + mode: str, +) -> Tensor: + x = input if input.dtype == torch.bool and input.is_contiguous() else (input != 0).contiguous() + kernel = _C.binary_erosion_cuda if mode == "erosion" else _C.binary_dilation_cuda + return kernel(x, structure, list(origin), bool(border_value)) + + +def _binary_morphology( + input: Tensor, + structure: Tensor | None, + iterations: int, + mask: Tensor | None, + output: Tensor | None, + border_value: bool, + origin: int | tuple[int, ...], + *, + mode: str, +) -> Tensor: + iterate_until_stable = iterations < 1 + spatial_ndim = validate_bcs_input(input) + validate_output(input, output) + + if structure is None: + structure = generate_binary_structure(spatial_ndim, 1) + structure = _normalize_structure(structure, spatial_ndim) + + origin = _normalize_origin(origin, spatial_ndim) + _validate_origin(origin, structure) + x = input != 0 + input_bool = x + if mask is not None and mask.shape != input.shape: + raise ValueError(f"mask shape {mask.shape} must match input shape {input.shape}") + mask_bool = mask.to(device=input.device, dtype=torch.bool) if mask is not None else None + structure_is_empty = not structure.any().item() + + def step(value: Tensor) -> Tensor: + if structure_is_empty: + return torch.full_like(value, mode == "erosion", dtype=torch.bool) + return _binary_morphology_cuda_step( + value, + structure, + border_value, + origin, + mode=mode, + ) + + if iterate_until_stable: + old = None + while True: + x = step(x) + if mask_bool is not None: + x = torch.where(mask_bool, x, input_bool) + + if old is not None and torch.equal(x, old): + break + + old = x + else: + for _ in range(iterations): + x = step(x) + if mask_bool is not None: + x = torch.where(mask_bool, x, input_bool) + + result = x.to(dtype=torch.bool) + if output is not None: + output.copy_(result) + return output + return result + + +def binary_erosion( + input: Tensor, + structure: Tensor | None = None, + iterations: int = 1, + mask: Tensor | None = None, + output: Tensor | None = None, + border_value: bool = False, + origin: int | tuple[int, ...] = 0, ) -> Tensor: + """Erode binary objects in a batched CUDA tensor + + A value is foreground when it is nonzero. The operation supports one to + eight spatial dimensions and always treats the first two dimensions as + batch and channel. When ``iterations`` is less than one, erosion continues + until the result no longer changes. + + Args: + input (torch.Tensor): CUDA tensor with shape ``(B, C, Spatial...)``. + structure (torch.Tensor, optional): Spatial structuring element. Nonzero + entries participate in the neighborhood. If ``None``, uses + :func:`generate_binary_structure` with connectivity ``1``. + iterations (int): Number of sequential erosions. Values less than ``1`` + repeat the operation until convergence. + mask (torch.Tensor, optional): Tensor with the same shape as ``input``. + Only locations where the mask is nonzero may change. + output (torch.Tensor, optional): Preallocated tensor with the same shape + and device as ``input``. When supplied, receives and is returned as + the result. + border_value (bool): Value used outside the input boundary. + origin (int or tuple[int, ...]): Offset of the structuring-element + anchor. A scalar is applied to every spatial dimension. + + Returns: + torch.Tensor: Boolean erosion result, or ``output`` when it is supplied. + + Example: + ```pycon + >>> import torch + >>> import torchmorph as tm + >>> x = torch.ones((1, 1, 5, 5), device="cuda") + >>> y = tm.binary_erosion(x, structure=torch.ones(3, 3)) + >>> y.dtype, y.sum().item() + (torch.bool, 9) + ``` """ - N-dimensional binary erosion for `(B, C, Spatial...)` CUDA tensors. - - For a single image or volume, add batch and channel dimensions first. - """ - return _binary_morphology( - input=input, - structure=structure, - iterations=iterations, - mask=mask, - output=output, - border_value=border_value, - origin=origin, - mode="erosion", - ) - - -def binary_dilation( - input: Tensor, - structure: Tensor | None = None, - iterations: int = 1, - mask: Tensor | None = None, - output: Tensor | None = None, - border_value: bool = False, - origin: int | tuple[int, ...] = 0, + return _binary_morphology( + input=input, + structure=structure, + iterations=iterations, + mask=mask, + output=output, + border_value=border_value, + origin=origin, + mode="erosion", + ) + + +def binary_dilation( + input: Tensor, + structure: Tensor | None = None, + iterations: int = 1, + mask: Tensor | None = None, + output: Tensor | None = None, + border_value: bool = False, + origin: int | tuple[int, ...] = 0, ) -> Tensor: - """N-dimensional binary dilation for `(B, C, Spatial...)` CUDA tensors.""" - return _binary_morphology( - input, - structure, - iterations, - mask, - output, - border_value, - origin, - mode="dilation", - ) - - -def binary_propagation( - input: Tensor, - structure: Tensor | None = None, - mask: Tensor | None = None, - output: Tensor | None = None, - border_value: bool = False, - origin: int | tuple[int, ...] = 0, + """Dilate binary objects in a batched CUDA tensor + + A value is foreground when it is nonzero. The operation supports one to + eight spatial dimensions and always treats the first two dimensions as + batch and channel. When ``iterations`` is less than one, dilation continues + until the result no longer changes. + + Args: + input (torch.Tensor): CUDA tensor with shape ``(B, C, Spatial...)``. + structure (torch.Tensor, optional): Spatial structuring element. Nonzero + entries participate in the neighborhood. If ``None``, uses + :func:`generate_binary_structure` with connectivity ``1``. + iterations (int): Number of sequential dilations. Values less than ``1`` + repeat the operation until convergence. + mask (torch.Tensor, optional): Tensor with the same shape as ``input``. + Only locations where the mask is nonzero may change. + output (torch.Tensor, optional): Preallocated tensor with the same shape + and device as ``input``. When supplied, receives and is returned as + the result. + border_value (bool): Value used outside the input boundary. + origin (int or tuple[int, ...]): Offset of the structuring-element + anchor. A scalar is applied to every spatial dimension. + + Returns: + torch.Tensor: Boolean dilation result, or ``output`` when it is supplied. + + Example: + ```pycon + >>> import torch + >>> import torchmorph as tm + >>> x = torch.zeros((1, 1, 5, 5), device="cuda") + >>> x[0, 0, 2, 2] = 1 + >>> y = tm.binary_dilation(x, structure=torch.ones(3, 3)) + >>> y.dtype, y.sum().item() + (torch.bool, 9) + ``` + """ + return _binary_morphology( + input, + structure, + iterations, + mask, + output, + border_value, + origin, + mode="dilation", + ) + + +def binary_propagation( + input: Tensor, + structure: Tensor | None = None, + mask: Tensor | None = None, + output: Tensor | None = None, + border_value: bool = False, + origin: int | tuple[int, ...] = 0, ) -> Tensor: - """N-dimensional binary propagation for `(B, C, Spatial...)` CUDA tensors.""" - return _binary_morphology( - input, - structure, - -1, - mask, - output, - border_value, - origin, - mode="dilation", - ) - - -def binary_fill_holes( - input: Tensor, - structure: Tensor | None = None, - output: Tensor | None = None, - origin: int | tuple[int, ...] = 0, + """Propagate a binary seed through a mask until convergence + + This repeatedly dilates ``input`` while allowing changes only where + ``mask`` is nonzero. It is useful for morphological reconstruction. + + Args: + input (torch.Tensor): Binary seed CUDA tensor with shape + ``(B, C, Spatial...)``; nonzero values are foreground. + structure (torch.Tensor, optional): Spatial structuring element. If + ``None``, uses connectivity ``1``. + mask (torch.Tensor, optional): Tensor with the same shape as ``input``. + Only nonzero mask locations may change. If ``None``, propagation is + unconstrained. + output (torch.Tensor, optional): Preallocated result tensor with the + same shape and device as ``input``. + border_value (bool): Value used outside the input boundary. + origin (int or tuple[int, ...]): Structuring-element anchor offset. + + Returns: + torch.Tensor: Boolean converged result, or ``output`` when supplied. + + Example: + ```pycon + >>> import torch + >>> import torchmorph as tm + >>> seed = torch.zeros((1, 1, 5, 5), device="cuda") + >>> seed[0, 0, 2, 2] = 1 + >>> mask = torch.ones_like(seed) + >>> tm.binary_propagation(seed, mask=mask).all().item() + True + ``` + """ + return _binary_morphology( + input, + structure, + -1, + mask, + output, + border_value, + origin, + mode="dilation", + ) + + +def binary_fill_holes( + input: Tensor, + structure: Tensor | None = None, + output: Tensor | None = None, + origin: int | tuple[int, ...] = 0, ) -> Tensor: - """Fill holes in binary objects for `(B, C, Spatial...)` CUDA tensors.""" - validate_bcs_input(input) - validate_output(input, output) - - mask = input == 0 - seed = torch.zeros_like(mask, dtype=torch.bool) - background = binary_propagation( - seed, - structure=structure, - mask=mask, - output=None, - border_value=True, - origin=origin, - ) - result = torch.logical_not(background) - - if output is not None: - output.copy_(result) - return output - return result - - -def binary_hit_or_miss( - input: Tensor, - structure1: Tensor | None = None, - structure2: Tensor | None = None, - output: Tensor | None = None, - origin1: int | tuple[int, ...] = 0, - origin2: int | tuple[int, ...] | None = None, + """Fill holes enclosed by binary objects + + Background connected to the tensor boundary is reconstructed and removed; + enclosed background regions become foreground. + + Args: + input (torch.Tensor): CUDA tensor with shape ``(B, C, Spatial...)``; + nonzero values are foreground. + structure (torch.Tensor, optional): Spatial structuring element used to + determine background connectivity. If ``None``, uses connectivity + ``1``. + output (torch.Tensor, optional): Preallocated result tensor with the + same shape and device as ``input``. + origin (int or tuple[int, ...]): Structuring-element anchor offset. + + Returns: + torch.Tensor: Boolean tensor with enclosed holes filled, or ``output``. + + Example: + ```pycon + >>> import torch + >>> import torchmorph as tm + >>> x = torch.ones((1, 1, 5, 5), device="cuda") + >>> x[0, 0, 2, 2] = 0 + >>> tm.binary_fill_holes(x)[0, 0, 2, 2].item() + True + ``` + """ + validate_bcs_input(input) + validate_output(input, output) + + mask = input == 0 + seed = torch.zeros_like(mask, dtype=torch.bool) + background = binary_propagation( + seed, + structure=structure, + mask=mask, + output=None, + border_value=True, + origin=origin, + ) + result = torch.logical_not(background) + + if output is not None: + output.copy_(result) + return output + return result + + +def binary_hit_or_miss( + input: Tensor, + structure1: Tensor | None = None, + structure2: Tensor | None = None, + output: Tensor | None = None, + origin1: int | tuple[int, ...] = 0, + origin2: int | tuple[int, ...] | None = None, ) -> Tensor: - """N-dimensional binary hit-or-miss transform for `(B, C, Spatial...)` CUDA tensors.""" - spatial_ndim = validate_bcs_input(input) - validate_output(input, output) - origin1 = _normalize_origin(origin1, spatial_ndim) - origin2 = origin1 if origin2 is None else _normalize_origin(origin2, spatial_ndim) - - if structure1 is None: - structure1 = generate_binary_structure(spatial_ndim, 1) - structure1 = _normalize_structure(structure1, spatial_ndim, "structure1") - _validate_origin(origin1, structure1, "origin1") - - if structure2 is None: - structure2 = torch.logical_not(structure1).contiguous() - else: - structure2 = _normalize_structure(structure2, spatial_ndim, "structure2") - _validate_origin(origin2, structure2, "origin2") - - input_bool = input != 0 - if structure1.any().item(): - hit = binary_erosion(input_bool, structure=structure1, origin=origin1) - else: - hit = torch.ones_like(input_bool, dtype=torch.bool) - - if structure2.any().item(): - miss = binary_erosion(input_bool == 0, structure=structure2, origin=origin2) - result = torch.logical_and(hit, miss) - else: - result = hit - - if output is not None: - output.copy_(result) - return output - return result - - -def binary_opening( - input: Tensor, - structure: Tensor | None = None, - iterations: int = 1, - mask: Tensor | None = None, - output: Tensor | None = None, - border_value: bool = False, - origin: int | tuple[int, ...] = 0, + """Find binary configurations with the hit-or-miss transform + + A location matches when ``structure1`` fits the foreground and + ``structure2`` fits the background at the same anchor. + + Args: + input (torch.Tensor): CUDA tensor with shape ``(B, C, Spatial...)``; + nonzero values are foreground. + structure1 (torch.Tensor, optional): Foreground structuring element. If + ``None``, uses connectivity ``1``. + structure2 (torch.Tensor, optional): Background structuring element. If + ``None``, uses the logical complement of ``structure1``. + output (torch.Tensor, optional): Preallocated result tensor with the + same shape and device as ``input``. + origin1 (int or tuple[int, ...]): Anchor offset for ``structure1``. + origin2 (int or tuple[int, ...], optional): Anchor offset for + ``structure2``. If ``None``, uses ``origin1``. + + Returns: + torch.Tensor: Boolean tensor marking matching locations, or ``output``. + + Example: + ```pycon + >>> import torch + >>> import torchmorph as tm + >>> x = torch.zeros((1, 1, 3, 3), device="cuda") + >>> x[0, 0, 1, 1] = 1 + >>> hit = tm.binary_hit_or_miss(x, structure1=torch.ones(1, 1)) + >>> hit[0, 0, 1, 1].item() + True + ``` + """ + spatial_ndim = validate_bcs_input(input) + validate_output(input, output) + origin1 = _normalize_origin(origin1, spatial_ndim) + origin2 = origin1 if origin2 is None else _normalize_origin(origin2, spatial_ndim) + + if structure1 is None: + structure1 = generate_binary_structure(spatial_ndim, 1) + structure1 = _normalize_structure(structure1, spatial_ndim, "structure1") + _validate_origin(origin1, structure1, "origin1") + + if structure2 is None: + structure2 = torch.logical_not(structure1).contiguous() + else: + structure2 = _normalize_structure(structure2, spatial_ndim, "structure2") + _validate_origin(origin2, structure2, "origin2") + + input_bool = input != 0 + if structure1.any().item(): + hit = binary_erosion(input_bool, structure=structure1, origin=origin1) + else: + hit = torch.ones_like(input_bool, dtype=torch.bool) + + if structure2.any().item(): + miss = binary_erosion(input_bool == 0, structure=structure2, origin=origin2) + result = torch.logical_and(hit, miss) + else: + result = hit + + if output is not None: + output.copy_(result) + return output + return result + + +def binary_opening( + input: Tensor, + structure: Tensor | None = None, + iterations: int = 1, + mask: Tensor | None = None, + output: Tensor | None = None, + border_value: bool = False, + origin: int | tuple[int, ...] = 0, ) -> Tensor: - """N-dimensional binary opening for `(B, C, Spatial...)` CUDA tensors.""" - x = binary_erosion( - input, - structure, - iterations, - mask, - None, - border_value, - origin, - ) - x = binary_dilation( - x, - structure, - iterations, - mask, - output, - border_value, - origin, - ) - return x - - -def binary_closing( - input: Tensor, - structure: Tensor | None = None, - iterations: int = 1, - mask: Tensor | None = None, - output: Tensor | None = None, - border_value: bool = False, - origin: int | tuple[int, ...] = 0, + """Open binary objects by erosion followed by dilation + + Opening removes foreground features that cannot contain the structuring + element while largely preserving larger objects. + + Args: + input (torch.Tensor): CUDA tensor with shape ``(B, C, Spatial...)``; + nonzero values are foreground. + structure (torch.Tensor, optional): Spatial structuring element. If + ``None``, uses connectivity ``1``. + iterations (int): Number of erosions followed by the same number of + dilations. Values less than ``1`` iterate each stage to convergence. + mask (torch.Tensor, optional): Same-shaped tensor whose nonzero entries + identify locations that may change. + output (torch.Tensor, optional): Preallocated result tensor with the + same shape and device as ``input``. + border_value (bool): Value used outside the input boundary. + origin (int or tuple[int, ...]): Structuring-element anchor offset. + + Returns: + torch.Tensor: Boolean opened tensor, or ``output`` when supplied. + + Example: + ```pycon + >>> import torch + >>> import torchmorph as tm + >>> x = torch.zeros((1, 1, 5, 5), device="cuda") + >>> x[0, 0, 1:4, 1:4] = 1 + >>> x[0, 0, 0, 0] = 1 + >>> tm.binary_opening(x).sum().item() + 5 + ``` + """ + x = binary_erosion( + input, + structure, + iterations, + mask, + None, + border_value, + origin, + ) + x = binary_dilation( + x, + structure, + iterations, + mask, + output, + border_value, + origin, + ) + return x + + +def binary_closing( + input: Tensor, + structure: Tensor | None = None, + iterations: int = 1, + mask: Tensor | None = None, + output: Tensor | None = None, + border_value: bool = False, + origin: int | tuple[int, ...] = 0, ) -> Tensor: - """N-dimensional binary closing for `(B, C, Spatial...)` CUDA tensors.""" - x = binary_dilation( - input, - structure, - iterations, - mask, - None, - border_value, - origin, - ) - x = binary_erosion( - x, - structure, - iterations, - mask, - output, - border_value, - origin, - ) - return x + """Close binary objects by dilation followed by erosion + + Closing fills small gaps and joins nearby foreground regions while largely + preserving the extent of larger objects. + + Args: + input (torch.Tensor): CUDA tensor with shape ``(B, C, Spatial...)``; + nonzero values are foreground. + structure (torch.Tensor, optional): Spatial structuring element. If + ``None``, uses connectivity ``1``. + iterations (int): Number of dilations followed by the same number of + erosions. Values less than ``1`` iterate each stage to convergence. + mask (torch.Tensor, optional): Same-shaped tensor whose nonzero entries + identify locations that may change. + output (torch.Tensor, optional): Preallocated result tensor with the + same shape and device as ``input``. + border_value (bool): Value used outside the input boundary. + origin (int or tuple[int, ...]): Structuring-element anchor offset. + + Returns: + torch.Tensor: Boolean closed tensor, or ``output`` when supplied. + + Example: + ```pycon + >>> import torch + >>> import torchmorph as tm + >>> x = torch.zeros((1, 1, 5, 5), device="cuda") + >>> x[0, 0, 2, 1] = x[0, 0, 2, 3] = 1 + >>> tm.binary_closing(x)[0, 0, 2, 2].item() + True + ``` + """ + x = binary_dilation( + input, + structure, + iterations, + mask, + None, + border_value, + origin, + ) + x = binary_erosion( + x, + structure, + iterations, + mask, + output, + border_value, + origin, + ) + return x diff --git a/torchmorph/morphology/grey.py b/torchmorph/morphology/grey.py index 0f53b37..9f55494 100644 --- a/torchmorph/morphology/grey.py +++ b/torchmorph/morphology/grey.py @@ -115,9 +115,37 @@ def grey_erosion( cval: float = 0.0, origin: int | tuple[int, ...] = 0, ) -> Tensor: - """N-dimensional grey erosion for ``(B, C, Spatial...)`` CUDA tensors. - - Computation uses float32. The result is float32 unless ``output`` is given. + """Erode a batched grayscale CUDA tensor + + Computes the local minimum after subtracting the non-flat ``structure``. + Computation uses ``float32`` and supports one to eight spatial dimensions. + + Args: + input (torch.Tensor): CUDA tensor with shape ``(B, C, Spatial...)``. + size (int or tuple[int, ...], optional): Shape of a flat, full + structuring element. A scalar is broadcast to every spatial axis. + footprint (torch.Tensor, optional): Boolean mask selecting participating + positions. Must match ``structure`` when both are supplied. + structure (torch.Tensor, optional): Non-flat additive structuring + element. It takes precedence over ``footprint`` and ``size``. + output (torch.Tensor, optional): Preallocated result tensor with the + same shape and device as ``input``. + mode (str): Boundary mode: ``"reflect"``, ``"constant"``, ``"nearest"``, + ``"mirror"``, or ``"wrap"``. + cval (float): Boundary value used when ``mode="constant"``. + origin (int or tuple[int, ...]): Structuring-element anchor offset. + + Returns: + torch.Tensor: Eroded tensor in ``float32``, or ``output`` when supplied. + + Example: + ```pycon + >>> import torch + >>> import torchmorph as tm + >>> x = torch.arange(9, device="cuda").reshape(1, 1, 3, 3) + >>> tm.grey_erosion(x, size=3).shape + torch.Size([1, 1, 3, 3]) + ``` """ return _grey_morphology( input, @@ -142,9 +170,37 @@ def grey_dilation( cval: float = 0.0, origin: int | tuple[int, ...] = 0, ) -> Tensor: - """N-dimensional grey dilation for ``(B, C, Spatial...)`` CUDA tensors. - - Computation uses float32. The result is float32 unless ``output`` is given. + """Dilate a batched grayscale CUDA tensor + + Computes the local maximum after adding the non-flat ``structure``. + Computation uses ``float32`` and supports one to eight spatial dimensions. + + Args: + input (torch.Tensor): CUDA tensor with shape ``(B, C, Spatial...)``. + size (int or tuple[int, ...], optional): Shape of a flat, full + structuring element. A scalar is broadcast to every spatial axis. + footprint (torch.Tensor, optional): Boolean mask selecting participating + positions. Must match ``structure`` when both are supplied. + structure (torch.Tensor, optional): Non-flat additive structuring + element. It takes precedence over ``footprint`` and ``size``. + output (torch.Tensor, optional): Preallocated result tensor with the + same shape and device as ``input``. + mode (str): Boundary mode: ``"reflect"``, ``"constant"``, ``"nearest"``, + ``"mirror"``, or ``"wrap"``. + cval (float): Boundary value used when ``mode="constant"``. + origin (int or tuple[int, ...]): Structuring-element anchor offset. + + Returns: + torch.Tensor: Dilated tensor in ``float32``, or ``output`` when supplied. + + Example: + ```pycon + >>> import torch + >>> import torchmorph as tm + >>> x = torch.arange(9, device="cuda").reshape(1, 1, 3, 3) + >>> tm.grey_dilation(x, size=3).dtype + torch.float32 + ``` """ return _grey_morphology( input, @@ -169,9 +225,38 @@ def grey_opening( cval: float = 0.0, origin: int | tuple[int, ...] = 0, ) -> Tensor: - """N-dimensional grey opening for ``(B, C, Spatial...)`` CUDA tensors. - - Computation uses float32. The result is float32 unless ``output`` is given. + """Open a grayscale tensor by erosion followed by dilation + + Opening suppresses bright features smaller than the selected structuring + element. Computation uses ``float32``. + + Args: + input (torch.Tensor): CUDA tensor with shape ``(B, C, Spatial...)``. + size (int or tuple[int, ...], optional): Shape of a flat, full + structuring element. + footprint (torch.Tensor, optional): Boolean mask selecting participating + positions. + structure (torch.Tensor, optional): Non-flat additive structuring + element. It takes precedence over ``footprint`` and ``size``. + output (torch.Tensor, optional): Preallocated result tensor with the + same shape and device as ``input``. + mode (str): Boundary mode: ``"reflect"``, ``"constant"``, ``"nearest"``, + ``"mirror"``, or ``"wrap"``. + cval (float): Boundary value used when ``mode="constant"``. + origin (int or tuple[int, ...]): Structuring-element anchor offset. + + Returns: + torch.Tensor: Opened tensor in ``float32``, or ``output`` when supplied. + + Example: + ```pycon + >>> import torch + >>> import torchmorph as tm + >>> x = torch.zeros((1, 1, 5, 5), device="cuda") + >>> x[0, 0, 2, 2] = 1 + >>> tm.grey_opening(x, size=3).max().item() + 0.0 + ``` """ kwargs = _element_kwargs(size, footprint, structure, mode, cval, origin) eroded = grey_erosion(input, **kwargs) @@ -188,9 +273,38 @@ def grey_closing( cval: float = 0.0, origin: int | tuple[int, ...] = 0, ) -> Tensor: - """N-dimensional grey closing for ``(B, C, Spatial...)`` CUDA tensors. - - Computation uses float32. The result is float32 unless ``output`` is given. + """Close a grayscale tensor by dilation followed by erosion + + Closing suppresses dark features smaller than the selected structuring + element. Computation uses ``float32``. + + Args: + input (torch.Tensor): CUDA tensor with shape ``(B, C, Spatial...)``. + size (int or tuple[int, ...], optional): Shape of a flat, full + structuring element. + footprint (torch.Tensor, optional): Boolean mask selecting participating + positions. + structure (torch.Tensor, optional): Non-flat additive structuring + element. It takes precedence over ``footprint`` and ``size``. + output (torch.Tensor, optional): Preallocated result tensor with the + same shape and device as ``input``. + mode (str): Boundary mode: ``"reflect"``, ``"constant"``, ``"nearest"``, + ``"mirror"``, or ``"wrap"``. + cval (float): Boundary value used when ``mode="constant"``. + origin (int or tuple[int, ...]): Structuring-element anchor offset. + + Returns: + torch.Tensor: Closed tensor in ``float32``, or ``output`` when supplied. + + Example: + ```pycon + >>> import torch + >>> import torchmorph as tm + >>> x = torch.ones((1, 1, 5, 5), device="cuda") + >>> x[0, 0, 2, 2] = 0 + >>> tm.grey_closing(x, size=3).min().item() + 1.0 + ``` """ kwargs = _element_kwargs(size, footprint, structure, mode, cval, origin) dilated = grey_dilation(input, **kwargs) @@ -207,7 +321,38 @@ def morphological_gradient( cval: float = 0.0, origin: int | tuple[int, ...] = 0, ) -> Tensor: - """N-dimensional morphological gradient for ``(B, C, Spatial...)`` CUDA tensors.""" + """Compute the morphological gradient of a grayscale tensor + + The gradient is ``dilation(input) - erosion(input)`` and emphasizes local + intensity transitions. + + Args: + input (torch.Tensor): CUDA tensor with shape ``(B, C, Spatial...)``. + size (int or tuple[int, ...], optional): Shape of a flat, full + structuring element. + footprint (torch.Tensor, optional): Boolean mask selecting participating + positions. + structure (torch.Tensor, optional): Non-flat additive structuring + element. It takes precedence over ``footprint`` and ``size``. + output (torch.Tensor, optional): Preallocated result tensor with the + same shape and device as ``input``. + mode (str): Boundary mode: ``"reflect"``, ``"constant"``, ``"nearest"``, + ``"mirror"``, or ``"wrap"``. + cval (float): Boundary value used when ``mode="constant"``. + origin (int or tuple[int, ...]): Structuring-element anchor offset. + + Returns: + torch.Tensor: Morphological gradient in ``float32``, or ``output``. + + Example: + ```pycon + >>> import torch + >>> import torchmorph as tm + >>> x = torch.arange(25, device="cuda").reshape(1, 1, 5, 5) + >>> tm.morphological_gradient(x, size=3).shape + torch.Size([1, 1, 5, 5]) + ``` + """ validate_bcs_input(input) validate_output(input, output) @@ -231,7 +376,38 @@ def white_tophat( cval: float = 0.0, origin: int | tuple[int, ...] = 0, ) -> Tensor: - """N-dimensional white top-hat filter for ``(B, C, Spatial...)`` CUDA tensors.""" + """Extract small bright features with a white top-hat filter + + Computes ``input - grey_opening(input)`` using ``float32`` arithmetic. + + Args: + input (torch.Tensor): CUDA tensor with shape ``(B, C, Spatial...)``. + size (int or tuple[int, ...], optional): Shape of a flat, full + structuring element. + footprint (torch.Tensor, optional): Boolean mask selecting participating + positions. + structure (torch.Tensor, optional): Non-flat additive structuring + element. It takes precedence over ``footprint`` and ``size``. + output (torch.Tensor, optional): Preallocated result tensor with the + same shape and device as ``input``. + mode (str): Boundary mode: ``"reflect"``, ``"constant"``, ``"nearest"``, + ``"mirror"``, or ``"wrap"``. + cval (float): Boundary value used when ``mode="constant"``. + origin (int or tuple[int, ...]): Structuring-element anchor offset. + + Returns: + torch.Tensor: White top-hat response in ``float32``, or ``output``. + + Example: + ```pycon + >>> import torch + >>> import torchmorph as tm + >>> x = torch.zeros((1, 1, 5, 5), device="cuda") + >>> x[0, 0, 2, 2] = 2 + >>> tm.white_tophat(x, size=3).max().item() + 2.0 + ``` + """ validate_bcs_input(input) validate_output(input, output) @@ -253,7 +429,38 @@ def black_tophat( cval: float = 0.0, origin: int | tuple[int, ...] = 0, ) -> Tensor: - """N-dimensional black top-hat filter for ``(B, C, Spatial...)`` CUDA tensors.""" + """Extract small dark features with a black top-hat filter + + Computes ``grey_closing(input) - input`` using ``float32`` arithmetic. + + Args: + input (torch.Tensor): CUDA tensor with shape ``(B, C, Spatial...)``. + size (int or tuple[int, ...], optional): Shape of a flat, full + structuring element. + footprint (torch.Tensor, optional): Boolean mask selecting participating + positions. + structure (torch.Tensor, optional): Non-flat additive structuring + element. It takes precedence over ``footprint`` and ``size``. + output (torch.Tensor, optional): Preallocated result tensor with the + same shape and device as ``input``. + mode (str): Boundary mode: ``"reflect"``, ``"constant"``, ``"nearest"``, + ``"mirror"``, or ``"wrap"``. + cval (float): Boundary value used when ``mode="constant"``. + origin (int or tuple[int, ...]): Structuring-element anchor offset. + + Returns: + torch.Tensor: Black top-hat response in ``float32``, or ``output``. + + Example: + ```pycon + >>> import torch + >>> import torchmorph as tm + >>> x = torch.ones((1, 1, 5, 5), device="cuda") + >>> x[0, 0, 2, 2] = 0 + >>> tm.black_tophat(x, size=3).max().item() + 1.0 + ``` + """ validate_bcs_input(input) validate_output(input, output) @@ -275,7 +482,38 @@ def morphological_laplace( cval: float = 0.0, origin: int | tuple[int, ...] = 0, ) -> Tensor: - """N-dimensional morphological Laplace for ``(B, C, Spatial...)`` CUDA tensors.""" + """Compute the morphological Laplacian of a grayscale tensor + + Computes ``dilation(input) + erosion(input) - 2 * input`` using ``float32`` + arithmetic. + + Args: + input (torch.Tensor): CUDA tensor with shape ``(B, C, Spatial...)``. + size (int or tuple[int, ...], optional): Shape of a flat, full + structuring element. + footprint (torch.Tensor, optional): Boolean mask selecting participating + positions. + structure (torch.Tensor, optional): Non-flat additive structuring + element. It takes precedence over ``footprint`` and ``size``. + output (torch.Tensor, optional): Preallocated result tensor with the + same shape and device as ``input``. + mode (str): Boundary mode: ``"reflect"``, ``"constant"``, ``"nearest"``, + ``"mirror"``, or ``"wrap"``. + cval (float): Boundary value used when ``mode="constant"``. + origin (int or tuple[int, ...]): Structuring-element anchor offset. + + Returns: + torch.Tensor: Morphological Laplacian in ``float32``, or ``output``. + + Example: + ```pycon + >>> import torch + >>> import torchmorph as tm + >>> x = torch.arange(25, device="cuda").reshape(1, 1, 5, 5) + >>> tm.morphological_laplace(x, size=3).dtype + torch.float32 + ``` + """ validate_bcs_input(input) validate_output(input, output) diff --git a/torchmorph/morphology/structure.py b/torchmorph/morphology/structure.py index 71427fc..2f500e3 100644 --- a/torchmorph/morphology/structure.py +++ b/torchmorph/morphology/structure.py @@ -3,7 +3,37 @@ def generate_binary_structure(rank: int, connectivity: int) -> Tensor: - """N-D generate binary structure""" + """Generate an N-dimensional binary structuring element + + The returned tensor has shape ``(3,) * rank``. Elements whose offset differs + from the center along at most ``connectivity`` axes are ``True``; all other + elements are ``False``. This matches SciPy's + ``ndimage.generate_binary_structure`` connectivity convention. + + Args: + rank (int): Number of spatial dimensions in the structuring element. + Must be at least ``1``. + connectivity (int): Neighborhood connectivity from ``1`` to ``rank``. + ``1`` includes axis-adjacent neighbors; ``rank`` includes the full + ``3 ** rank`` neighborhood. + + Returns: + torch.Tensor: Boolean tensor with shape ``(3,) * rank``. + + Example: + ```pycon + >>> import torch + >>> import torchmorph as tm + >>> tm.generate_binary_structure(2, 1).to(dtype=torch.int32) + tensor([[0, 1, 0], + [1, 1, 1], + [0, 1, 0]], dtype=torch.int32) + >>> tm.generate_binary_structure(2, 2).to(dtype=torch.int32) + tensor([[1, 1, 1], + [1, 1, 1], + [1, 1, 1]], dtype=torch.int32) + ``` + """ if connectivity < 1 or connectivity > rank: raise ValueError(f"connectivity must be in [1, rank], got {connectivity}") @@ -62,7 +92,38 @@ def iterate_structure( iterations: int, origin: int | tuple[int, ...] | None = None, ) -> Tensor | tuple[Tensor, list[int]]: - """Iterate a binary structure by dilating it with itself.""" + """Dilate a binary structuring element with itself repeatedly + + The result is equivalent to applying the original structuring element + ``iterations`` times in a morphology operation. Nonzero input values are + treated as part of the structure. + + Args: + structure (torch.Tensor): N-dimensional binary structuring element. + iterations (int): Number of copies to combine. Values below ``2`` + return a boolean clone of ``structure``. + origin (int or tuple[int, ...], optional): Original anchor offset. A + scalar is applied to every dimension. When supplied, the adjusted + origin is returned with the iterated structure. + + Returns: + torch.Tensor or tuple[torch.Tensor, list[int]]: Boolean iterated + structure. If ``origin`` is supplied, returns ``(structure, origin)`` + with each origin component multiplied by ``iterations``. + + Example: + ```pycon + >>> import torch + >>> import torchmorph as tm + >>> structure = tm.generate_binary_structure(2, 1) + >>> tm.iterate_structure(structure, 2).to(dtype=torch.int32) + tensor([[0, 0, 1, 0, 0], + [0, 1, 1, 1, 0], + [1, 1, 1, 1, 1], + [0, 1, 1, 1, 0], + [0, 0, 1, 0, 0]], dtype=torch.int32) + ``` + """ structure = structure != 0 if iterations < 2: diff --git a/torchmorph/optimal_transport.py b/torchmorph/optimal_transport.py index 3652845..86315af 100644 --- a/torchmorph/optimal_transport.py +++ b/torchmorph/optimal_transport.py @@ -5,11 +5,29 @@ def build_cost_matrix(shape, p=2, device=None) -> Tensor: - """Pairwise L^p distance matrix between the points of a spatial grid. - - Returns a (d, d) cost matrix for the d flattened grid points of `shape`. - Flatten grid-shaped distributions to rows of an (n, d) tensor and pass the - result as `cost_matrix` to :class:`SinkhornSolver`. + """Build pairwise distances between points on a spatial grid + + The grid contains ``d = prod(shape)`` points in row-major order. Flatten + grid-shaped distributions to ``(n, d)`` before passing this matrix to + :class:`SinkhornSolver`. + + Args: + shape (sequence[int]): Spatial grid shape. + p (float): Norm degree passed to :func:`torch.cdist`. + device (torch.device or str, optional): Device for the returned matrix. + + Returns: + torch.Tensor: Symmetric ``float32`` cost matrix with shape ``(d, d)``. + + Example: + ```pycon + >>> import torchmorph as tm + >>> cost = tm.build_cost_matrix((2, 3)) + >>> cost.shape + torch.Size([6, 6]) + >>> cost[0, 1].item() + 1.0 + ``` """ coords = torch.stack( torch.meshgrid([torch.arange(s, device=device) for s in shape], indexing="ij"), dim=-1 @@ -151,19 +169,32 @@ def backward(ctx, grad_output): class SinkhornSolver(nn.Module): - """Entropy-regularized balanced optimal transport as a differentiable module. - - forward(source, target, cost_matrix=None) takes two (n, d) batches of - histograms sharing one (d, d) cost matrix and returns the (n,) transport - distances , with K = exp(-C / epsilon). Gradients w.r.t. source and - target are the centered dual potentials (envelope theorem), so the module - can be used directly as a loss. - - log_space=True runs the iterations entirely in the log domain, which is - numerically stable for small epsilon. The implementation is picked from - the input automatically: CUDA float32 tensors use the fused kernels, - everything else uses pure torch ops (where `threshold` > 0 enables early - stopping; the fused kernels always run `max_iter` iterations). + """Solve entropy-regularized balanced optimal transport + + The module accepts batches of flattened histograms and returns transport + costs that are differentiable with respect to both marginals. CUDA + ``float32`` inputs use fused kernels; other device and dtype combinations + use PyTorch operations. Set ``log_space=True`` for better stability at small + regularization values. + + Args: + epsilon (float): Positive entropy-regularization strength. + max_iter (int): Positive maximum number of Sinkhorn iterations. + threshold (float): Nonnegative early-stopping tolerance checked by the + PyTorch implementation. Fused CUDA kernels always run ``max_iter``. + p (float): Norm used by the default one-dimensional cost matrix. + log_space (bool): Run Sinkhorn iterations in the log domain. + + Example: + ```pycon + >>> import torch + >>> import torchmorph as tm + >>> source = torch.tensor([[1.0, 0.0, 0.0]]) + >>> target = torch.tensor([[0.0, 0.0, 1.0]]) + >>> solver = tm.SinkhornSolver(epsilon=1.0, max_iter=200) + >>> solver(source, target).shape + torch.Size([1]) + ``` """ def __init__(self, epsilon=1.0, max_iter=100, threshold=0.0, p=2, log_space=False): @@ -186,12 +217,42 @@ def extra_repr(self): f"threshold={self.threshold}, log_space={self.log_space}" ) - def data_preprocess(self, source, target, cost_matrix: Optional[Tensor] = None): - """Validate (n, d) inputs, clamp negatives, and normalize each row to unit mass. - - All ops are differentiable, so forward gradients flow through the - normalization. When `cost_matrix` is None the d bins are treated as - points on a line. + def data_preprocess( + self, + source, + target, + cost_matrix: Optional[Tensor] = None, + ) -> tuple[Tensor, Tensor, Tensor]: + """Validate and normalize batched transport marginals + + Negative values are clamped to zero and every row is normalized to unit + mass. These operations remain in the autograd graph. If no cost matrix + is supplied, the ``d`` bins are treated as points on a line. + + Args: + source (torch.Tensor): Floating-point source marginals with shape + ``(n, d)``. + target (torch.Tensor): Floating-point target marginals with the same + shape as ``source``. + cost_matrix (torch.Tensor, optional): Shared pairwise cost matrix + with shape ``(d, d)``. + + Returns: + tuple[torch.Tensor, torch.Tensor, torch.Tensor]: Normalized source, + normalized target, and contiguous cost matrix on the source device + and with the source dtype. + + Example: + ```pycon + >>> import torch + >>> import torchmorph as tm + >>> solver = tm.SinkhornSolver() + >>> a, b, cost = solver.data_preprocess( + ... torch.tensor([[1.0, 1.0]]), torch.tensor([[3.0, 1.0]]) + ... ) + >>> a.sum().item(), b.sum().item(), cost.shape + (1.0, 1.0, torch.Size([2, 2])) + ``` """ if source.ndim != 2 or source.shape != target.shape: raise ValueError("source and target must be (n, d) tensors of the same shape.") @@ -212,8 +273,34 @@ def data_preprocess(self, source, target, cost_matrix: Optional[Tensor] = None): target = target / target.sum(dim=-1, keepdim=True).clamp(min=1e-12) return source, target, cost_matrix - def forward(self, source: Tensor, target: Tensor, cost_matrix: Optional[Tensor] = None): - """(n,) transport distances between the rows of source and target.""" + def forward( + self, + source: Tensor, + target: Tensor, + cost_matrix: Optional[Tensor] = None, + ) -> Tensor: + """Compute transport costs between corresponding marginal rows + + Args: + source (torch.Tensor): Floating-point source marginals with shape + ``(n, d)``. + target (torch.Tensor): Floating-point target marginals with shape + ``(n, d)``. + cost_matrix (torch.Tensor, optional): Shared cost matrix with shape + ``(d, d)``. If ``None``, uses a one-dimensional grid cost. + + Returns: + torch.Tensor: Differentiable transport costs with shape ``(n,)``. + + Example: + ```pycon + >>> import torch + >>> import torchmorph as tm + >>> solver = tm.SinkhornSolver(max_iter=200) + >>> solver(torch.ones(2, 4), torch.ones(2, 4)).shape + torch.Size([2]) + ``` + """ source, target, cost_matrix = self.data_preprocess(source, target, cost_matrix) return _SinkhornDistance.apply( source, @@ -226,17 +313,74 @@ def forward(self, source: Tensor, target: Tensor, cost_matrix: Optional[Tensor] ) @torch.no_grad() - def plan(self, source: Tensor, target: Tensor, cost_matrix: Optional[Tensor] = None): - """(n, d, d) transport plans, reconstructed in log space.""" + def plan( + self, + source: Tensor, + target: Tensor, + cost_matrix: Optional[Tensor] = None, + ) -> Tensor: + """Reconstruct transport plans in the log domain + + This method runs without gradient tracking. + + Args: + source (torch.Tensor): Floating-point source marginals with shape + ``(n, d)``. + target (torch.Tensor): Floating-point target marginals with shape + ``(n, d)``. + cost_matrix (torch.Tensor, optional): Shared cost matrix with shape + ``(d, d)``. + + Returns: + torch.Tensor: Transport plans with shape ``(n, d, d)``. + + Example: + ```pycon + >>> import torch + >>> import torchmorph as tm + >>> solver = tm.SinkhornSolver(max_iter=200) + >>> solver.plan(torch.ones(1, 3), torch.ones(1, 3)).shape + torch.Size([1, 3, 3]) + ``` + """ source, target, cost_matrix = self.data_preprocess(source, target, cost_matrix) log_u, log_v = self._iterate(source, target, cost_matrix) return torch.exp(log_u.unsqueeze(-1) - cost_matrix / self.epsilon + log_v.unsqueeze(-2)) @torch.no_grad() - def potentials(self, source: Tensor, target: Tensor, cost_matrix: Optional[Tensor] = None): - """Centered dual potentials (f, g), each (n, d). - - These are the gradients of the entropic OT cost w.r.t. the marginals. + def potentials( + self, + source: Tensor, + target: Tensor, + cost_matrix: Optional[Tensor] = None, + ) -> tuple[Tensor, Tensor]: + """Compute centered dual potentials for both marginals + + The potentials are the envelope-theorem gradients of the entropic + transport cost with respect to normalized marginals. This method runs + without gradient tracking. + + Args: + source (torch.Tensor): Floating-point source marginals with shape + ``(n, d)``. + target (torch.Tensor): Floating-point target marginals with shape + ``(n, d)``. + cost_matrix (torch.Tensor, optional): Shared cost matrix with shape + ``(d, d)``. + + Returns: + tuple[torch.Tensor, torch.Tensor]: Source and target potentials, + each with shape ``(n, d)`` and zero mean along the last axis. + + Example: + ```pycon + >>> import torch + >>> import torchmorph as tm + >>> solver = tm.SinkhornSolver(max_iter=200) + >>> f, g = solver.potentials(torch.ones(1, 3), torch.ones(1, 3)) + >>> f.shape, g.shape + (torch.Size([1, 3]), torch.Size([1, 3])) + ``` """ source, target, cost_matrix = self.data_preprocess(source, target, cost_matrix) log_u, log_v = self._iterate(source, target, cost_matrix)