Setup documentation

.gitignore

@@ -55,6 +55,6 @@ Thumbs.db
 .coverage
 htmlcov/
 
-# Local development documentation
+# Local documentation artifacts
 /CONTEXT.md
-/docs/
+/site/

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.

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.

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

docs/distance_transforms.md

@@ -0,0 +1,7 @@
+# Distance Transforms
+
+::: torchmorph.euclidean_distance_transform
+
+::: torchmorph.chamfer_distance_transform
+
+::: torchmorph.brute_force_distance_transform

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

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

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

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

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;
+  }
+}

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

mkdocs_hooks.py

@@ -0,0 +1,16 @@
+"""Small build-time enhancements for the generated API reference."""
+
+import re
+
+_API_HEADING = re.compile(r'(<h[2-6]\s+id="torchmorph\.([^"]+)"[^>]*>)')
+
+
+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'<a id="{short_name}" class="doc-short-anchor" aria-hidden="true"></a>'
+        return f"{anchor}\n{match.group(1)}"
+
+    return _API_HEADING.sub(add_short_anchor, html)

requirements-docs.txt

@@ -0,0 +1,5 @@
+-r requirements-test.txt
+mkdocs-material
+mkdocstrings[python]
+pymdown-extensions
+

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",