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.
- serron-0.0.1/CMakeLists.txt +97 -0
- serron-0.0.1/PKG-INFO +147 -0
- serron-0.0.1/README.md +110 -0
- serron-0.0.1/csrc/bindings.cpp +18 -0
- serron-0.0.1/csrc/compat/registration.h +16 -0
- serron-0.0.1/csrc/cuda/morphology.cu +17 -0
- serron-0.0.1/csrc/ops.h +13 -0
- serron-0.0.1/pyproject.toml +82 -0
- serron-0.0.1/torch-ext/serron/__init__.py +36 -0
- serron-0.0.1/torch-ext/serron/_cmake_ops.py +33 -0
- serron-0.0.1/torch-ext/serron/autograd.py +12 -0
- serron-0.0.1/torch-ext/serron/enums.py +31 -0
- serron-0.0.1/torch-ext/serron/functional.py +86 -0
- serron-0.0.1/torch-ext/serron/modules.py +66 -0
- serron-0.0.1/torch-ext/serron/py.typed +0 -0
- serron-0.0.1/torch-ext/serron/structuring_element.py +55 -0
|
@@ -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
|
serron-0.0.1/csrc/ops.h
ADDED
|
@@ -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
|