serron 0.0.1__tar.gz

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
@@ -0,0 +1,97 @@
1
+ cmake_minimum_required(VERSION 3.24)
2
+
3
+ if(NOT DEFINED CMAKE_CUDA_COMPILER)
4
+ if(DEFINED ENV{CUDA_HOME})
5
+ set(CMAKE_CUDA_COMPILER "$ENV{CUDA_HOME}/bin/nvcc" CACHE FILEPATH "" FORCE)
6
+ else()
7
+ find_program(_NVCC nvcc)
8
+ if(_NVCC)
9
+ set(CMAKE_CUDA_COMPILER "${_NVCC}" CACHE FILEPATH "" FORCE)
10
+ endif()
11
+ endif()
12
+ endif()
13
+
14
+
15
+ project(morphottention LANGUAGES CXX CUDA)
16
+
17
+ # 23 host / 20 device
18
+ set(CMAKE_CXX_STANDARD 23)
19
+ set(CMAKE_CXX_STANDARD_REQUIRED ON)
20
+ set(CMAKE_CUDA_STANDARD 20)
21
+ set(CMAKE_CUDA_STANDARD_REQUIRED ON)
22
+
23
+ set(CMAKE_POSITION_INDEPENDENT_CODE ON)
24
+ set(CMAKE_CUDA_SEPARABLE_COMPILATION OFF)
25
+
26
+ set(CMAKE_CUDA_ARCHITECTURES 90 100 120)
27
+ set(TORCH_CUDA_ARCH_LIST "9.0;10.0;12.0")
28
+
29
+ set(CMAKE_EXPORT_COMPILE_COMMANDS ON)
30
+
31
+ set(Python_FIND_VIRTUALENV FIRST)
32
+
33
+ find_package(Python REQUIRED COMPONENTS Interpreter Development.Module)
34
+
35
+ execute_process(
36
+ COMMAND ${Python_EXECUTABLE} -c "import torch; print(torch.utils.cmake_prefix_path)"
37
+ OUTPUT_VARIABLE TORCH_CMAKE_PREFIX
38
+ OUTPUT_STRIP_TRAILING_WHITESPACE
39
+ )
40
+ list(APPEND CMAKE_PREFIX_PATH "${TORCH_CMAKE_PREFIX}")
41
+ find_package(Torch REQUIRED CONFIG)
42
+
43
+
44
+ execute_process(
45
+ COMMAND ${Python_EXECUTABLE} -c "import pathlib, torch; print(pathlib.Path(torch.__file__).resolve().parent / 'lib')"
46
+ OUTPUT_VARIABLE TORCH_LIB_DIR
47
+ OUTPUT_STRIP_TRAILING_WHITESPACE
48
+ )
49
+ if(NOT EXISTS "${TORCH_LIB_DIR}")
50
+ message(FATAL_ERROR "Could not locate torch lib dir at: ${TORCH_LIB_DIR}")
51
+ endif()
52
+
53
+ find_package(CUDAToolkit REQUIRED)
54
+
55
+ Python_add_library(_C MODULE WITH_SOABI
56
+ csrc/bindings.cpp
57
+ csrc/cuda/morphology.cu
58
+ )
59
+
60
+ # -lineinfo on Debug; full device debug (-G) only when explicitly sanitizing.
61
+ option(CUDA_SANITIZE "Build device code with -G for compute-sanitizer" OFF)
62
+ target_compile_options(_C PRIVATE
63
+ $<$<COMPILE_LANG_AND_ID:CUDA,NVIDIA>:--threads=0>
64
+ $<$<AND:$<COMPILE_LANG_AND_ID:CUDA,NVIDIA>,$<CONFIG:Debug>>:-lineinfo>
65
+ $<$<AND:$<COMPILE_LANG_AND_ID:CUDA,NVIDIA>,$<CONFIG:Debug>,$<BOOL:${CUDA_SANITIZE}>>:-G>
66
+ )
67
+
68
+
69
+ if(MSVC)
70
+ target_compile_options(_C PRIVATE
71
+ $<$<COMPILE_LANGUAGE:CXX>:/Zc:preprocessor>
72
+ $<$<COMPILE_LANGUAGE:CUDA>:-Xcompiler=/Zc:preprocessor>
73
+ )
74
+ endif()
75
+
76
+ target_include_directories(_C PRIVATE
77
+ ${CMAKE_CURRENT_SOURCE_DIR}/csrc
78
+ ${CMAKE_CURRENT_SOURCE_DIR}/csrc/compat
79
+ ${TORCH_INCLUDE_DIRS}
80
+ ${Python_INCLUDE_DIRS}
81
+ ${CUDAToolkit_INCLUDE_DIRS}
82
+ )
83
+
84
+ target_link_libraries(_C PRIVATE
85
+ ${TORCH_LIBRARIES}
86
+ CUDA::cudart
87
+ CUDA::cublas
88
+ )
89
+
90
+ if(NOT WIN32)
91
+ set_target_properties(_C PROPERTIES
92
+ BUILD_RPATH "${TORCH_LIB_DIR}"
93
+ INSTALL_RPATH "$ORIGIN/../torch/lib"
94
+ )
95
+ endif()
96
+
97
+ install(TARGETS _C LIBRARY DESTINATION serron)
serron-0.0.1/PKG-INFO ADDED
@@ -0,0 +1,147 @@
1
+ Metadata-Version: 2.2
2
+ Name: serron
3
+ Version: 0.0.1
4
+ Summary: Mathematical Morphology module for PyTorch (CUDA), providing differentiable operators and learnable network layers.
5
+ Keywords: cuda,pytorch,morphology
6
+ Author-Email: Vedran Hrabar <vedran.hrabar@outlook.com>
7
+ License: MIT
8
+ Classifier: Development Status :: 1 - Planning
9
+ Classifier: Environment :: GPU
10
+ Classifier: Environment :: GPU :: NVIDIA CUDA
11
+ Classifier: Environment :: GPU :: NVIDIA CUDA :: 13
12
+ Classifier: Intended Audience :: Developers
13
+ Classifier: Intended Audience :: Science/Research
14
+ Classifier: License :: OSI Approved :: MIT License
15
+ Classifier: Operating System :: POSIX :: Linux
16
+ Classifier: Operating System :: Microsoft :: Windows
17
+ Classifier: Programming Language :: C++
18
+ Classifier: Programming Language :: Python :: 3
19
+ Classifier: Programming Language :: Python :: 3 :: Only
20
+ Classifier: Programming Language :: Python :: 3.12
21
+ Classifier: Programming Language :: Python :: 3.13
22
+ Classifier: Programming Language :: Python :: 3.14
23
+ Classifier: Programming Language :: Python :: Implementation :: CPython
24
+ Classifier: Topic :: Scientific/Engineering
25
+ Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
26
+ Classifier: Topic :: Scientific/Engineering :: Image Recognition
27
+ Classifier: Topic :: Scientific/Engineering :: Mathematics
28
+ Classifier: Topic :: Software Development :: Libraries
29
+ Classifier: Topic :: Software Development :: Libraries :: Python Modules
30
+ Classifier: Typing :: Typed
31
+ Project-URL: repository, https://github.com/vhrabar/serron
32
+ Project-URL: documentation, https://github.com/vhrabar/serron/wiki
33
+ Project-URL: Bug Tracker, https://github.com/vhrabar/serron/issues
34
+ Requires-Python: >=3.12
35
+ Requires-Dist: torch>=2.12
36
+ Description-Content-Type: text/markdown
37
+
38
+ # Serron
39
+
40
+ Mathematical Morphology module for PyTorch (CUDA), providing differentiable operators and learnable network layers.
41
+ ## Install
42
+
43
+ Prebuilt wheels are published for CPython 3.14 on Linux (x86_64, aarch64) and
44
+ Windows (x86_64). A working CUDA-enabled PyTorch (`torch >= 2.12`) must already
45
+ be installed in the environment.
46
+
47
+ ```bash
48
+ pip install serron
49
+ ```
50
+
51
+ ## Usage
52
+
53
+ > **Alpha:** the CUDA kernels are not wired up yet, so the operators below currently
54
+ > raise `NotImplementedError`.
55
+
56
+ Serron works on standard PyTorch tensors laid out as `(N, C, H, W)` and living on a CUDA
57
+ device.
58
+
59
+ ### Functional operators
60
+
61
+ Stateless operators live in `serron.functional` and are re-exported at the top level.
62
+ Each takes an input tensor, a structuring element, and an optional `border` mode:
63
+
64
+ ```python
65
+ import torch
66
+ import serron
67
+ from serron import BorderMode
68
+ from serron import structuring_element as se
69
+
70
+ x = torch.rand(1, 1, 256, 256, device="cuda")
71
+ kernel = se.disk(3, device="cuda")
72
+
73
+ eroded = serron.erosion(x, kernel)
74
+ dilated = serron.dilatation(x, kernel)
75
+ opened = serron.opening(x, kernel)
76
+ closed = serron.closing(x, kernel)
77
+
78
+ # Derived operators
79
+ grad = serron.gradient(x, kernel) # dilate(x) - erode(x)
80
+ white = serron.top_hat(x, kernel) # x - open(x)
81
+ black = serron.black_hat(x, kernel) # close(x) - x
82
+
83
+ # Control border handling
84
+ eroded_reflect = serron.erosion(x, kernel, border=BorderMode.REFLECT)
85
+ ```
86
+
87
+ Available operators: `erosion`, `dilatation`, `opening`, `closing`, `gradient`,
88
+ `top_hat`, `black_hat`.
89
+
90
+ ### Structuring elements
91
+
92
+ `serron.structuring_element` builds common SE shapes on the requested device:
93
+
94
+ ```python
95
+ from serron import structuring_element as se
96
+
97
+ se.square(5, device="cuda") # (5, 5) full square
98
+ se.cross(5, device="cuda") # (5, 5) plus shape
99
+ se.disk(3, device="cuda") # (7, 7) disk, radius 3
100
+ se.diamond(3, device="cuda") # (7, 7) diamond, radius 3
101
+ se.from_tensor(my_weights) # wrap an arbitrary 2-D tensor as a grayscale SE
102
+ ```
103
+
104
+ ### Border modes
105
+
106
+ `BorderMode` controls how out-of-bounds neighbors are handled:
107
+
108
+ | Mode | Behavior |
109
+ | --- | --- |
110
+ | `BorderMode.REPLICATE` | Repeat the edge value (default) |
111
+ | `BorderMode.REFLECT` | Mirror across the edge |
112
+ | `BorderMode.CONSTANT` | Pad with a constant |
113
+
114
+ ### Learnable layers
115
+
116
+ `serron` also exposes `torch.nn.Module` layers with a learnable structuring element,
117
+ so morphology can be trained end-to-end inside a network. Each layer takes the number
118
+ of `channels` and a `kernel_size`:
119
+
120
+ ```python
121
+ import torch
122
+ from serron import Erosion2d, Dilation2d, Opening2d, Closing2d
123
+ from serron import BorderMode
124
+
125
+ layer = Erosion2d(channels=3, kernel_size=5, border=BorderMode.REPLICATE).cuda()
126
+
127
+ x = torch.rand(8, 3, 64, 64, device="cuda")
128
+ y = layer(x) # forward pass; layer.weight is a trainable (C, k, k) SE
129
+ y.sum().backward() # gradients flow into layer.weight
130
+ ```
131
+
132
+ Available layers: `Erosion2d`, `Dilation2d`, `Opening2d`, `Closing2d`.
133
+
134
+ ## Building from source
135
+
136
+ Requires the CUDA 13.X toolkit (`nvcc`) and a matching `torch` build:
137
+
138
+ ```bash
139
+ uv sync --package serron --no-dev --group build
140
+ uv build --package serron --wheel --no-build-isolation
141
+ ```
142
+
143
+ ## License
144
+
145
+ MIT
146
+
147
+ Copyright © 2026 Vedran Hrabar.
serron-0.0.1/README.md ADDED
@@ -0,0 +1,110 @@
1
+ # Serron
2
+
3
+ Mathematical Morphology module for PyTorch (CUDA), providing differentiable operators and learnable network layers.
4
+ ## Install
5
+
6
+ Prebuilt wheels are published for CPython 3.14 on Linux (x86_64, aarch64) and
7
+ Windows (x86_64). A working CUDA-enabled PyTorch (`torch >= 2.12`) must already
8
+ be installed in the environment.
9
+
10
+ ```bash
11
+ pip install serron
12
+ ```
13
+
14
+ ## Usage
15
+
16
+ > **Alpha:** the CUDA kernels are not wired up yet, so the operators below currently
17
+ > raise `NotImplementedError`.
18
+
19
+ Serron works on standard PyTorch tensors laid out as `(N, C, H, W)` and living on a CUDA
20
+ device.
21
+
22
+ ### Functional operators
23
+
24
+ Stateless operators live in `serron.functional` and are re-exported at the top level.
25
+ Each takes an input tensor, a structuring element, and an optional `border` mode:
26
+
27
+ ```python
28
+ import torch
29
+ import serron
30
+ from serron import BorderMode
31
+ from serron import structuring_element as se
32
+
33
+ x = torch.rand(1, 1, 256, 256, device="cuda")
34
+ kernel = se.disk(3, device="cuda")
35
+
36
+ eroded = serron.erosion(x, kernel)
37
+ dilated = serron.dilatation(x, kernel)
38
+ opened = serron.opening(x, kernel)
39
+ closed = serron.closing(x, kernel)
40
+
41
+ # Derived operators
42
+ grad = serron.gradient(x, kernel) # dilate(x) - erode(x)
43
+ white = serron.top_hat(x, kernel) # x - open(x)
44
+ black = serron.black_hat(x, kernel) # close(x) - x
45
+
46
+ # Control border handling
47
+ eroded_reflect = serron.erosion(x, kernel, border=BorderMode.REFLECT)
48
+ ```
49
+
50
+ Available operators: `erosion`, `dilatation`, `opening`, `closing`, `gradient`,
51
+ `top_hat`, `black_hat`.
52
+
53
+ ### Structuring elements
54
+
55
+ `serron.structuring_element` builds common SE shapes on the requested device:
56
+
57
+ ```python
58
+ from serron import structuring_element as se
59
+
60
+ se.square(5, device="cuda") # (5, 5) full square
61
+ se.cross(5, device="cuda") # (5, 5) plus shape
62
+ se.disk(3, device="cuda") # (7, 7) disk, radius 3
63
+ se.diamond(3, device="cuda") # (7, 7) diamond, radius 3
64
+ se.from_tensor(my_weights) # wrap an arbitrary 2-D tensor as a grayscale SE
65
+ ```
66
+
67
+ ### Border modes
68
+
69
+ `BorderMode` controls how out-of-bounds neighbors are handled:
70
+
71
+ | Mode | Behavior |
72
+ | --- | --- |
73
+ | `BorderMode.REPLICATE` | Repeat the edge value (default) |
74
+ | `BorderMode.REFLECT` | Mirror across the edge |
75
+ | `BorderMode.CONSTANT` | Pad with a constant |
76
+
77
+ ### Learnable layers
78
+
79
+ `serron` also exposes `torch.nn.Module` layers with a learnable structuring element,
80
+ so morphology can be trained end-to-end inside a network. Each layer takes the number
81
+ of `channels` and a `kernel_size`:
82
+
83
+ ```python
84
+ import torch
85
+ from serron import Erosion2d, Dilation2d, Opening2d, Closing2d
86
+ from serron import BorderMode
87
+
88
+ layer = Erosion2d(channels=3, kernel_size=5, border=BorderMode.REPLICATE).cuda()
89
+
90
+ x = torch.rand(8, 3, 64, 64, device="cuda")
91
+ y = layer(x) # forward pass; layer.weight is a trainable (C, k, k) SE
92
+ y.sum().backward() # gradients flow into layer.weight
93
+ ```
94
+
95
+ Available layers: `Erosion2d`, `Dilation2d`, `Opening2d`, `Closing2d`.
96
+
97
+ ## Building from source
98
+
99
+ Requires the CUDA 13.X toolkit (`nvcc`) and a matching `torch` build:
100
+
101
+ ```bash
102
+ uv sync --package serron --no-dev --group build
103
+ uv build --package serron --wheel --no-build-isolation
104
+ ```
105
+
106
+ ## License
107
+
108
+ MIT
109
+
110
+ Copyright © 2026 Vedran Hrabar.
@@ -0,0 +1,18 @@
1
+ #include "ops.h"
2
+ #include "registration.h"
3
+
4
+ #include <torch/library.h>
5
+
6
+ // Operator schemas, registered under the serron:: namespace.
7
+ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
8
+ ops.def("erode(Tensor input, Tensor kernel, int border) -> Tensor");
9
+ ops.def("dilate(Tensor input, Tensor kernel, int border) -> Tensor");
10
+ }
11
+
12
+ // CUDA implementations.
13
+ TORCH_LIBRARY_IMPL_EXPAND(TORCH_EXTENSION_NAME, CUDA, ops) {
14
+ ops.impl("erode", &serron::erode);
15
+ ops.impl("dilate", &serron::dilate);
16
+ }
17
+
18
+ REGISTER_EXTENSION(TORCH_EXTENSION_NAME)
@@ -0,0 +1,16 @@
1
+ #ifndef SERRON_COMPAT_REGISTRATION_H
2
+ #define SERRON_COMPAT_REGISTRATION_H
3
+
4
+ // Compatibility shim for the local CMake / PyPI build.
5
+
6
+ #include <torch/library.h>
7
+
8
+ #define TORCH_LIBRARY_EXPAND(NAME, MODULE) TORCH_LIBRARY(NAME, MODULE)
9
+
10
+ #define TORCH_LIBRARY_IMPL_EXPAND(NAME, DEVICE, MODULE) TORCH_LIBRARY_IMPL(NAME, DEVICE, MODULE)
11
+
12
+ #define TORCH_EXTENSION_NAME serron
13
+
14
+ #define REGISTER_EXTENSION(NAME)
15
+
16
+ #endif // SERRON_COMPAT_REGISTRATION_H
@@ -0,0 +1,17 @@
1
+ #include "ops.h"
2
+
3
+ #include <c10/util/Exception.h>
4
+
5
+
6
+ namespace serron {
7
+
8
+ at::Tensor erode(const at::Tensor& input, const at::Tensor& kernel, int64_t border) {
9
+ TORCH_CHECK_NOT_IMPLEMENTED(false, "serron::erode is not implemented yet");
10
+ }
11
+
12
+ at::Tensor dilate(const at::Tensor& input, const at::Tensor& kernel, int64_t border) {
13
+
14
+ TORCH_CHECK_NOT_IMPLEMENTED(false, "serron::dilate is not implemented yet");
15
+ }
16
+
17
+ } // namespace serron
@@ -0,0 +1,13 @@
1
+ #ifndef SERRON_OPS_H
2
+ #define SERRON_OPS_H
3
+
4
+ #include <torch/torch.h>
5
+
6
+ namespace serron {
7
+
8
+ at::Tensor erode(const at::Tensor& input, const at::Tensor& kernel, int64_t border);
9
+ at::Tensor dilate(const at::Tensor& input, const at::Tensor& kernel, int64_t border);
10
+
11
+ } // namespace serron
12
+
13
+ #endif // SERRON_OPS_H
@@ -0,0 +1,82 @@
1
+ [project]
2
+ name = "serron"
3
+ version = "0.0.1"
4
+ description = " Mathematical Morphology module for PyTorch (CUDA), providing differentiable operators and learnable network layers."
5
+ readme = "README.md"
6
+ requires-python = ">=3.12"
7
+ authors = [
8
+ { name = "Vedran Hrabar", email = "vedran.hrabar@outlook.com" },
9
+ ]
10
+ license = { text = "MIT" }
11
+ keywords = ["cuda", "pytorch", "morphology",]
12
+
13
+ classifiers = [
14
+ "Development Status :: 1 - Planning",
15
+
16
+ "Environment :: GPU",
17
+ "Environment :: GPU :: NVIDIA CUDA",
18
+ "Environment :: GPU :: NVIDIA CUDA :: 13",
19
+
20
+ "Intended Audience :: Developers",
21
+ "Intended Audience :: Science/Research",
22
+
23
+ "License :: OSI Approved :: MIT License",
24
+
25
+ "Operating System :: POSIX :: Linux",
26
+ "Operating System :: Microsoft :: Windows",
27
+
28
+ "Programming Language :: C++",
29
+ "Programming Language :: Python :: 3",
30
+ "Programming Language :: Python :: 3 :: Only",
31
+ "Programming Language :: Python :: 3.12",
32
+ "Programming Language :: Python :: 3.13",
33
+ "Programming Language :: Python :: 3.14",
34
+ "Programming Language :: Python :: Implementation :: CPython",
35
+
36
+ "Topic :: Scientific/Engineering",
37
+ "Topic :: Scientific/Engineering :: Artificial Intelligence",
38
+ "Topic :: Scientific/Engineering :: Image Recognition",
39
+ "Topic :: Scientific/Engineering :: Mathematics",
40
+ "Topic :: Software Development :: Libraries",
41
+ "Topic :: Software Development :: Libraries :: Python Modules",
42
+
43
+ "Typing :: Typed",
44
+ ]
45
+
46
+ dependencies = [
47
+ "torch>=2.12",
48
+ ]
49
+
50
+ [project.urls]
51
+ repository = "https://github.com/vhrabar/serron"
52
+ documentation = "https://github.com/vhrabar/serron/wiki"
53
+ "Bug Tracker" = "https://github.com/vhrabar/serron/issues"
54
+
55
+
56
+ [[tool.uv.index]]
57
+ name = "pytorch"
58
+ url = "https://download.pytorch.org/whl/cu132"
59
+ explicit = true
60
+
61
+ [tool.uv.sources]
62
+ torch = { index = "pytorch" }
63
+ torchvision = { index = "pytorch" }
64
+
65
+
66
+ [build-system]
67
+ requires = ["scikit-build-core==1.0.2"]
68
+ build-backend = "scikit_build_core.build"
69
+
70
+
71
+ [dependency-groups]
72
+ build = ["scikit-build-core==1.0.2", "cmake", "ninja"]
73
+
74
+
75
+
76
+ [tool.scikit-build]
77
+ minimum-version = "build-system.requires"
78
+ cmake.version = ">=3.24"
79
+ build-dir = "build/{wheel_tag}"
80
+
81
+ wheel.packages = ["torch-ext/serron"]
82
+ editable.rebuild = true
@@ -0,0 +1,36 @@
1
+ """
2
+ Serron: CUDA-accelerated mathematical morphology for PyTorch.
3
+ """
4
+
5
+ from __future__ import annotations
6
+
7
+ from . import functional, structuring_element
8
+ from .enums import BorderMode, Operation
9
+ from .functional import (
10
+ black_hat,
11
+ closing,
12
+ dilatation,
13
+ erosion,
14
+ gradient,
15
+ opening,
16
+ top_hat,
17
+ )
18
+ from .modules import Closing2d, Dilation2d, Erosion2d, Opening2d
19
+
20
+ __all__ = [
21
+ "BorderMode",
22
+ "Closing2d",
23
+ "Dilation2d",
24
+ "Erosion2d",
25
+ "Opening2d",
26
+ "Operation",
27
+ "black_hat",
28
+ "closing",
29
+ "dilatation",
30
+ "erosion",
31
+ "functional",
32
+ "gradient",
33
+ "opening",
34
+ "structuring_element",
35
+ "top_hat",
36
+ ]
@@ -0,0 +1,33 @@
1
+ """
2
+ Op-namespace loader for the local CMake / PyPI build.
3
+ """
4
+
5
+ from __future__ import annotations
6
+
7
+ import glob
8
+ import os
9
+
10
+ import torch
11
+
12
+ _NAMESPACE = "serron"
13
+
14
+
15
+ def _load_c_extension() -> None:
16
+ pkg_dir = os.path.dirname(__file__)
17
+ for pattern in ("_C*.so", "_C*.pyd", "_C*.dll"):
18
+ matches = glob.glob(os.path.join(pkg_dir, pattern))
19
+ if matches:
20
+ torch.ops.load_library(matches[0]) # type: ignore[no-untyped-call]
21
+ return
22
+ raise ImportError(
23
+ f"Could not find the compiled Serron '_C' extension in {pkg_dir}. Reinstall the package so the CUDA kernels are built."
24
+ )
25
+
26
+
27
+ _load_c_extension()
28
+
29
+ ops = getattr(torch.ops, _NAMESPACE)
30
+
31
+
32
+ def add_op_namespace_prefix(op_name: str) -> str:
33
+ return f"{_NAMESPACE}::{op_name}"
@@ -0,0 +1,12 @@
1
+ """
2
+ Autograd wrapper and nn.Module around the compiled serron CUDA kernels.
3
+ """
4
+
5
+ from __future__ import annotations
6
+
7
+ import contextlib
8
+
9
+ with contextlib.suppress(ImportError):
10
+ from ._ops import add_op_namespace_prefix, ops # type: ignore[import-not-found]
11
+
12
+ __all__ = ["add_op_namespace_prefix", "ops"]
@@ -0,0 +1,31 @@
1
+ """
2
+ Enumerations shared across the morphology ops and layers.
3
+ """
4
+
5
+ from __future__ import annotations
6
+
7
+ from enum import Enum
8
+
9
+
10
+ class Operation(Enum):
11
+ """
12
+ Morphological operations
13
+ """
14
+
15
+ ERODE = "erode"
16
+ DILATE = "dilate"
17
+ OPEN = "open"
18
+ CLOSE = "close"
19
+ GRADIENT = "gradient"
20
+ TOP_HAT = "top_hat"
21
+ BLACK_HAT = "black_hat"
22
+
23
+
24
+ class BorderMode(Enum):
25
+ """
26
+ OuB neighborhood border mode.
27
+ """
28
+
29
+ REFLECT = "reflect"
30
+ REPLICATE = "replicate"
31
+ CONSTANT = "constant"
@@ -0,0 +1,86 @@
1
+ """
2
+ Stateless morphological operations backed by the CUDA kernels.
3
+ """
4
+
5
+ from __future__ import annotations
6
+
7
+ import torch
8
+
9
+ from .enums import BorderMode
10
+
11
+
12
+ def erosion(input_: torch.Tensor, kernel: torch.Tensor, *, border: BorderMode = BorderMode.REPLICATE) -> torch.Tensor:
13
+ """
14
+ Grayscale erosion (sliding minimum) of ``input`` by ``kernel``
15
+ :param input_: input tensor
16
+ :param kernel: kernel tensor
17
+ :param border: border mode
18
+ :return: top-hat tensor
19
+ """
20
+ raise NotImplementedError
21
+
22
+
23
+ def dilatation(input_: torch.Tensor, kernel: torch.Tensor, *, border: BorderMode = BorderMode.REPLICATE) -> torch.Tensor:
24
+ """
25
+ Grayscale dilation (sliding maximum) of ``input`` by ``kernel``
26
+ :param input_: input tensor
27
+ :param kernel: kernel tensor
28
+ :param border: border mode
29
+ :return: top-hat tensor
30
+ """
31
+ raise NotImplementedError
32
+
33
+
34
+ def opening(input_: torch.Tensor, kernel: torch.Tensor, *, border: BorderMode = BorderMode.REPLICATE) -> torch.Tensor:
35
+ """
36
+ Opening: erosion followed by dilation.
37
+ :param input_: input tensor
38
+ :param kernel: kernel tensor
39
+ :param border: border mode
40
+ :return: top-hat tensor
41
+ """
42
+ return dilatation(erosion(input_, kernel, border=border), kernel, border=border)
43
+
44
+
45
+ def closing(input_: torch.Tensor, kernel: torch.Tensor, *, border: BorderMode = BorderMode.REPLICATE) -> torch.Tensor:
46
+ """
47
+ "Closing: dilation followed by erosion.
48
+ :param input_: input tensor
49
+ :param kernel: kernel tensor
50
+ :param border: border mode
51
+ :return: top-hat tensor
52
+ """
53
+ return erosion(dilatation(input_, kernel, border=border), kernel, border=border)
54
+
55
+
56
+ def gradient(input_: torch.Tensor, kernel: torch.Tensor, *, border: BorderMode = BorderMode.REPLICATE) -> torch.Tensor:
57
+ """
58
+ Morphological gradient: ``dilate(input) - erode(input)``.
59
+ :param input_: input tensor
60
+ :param kernel: kernel tensor
61
+ :param border: border mode
62
+ :return: top-hat tensor
63
+ """
64
+ return dilatation(input_, kernel, border=border) - erosion(input_, kernel, border=border)
65
+
66
+
67
+ def top_hat(input_: torch.Tensor, kernel: torch.Tensor, *, border: BorderMode = BorderMode.REPLICATE) -> torch.Tensor:
68
+ """
69
+ White top-hat: ``input - open(input)``
70
+ :param input_: input tensor
71
+ :param kernel: kernel tensor
72
+ :param border: border mode
73
+ :return: top-hat tensor
74
+ """
75
+ return input_ - opening(input_, kernel, border=border)
76
+
77
+
78
+ def black_hat(input_: torch.Tensor, kernel: torch.Tensor, *, border: BorderMode = BorderMode.REPLICATE) -> torch.Tensor:
79
+ """
80
+ Black top-hat: ``close(input) - input``
81
+ :param input_: input tensor
82
+ :param kernel: kernel tensor
83
+ :param border: border mode
84
+ :return: top-hat tensor
85
+ """
86
+ return closing(input_, kernel, border=border) - input_
@@ -0,0 +1,66 @@
1
+ """
2
+ Learnable morphology layers (``torch.nn.Module``).
3
+ """
4
+
5
+ from __future__ import annotations
6
+
7
+ import torch
8
+ from torch import nn
9
+
10
+ from . import functional as F
11
+ from .enums import BorderMode
12
+
13
+
14
+ class _MorphologyNd(nn.Module):
15
+ """Shared base: owns the learnable SE parameter."""
16
+
17
+ def __init__(
18
+ self,
19
+ channels: int,
20
+ kernel_size: int,
21
+ *,
22
+ border: BorderMode = BorderMode.REPLICATE,
23
+ dtype: torch.dtype = torch.float32,
24
+ device: torch.device | str | None = None,
25
+ ) -> None:
26
+ super().__init__()
27
+ self.channels = channels
28
+ self.kernel_size = kernel_size
29
+ self.border = border
30
+ self.weight = nn.Parameter(torch.empty(channels, kernel_size, kernel_size, dtype=dtype, device=device))
31
+ self.reset_parameters()
32
+
33
+ def reset_parameters(self) -> None:
34
+ with torch.no_grad():
35
+ self.weight.zero_()
36
+
37
+ def forward(self, x: torch.Tensor) -> torch.Tensor:
38
+ raise NotImplementedError
39
+
40
+
41
+ class Erosion2d(_MorphologyNd):
42
+ """Learnable 2-D erosion layer."""
43
+
44
+ def forward(self, x: torch.Tensor) -> torch.Tensor:
45
+ return F.erosion(x, self.weight, border=self.border)
46
+
47
+
48
+ class Dilation2d(_MorphologyNd):
49
+ """Learnable 2-D dilation layer."""
50
+
51
+ def forward(self, x: torch.Tensor) -> torch.Tensor:
52
+ return F.dilatation(x, self.weight, border=self.border)
53
+
54
+
55
+ class Opening2d(_MorphologyNd):
56
+ """Learnable 2-D opening layer."""
57
+
58
+ def forward(self, x: torch.Tensor) -> torch.Tensor:
59
+ return F.opening(x, self.weight, border=self.border)
60
+
61
+
62
+ class Closing2d(_MorphologyNd):
63
+ """Learnable 2-D closing layer."""
64
+
65
+ def forward(self, x: torch.Tensor) -> torch.Tensor:
66
+ return F.closing(x, self.weight, border=self.border)
File without changes
@@ -0,0 +1,55 @@
1
+ """
2
+ Structuring-element builders.
3
+ """
4
+
5
+ from __future__ import annotations
6
+
7
+ import torch
8
+
9
+
10
+ def square(size: int, *, device: torch.device | str | None = None) -> torch.Tensor:
11
+ """
12
+ Square (full) SE of shape ``(size, size)``.
13
+ :param size: SE size
14
+ :param device: device to create the SE on
15
+ :return: SE tensor
16
+ """
17
+ raise NotImplementedError
18
+
19
+
20
+ def cross(size: int, *, device: torch.device | str | None = None) -> torch.Tensor:
21
+ """
22
+ Cross/plus-shaped SE of shape ``(size, size)``.
23
+ :param size: SE size
24
+ :param device: device to create the SE on
25
+ :return: SE tensor"""
26
+ raise NotImplementedError
27
+
28
+
29
+ def disk(radius: int, *, device: torch.device | str | None = None) -> torch.Tensor:
30
+ """
31
+ Disk SE of shape ``(2*radius+1, 2*radius+1)``.
32
+ :param radius: SE radius
33
+ :param device: device to create the SE on
34
+ :return: SE tensor
35
+ """
36
+ raise NotImplementedError
37
+
38
+
39
+ def diamond(radius: int, *, device: torch.device | str | None = None) -> torch.Tensor:
40
+ """
41
+ Diamond SE of shape ``(2*radius+1, 2*radius+1)``.
42
+ :param radius: SE radius
43
+ :param device: device to create the SE on
44
+ :return: SE tensor
45
+ """
46
+ raise NotImplementedError
47
+
48
+
49
+ def from_tensor(weights: torch.Tensor) -> torch.Tensor:
50
+ """
51
+ Wrap an arbitrary 2-D tensor as a grayscale SE.
52
+ :param weights: 2-D tensor of weights
53
+ :return: SE tensor
54
+ """
55
+ raise NotImplementedError