From 432a00506549a5b2823388b21853bb20ff6569f8 Mon Sep 17 00:00:00 2001 From: Yuhan Deng Date: Wed, 24 Jun 2026 03:24:02 +0800 Subject: [PATCH 01/25] feature:add grey erosion kernel --- benchmark/grey.py | 97 ++++++++++ test/test_grey.py | 161 ++++++++++++++++ torchmorph/__init__.py | 2 + torchmorph/csrc/greyero_kernel.cu | 306 ++++++++++++++++++++++++++++++ torchmorph/csrc/torchmorph.cpp | 34 +++- torchmorph/morphology/__init__.py | 2 + torchmorph/morphology/grey.py | 142 ++++++++++++++ 7 files changed, 737 insertions(+), 7 deletions(-) create mode 100644 benchmark/grey.py create mode 100644 test/test_grey.py create mode 100644 torchmorph/csrc/greyero_kernel.cu create mode 100644 torchmorph/morphology/grey.py diff --git a/benchmark/grey.py b/benchmark/grey.py new file mode 100644 index 0000000..1496ef6 --- /dev/null +++ b/benchmark/grey.py @@ -0,0 +1,97 @@ +import scipy.ndimage as ndi +import torch +import torch.utils.benchmark as benchmark +from prettytable import PrettyTable + +import torchmorph as tm + +image_size = [64, 128, 256, 1024] +batch_size = [1, 2, 4, 16] +STRUCTURE_SIZE = 3 +MIN_RUN_TIME = 1.0 + + +def _sync_grey_erosion(x): + out = tm.grey_erosion(x, size=STRUCTURE_SIZE) + torch.cuda.synchronize() + return out + + +def bench_grey_erosion(): + print("\n============================================") + print(" Benchmark: grey erosion ") + print("============================================") + + for batch in batch_size: + table = PrettyTable() + table.field_names = [ + "Size", + "SciPy(ms)", + "Torch 1x(ms)", + "Torch batch(ms)", + "Speedup 1x", + "Speedup batch", + ] + for column in table.field_names: + table.align[column] = "r" + + for size in image_size: + x = torch.randn(batch, 1, size, size, device="cuda", dtype=torch.float32) + # scipy data + x_np_list = [x[i, 0].detach().cpu().numpy() for i in range(batch)] + # torch 1x data + x_one = [x[i : i + 1] for i in range(batch)] + + t_scipy = benchmark.Timer( + stmt="[scipy_erosion(data, size=structure_size) for data in x_np_list]", + globals={ + "scipy_erosion": ndi.grey_erosion, + "structure_size": STRUCTURE_SIZE, + "x_np_list": x_np_list, + }, + ).blocked_autorange(min_run_time=MIN_RUN_TIME) + + # warmup + for data in x_one: + tm.grey_erosion(data, size=STRUCTURE_SIZE) + torch.cuda.synchronize() + + t_torch_1x = benchmark.Timer( + stmt="[tm_erosion(data) for data in x_one]", + globals={ + "tm_erosion": _sync_grey_erosion, + "x_one": x_one, + }, + ).blocked_autorange(min_run_time=MIN_RUN_TIME) + + t_torch_batch = benchmark.Timer( + stmt="tm_erosion(x)", + globals={ + "tm_erosion": _sync_grey_erosion, + "x": x, + }, + ).blocked_autorange(min_run_time=MIN_RUN_TIME) + + scipy_ms = t_scipy.median * 1e3 / batch + torch_1x_ms = t_torch_1x.median * 1e3 / batch + torch_batch_ms = t_torch_batch.median * 1e3 / batch + + speedup1x = scipy_ms / torch_1x_ms + speedupbatch = scipy_ms / torch_batch_ms + + table.add_row( + [ + size, + f"{scipy_ms:.3f}", + f"{torch_1x_ms:.3f}", + f"{torch_batch_ms:.3f}", + f"{speedup1x:.1f}x", + f"{speedupbatch:.1f}x", + ] + ) + print(f"\n=========== Batch size : {batch} ===========") + print(table) + + +if __name__ == "__main__": + bench_grey_erosion() diff --git a/test/test_grey.py b/test/test_grey.py new file mode 100644 index 0000000..ad0feda --- /dev/null +++ b/test/test_grey.py @@ -0,0 +1,161 @@ +import numpy as np +import pytest +import torch +from scipy.ndimage import generate_binary_structure, grey_erosion + +import torchmorph as tm + +pytestmark = pytest.mark.skipif( + not torch.cuda.is_available(), reason="CUDA is required for torchmorph tests" +) + + +def apply_scipy_to_batch(np_input, scipy_func, **kwargs): + input_shape = np_input.shape + batch_shape = input_shape[:2] + spatial_shape = input_shape[2:] + batch_size = int(np.prod(batch_shape)) + flatten_input = np_input.reshape(batch_size, *spatial_shape) + + output = kwargs.pop("output", None) + if output is not None and output.shape != spatial_shape: + raise ValueError(f"output shape must be {spatial_shape}, got {output.shape}") + + results = [] + for sample in flatten_input: + sample_kwargs = dict(kwargs) + if output is not None: + sample_output = np.empty_like(output) + sample_kwargs["output"] = sample_output + result = scipy_func(sample, **sample_kwargs) + results.append(np.asarray(result).copy()) + return np.stack(results, axis=0).reshape(*batch_shape, *spatial_shape) + + +# test cases +case_2d = np.array( + [ + [ + [ + [1.0, 2.0, 3.0], + [4.0, 5.0, 6.0], + [7.0, 8.0, 9.0], + ] + ] + ], + dtype=np.float32, +) + +case_3d = np.arange(2 * 1 * 4 * 4 * 4, dtype=np.float32).reshape(2, 1, 4, 4, 4) + +# structures +footprint_cross = generate_binary_structure(rank=2, connectivity=1).astype(bool) +structure_cross = footprint_cross.astype(np.float32) + +footprint_3d = generate_binary_structure(rank=3, connectivity=2).astype(bool) + +# output tensors +output_2d_tm = torch.empty(1, 1, 3, 3) + + +@pytest.mark.parametrize( + ("np_input", "size", "footprint", "structure", "origin"), + [ + pytest.param(case_2d, 3, None, None, 0, id="2d_size_int"), + pytest.param(case_2d, (3, 3), None, None, 0, id="2d_size_tuple"), + pytest.param(case_2d, None, footprint_cross, None, 0, id="2d_footprint"), + pytest.param(case_2d, None, None, structure_cross, 0, id="2d_structure"), + pytest.param( + case_2d, None, footprint_cross, structure_cross, 0, + id="2d_structure_and_footprint", + ), + pytest.param(case_2d, 3, None, None, 1, id="2d_size_origin_1"), + pytest.param(case_3d, 3, None, None, 0, id="3d_size"), + ] + ) + + +def test_grey_erosion_basic(np_input, size, footprint, structure, origin): + x = torch.as_tensor(np_input, dtype=torch.float32).cuda() + ft = ( + torch.as_tensor(footprint, dtype=torch.bool).cuda() + if footprint is not None + else None + ) + st = ( + torch.as_tensor(structure, dtype=torch.float32).cuda() + if structure is not None + else None + ) + result = tm.grey_erosion( + x, + size=size, + footprint=ft, + structure=st, + origin=origin, + ) + expected_np = apply_scipy_to_batch( + np_input, + grey_erosion, + size=size, + footprint=footprint, + structure=structure, + origin=origin, + ) + expected = torch.as_tensor(expected_np) + torch.testing.assert_close(result.cpu(), expected) + + +@pytest.mark.parametrize( + ("mode", "cval"), + [ + pytest.param("reflect", 0.0, id="mode_reflect"), + pytest.param("constant", 0.0, id="mode_constant"), + pytest.param("constant", 99.0, id="mode_constant_cval99"), + pytest.param("nearest", 0.0, id="mode_nearest"), + pytest.param("mirror", 0.0, id="mode_mirror"), + pytest.param("wrap", 0.0, id="mode_wrap"), + ], +) +def test_grey_erosion_modes(mode, cval): + x = torch.as_tensor(case_2d, dtype=torch.float32).cuda() + result = tm.grey_erosion(x, size=3, mode=mode, cval=cval) + expected_np = apply_scipy_to_batch( + case_2d, grey_erosion, size=3, mode=mode, cval=cval + ) + expected = torch.as_tensor(expected_np) + torch.testing.assert_close(result.cpu(), expected) + + +def test_grey_erosion_output(): + x = torch.as_tensor(case_2d, dtype=torch.float32).cuda() + out = output_2d_tm.cuda() + result = tm.grey_erosion(x, size=3, output=out) + assert result is out + expected_np = apply_scipy_to_batch(case_2d, grey_erosion, size=3) + expected = torch.as_tensor(expected_np) + torch.testing.assert_close(result.cpu(), expected) + + +def test_grey_erosion_no_params(): + x = torch.as_tensor(case_2d, dtype=torch.float32).cuda() + with pytest.raises(ValueError, match="At least one"): + tm.grey_erosion(x) + + +def test_grey_erosion_invalid_device(): + x = torch.tensor([[[1.0, 2.0], [3.0, 4.0]]]) + with pytest.raises(ValueError, match="CUDA"): + tm.grey_erosion(x, size=3) + + +def test_grey_erosion_invalid_ndim(): + x = torch.tensor([1.0, 2.0]).cuda() + with pytest.raises(ValueError, match="at least 3 dimensions"): + tm.grey_erosion(x, size=3) + + +def test_grey_erosion_invalid_mode(): + x = torch.as_tensor(case_2d, dtype=torch.float32).cuda() + with pytest.raises(ValueError, match="Unknown mode"): + tm.grey_erosion(x, size=3, mode="invalid") diff --git a/torchmorph/__init__.py b/torchmorph/__init__.py index be99a6e..485b88e 100644 --- a/torchmorph/__init__.py +++ b/torchmorph/__init__.py @@ -10,6 +10,7 @@ binary_erosion, binary_opening, generate_binary_structure, + grey_erosion, ) __all__ = [ @@ -22,4 +23,5 @@ "binary_dilation", "binary_opening", "binary_closing", + "grey_erosion", ] diff --git a/torchmorph/csrc/greyero_kernel.cu b/torchmorph/csrc/greyero_kernel.cu new file mode 100644 index 0000000..73abb4e --- /dev/null +++ b/torchmorph/csrc/greyero_kernel.cu @@ -0,0 +1,306 @@ +#include +#include +#include +#include +#include +#include +#include + +// 支持的最大空间维度数 +#define GREYERO_MAX_NDIM 8 + +struct GreyErosionGeometry { + int64_t spatial_size[GREYERO_MAX_NDIM]; + int64_t spatial_stride[GREYERO_MAX_NDIM]; + int64_t min_deltas[GREYERO_MAX_NDIM]; + int64_t max_deltas[GREYERO_MAX_NDIM]; +}; + +// 边界填充模式枚举，与 scipy.ndimage 的 mode 参数对应 +enum BorderMode { + MODE_CONSTANT = 0, // 常数填充 + MODE_REFLECT = 1, // 镜像反射（不含边界像素本身） + MODE_NEAREST = 2, // 最近邻重复 + MODE_MIRROR = 3, // 镜像翻转（含边界像素本身） + MODE_WRAP = 4, // 环形/周期填充 +}; + +// 将越界坐标映射回有效范围 +__device__ __forceinline__ int64_t map_coord( + int64_t coord, int64_t size, BorderMode mode, bool* use_cval +) { + *use_cval = false; + if (coord >= 0 && coord < size) return coord; + + switch (mode) { + case MODE_CONSTANT: + // 常数模式：标记使用 cval，返回值无效 + *use_cval = true; + return 0; + case MODE_NEAREST: + // 最近邻：截断到边界 + return coord < 0 ? 0 : size - 1; + case MODE_REFLECT: { + // 反射模式：关于边界反射，不含边界像素 + // 例：[a b c d] -> [b a | a b c d | c d] + int64_t period = 2 * size; + int64_t c = ((coord % period) + period) % period; + if (c >= size) c = period - 1 - c; + return c; + } + case MODE_MIRROR: { + // 镜像模式：关于边界反射，含边界像素 + // 例：[a b c d] -> [c b | a b c d | b c] + int64_t period = 2 * (size - 1); + if (period <= 0) return 0; + int64_t c = ((coord % period) + period) % period; + if (c >= size) c = period - c; + return c; + } + case MODE_WRAP: { + // 环形模式：取模运算实现周期填充 + return ((coord % size) + size) % size; + } + default: + *use_cval = true; + return 0; + } +} + +// 融合 kernel：边界处理 + 最小值计算在一次遍历中完成 +// 每个线程处理一个输出像素 +// struct_meta: 每个结构元素位置的各维度偏移量，末尾附加 flat 偏移量 +__global__ void grey_erosion_fused_kernel( + const float* __restrict__ input, + float* __restrict__ output, + const float* __restrict__ struct_vals, // 结构元素的值（非平坦时使用） + const int64_t* __restrict__ struct_meta, // 每个结构元素的维度偏移量 + flat 偏移量 + const int num_struct, // 结构元素的元素总数 + const int ndim_spatial, // 空间维度数 + const GreyErosionGeometry geom, // 空间尺寸、步长、内部区域边界 + const int64_t total_spatial, // 空间元素总数 + const int64_t batch_channel, // batch * channel 总数 + const BorderMode mode, // 边界模式 + const float cval // 常数模式的填充值 +) { + int64_t idx = blockIdx.x * blockDim.x + threadIdx.x; + if (idx >= batch_channel * total_spatial) return; + + // 计算当前线程属于哪个 batch*channel 和空间位置 + int64_t bc = idx / total_spatial; + int64_t sp = idx % total_spatial; + + // 将展平的空间索引分解为 N 维坐标 + int64_t coords[GREYERO_MAX_NDIM]; + int64_t rem = sp; + for (int d = 0; d < ndim_spatial; d++) { + coords[d] = rem / geom.spatial_stride[d]; + rem %= geom.spatial_stride[d]; + } + + float min_val = FLT_MAX; + const float* in_ptr = input + bc * total_spatial; + const int meta_stride = ndim_spatial + 1; + + bool is_interior = true; + for (int d = 0; d < ndim_spatial; d++) { + if (coords[d] + geom.min_deltas[d] < 0 || + coords[d] + geom.max_deltas[d] >= geom.spatial_size[d]) { + is_interior = false; + break; + } + } + + if (is_interior) { + for (int i = 0; i < num_struct; i++) { + float val = in_ptr[sp + struct_meta[i * meta_stride + ndim_spatial]]; + min_val = fminf(min_val, val - struct_vals[i]); + } + output[idx] = min_val; + return; + } + + // 遍历结构元素的所有位置，取最小值 + for (int i = 0; i < num_struct; i++) { + const int64_t* meta = struct_meta + i * meta_stride; + + // 对每个维度，将坐标加上偏移后通过边界模式映射回有效范围 + bool use_cval = false; + int64_t flat_idx = 0; + for (int d = 0; d < ndim_spatial; d++) { + bool uc = false; + int64_t mapped = map_coord( + coords[d] + meta[d], geom.spatial_size[d], mode, &uc + ); + if (uc) { use_cval = true; break; } + flat_idx += mapped * geom.spatial_stride[d]; + } + + // 取输入值（或 cval）减去结构元素值，更新最小值 + float val = use_cval ? cval : in_ptr[flat_idx]; + min_val = fminf(min_val, val - struct_vals[i]); + } + + output[idx] = min_val; +} + +// ============================================================================== +// 主机端入口函数 +// ============================================================================== +torch::Tensor grey_erosion_cuda( + torch::Tensor input, // (B, C, *spatial) float32, 连续, CUDA + torch::Tensor structure, // (*spatial) float32, CPU + torch::Tensor footprint, // (*spatial) bool CPU；空 tensor 表示全有效 + std::vector origin_vec, + int mode_int, + float cval +) { + TORCH_CHECK(input.is_cuda(), "input 必须是 CUDA 张量"); + TORCH_CHECK(!structure.is_cuda(), "structure 必须是 CPU 张量"); + TORCH_CHECK(!footprint.is_cuda(), "footprint 必须是 CPU 张量"); + TORCH_CHECK(input.is_contiguous(), "input 必须是连续的"); + TORCH_CHECK(structure.is_contiguous(), "structure 必须是连续的"); + TORCH_CHECK(footprint.is_contiguous(), "footprint 必须是连续的"); + TORCH_CHECK(input.dtype() == torch::kFloat32, "input 必须是 float32"); + TORCH_CHECK(structure.dtype() == torch::kFloat32, "structure 必须是 float32"); + TORCH_CHECK(footprint.dtype() == torch::kBool, "footprint 必须是 bool"); + + int ndim_spatial = structure.dim(); + TORCH_CHECK(ndim_spatial > 0 && ndim_spatial <= GREYERO_MAX_NDIM, + "structure 维度必须在 1-", GREYERO_MAX_NDIM, + " 之间，实际为 ", ndim_spatial); + TORCH_CHECK((int)origin_vec.size() == ndim_spatial, + "origin 长度必须与 structure 维度一致"); + TORCH_CHECK(input.dim() == ndim_spatial + 2, + "input 空间维度必须与 structure 维度一致"); + + const bool use_footprint = footprint.numel() != 0; + if (use_footprint) { + TORCH_CHECK(footprint.dim() == ndim_spatial, + "footprint 维度必须与 structure 维度一致"); + for (int d = 0; d < ndim_spatial; d++) { + TORCH_CHECK(footprint.size(d) == structure.size(d), + "footprint 尺寸必须与 structure 尺寸一致"); + } + } + + auto output = torch::empty_like(input); + + int B = input.size(0); + int C = input.size(1); + int64_t batch_channel = (int64_t)B * C; + + // 计算逻辑空间尺寸和步长 + std::vector h_spatial_size(ndim_spatial); + std::vector h_spatial_stride(ndim_spatial); + for (int d = 0; d < ndim_spatial; d++) { + h_spatial_size[d] = input.size(d + 2); + } + h_spatial_stride[ndim_spatial - 1] = 1; + for (int d = ndim_spatial - 2; d >= 0; d--) { + h_spatial_stride[d] = h_spatial_stride[d + 1] * h_spatial_size[d + 1]; + } + + int64_t total_spatial = 1; + for (int d = 0; d < ndim_spatial; d++) { + total_spatial *= h_spatial_size[d]; + } + + GreyErosionGeometry h_geom; + for (int d = 0; d < GREYERO_MAX_NDIM; d++) { + h_geom.spatial_size[d] = 0; + h_geom.spatial_stride[d] = 0; + h_geom.min_deltas[d] = 0; + h_geom.max_deltas[d] = 0; + } + for (int d = 0; d < ndim_spatial; d++) { + h_geom.spatial_size[d] = h_spatial_size[d]; + h_geom.spatial_stride[d] = h_spatial_stride[d]; + h_geom.min_deltas[d] = INT64_MAX; + h_geom.max_deltas[d] = INT64_MIN; + } + + // 展平结构元素，仅保留 footprint 中有效的位置。 + auto struct_flat = structure.flatten().contiguous(); + const float* struct_ptr = struct_flat.data_ptr(); + const int64_t total_struct = struct_flat.numel(); + + torch::Tensor footprint_flat; + const bool* footprint_ptr = nullptr; + if (use_footprint) { + footprint_flat = footprint.flatten().contiguous(); + footprint_ptr = footprint_flat.data_ptr(); + } + + std::vector h_structure_stride(ndim_spatial); + h_structure_stride[ndim_spatial - 1] = 1; + for (int d = ndim_spatial - 2; d >= 0; d--) { + h_structure_stride[d] = + h_structure_stride[d + 1] * structure.size(d + 1); + } + + std::vector h_struct_vals; + std::vector h_struct_meta; + h_struct_vals.reserve(total_struct); + h_struct_meta.reserve(total_struct * (ndim_spatial + 1)); + + for (int64_t i = 0; i < total_struct; i++) { + if (use_footprint && !footprint_ptr[i]) { + continue; + } + + h_struct_vals.push_back(struct_ptr[i]); + + int64_t tmp = i; + int64_t flat_offset = 0; + for (int d = 0; d < ndim_spatial; d++) { + int64_t coord_d = tmp / h_structure_stride[d]; + tmp %= h_structure_stride[d]; + + // 偏移量 = 坐标 - 中心（中心考虑了 origin 偏移） + int64_t center_d = structure.size(d) / 2 + origin_vec[d]; + int64_t delta = coord_d - center_d; + h_struct_meta.push_back(delta); + flat_offset += delta * h_spatial_stride[d]; + h_geom.min_deltas[d] = std::min(h_geom.min_deltas[d], delta); + h_geom.max_deltas[d] = std::max(h_geom.max_deltas[d], delta); + } + h_struct_meta.push_back(flat_offset); + } + + int num_struct = static_cast(h_struct_vals.size()); + TORCH_CHECK(num_struct > 0, "footprint 至少需要包含一个有效位置"); + + // 将数据传到 GPU：先创建 CPU tensor，再 .to(cuda) 做真正的拷贝 + auto opts_f = torch::TensorOptions().dtype(torch::kFloat32); + auto opts_i = torch::TensorOptions().dtype(torch::kInt64); + + auto d_struct_vals = torch::from_blob( + h_struct_vals.data(), {num_struct}, opts_f).to(input.device()); + auto d_struct_meta = torch::from_blob( + h_struct_meta.data(), {num_struct * (ndim_spatial + 1)}, opts_i + ).to(input.device()); + + BorderMode bmode = static_cast(mode_int); + + // 启动 kernel + int64_t total_threads = batch_channel * total_spatial; + int threads = 256; + int blocks = (int)((total_threads + threads - 1) / threads); + + grey_erosion_fused_kernel<<>>( + input.data_ptr(), + output.data_ptr(), + d_struct_vals.data_ptr(), + d_struct_meta.data_ptr(), + num_struct, + ndim_spatial, + h_geom, + total_spatial, + batch_channel, + bmode, + cval + ); + + return output; +} diff --git a/torchmorph/csrc/torchmorph.cpp b/torchmorph/csrc/torchmorph.cpp index 0adac87..c397474 100644 --- a/torchmorph/csrc/torchmorph.cpp +++ b/torchmorph/csrc/torchmorph.cpp @@ -1,9 +1,19 @@ #include -// Declare CUDA implementations +// CUDA 实现声明 torch::Tensor add_cuda(torch::Tensor input, float scalar); -// Distance Transform functions +// 灰度腐蚀（融合 kernel，边界处理在 CUDA 内部完成） +torch::Tensor grey_erosion_cuda( + torch::Tensor input, + torch::Tensor structure, + torch::Tensor footprint, + std::vector origin, + int mode, + float cval +); + +// 距离变换函数 std::tuple edt_cuda( torch::Tensor input, std::vector sampling, @@ -29,25 +39,35 @@ std::tuple bfdt_cuda( PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { - m.def("add_cuda", &add_cuda, "Add tensor with scalar"); + m.def("add_cuda", &add_cuda, "标量加法"); + + // 灰度形态学 + m.def("grey_erosion_cuda", &grey_erosion_cuda, + "N维融合灰度腐蚀（CUDA，边界处理在内部完成）", + py::arg("input"), + py::arg("structure"), + py::arg("footprint"), + py::arg("origin"), + py::arg("mode"), + py::arg("cval")); - // Distance Transform + // 距离变换 m.def("edt_cuda", &edt_cuda, - "Exact Euclidean Distance Transform (Felzenszwalb algorithm)", + "精确欧氏距离变换（Felzenszwalb 算法）", py::arg("input"), py::arg("sampling"), py::arg("return_distances") = true, py::arg("return_indices") = false, py::arg("algorithm") = "exact"); m.def("cdt_cuda", &cdt_cuda, - "Chamfer Distance Transform", + "棋盘/曼哈顿距离变换", py::arg("input"), py::arg("metric") = "chessboard", py::arg("return_distances") = true, py::arg("return_indices") = false); m.def("bfdt_cuda", &bfdt_cuda, - "Brute-Force Distance Transform", + "暴力距离变换", py::arg("input"), py::arg("metric") = "euclidean", py::arg("sampling") = std::vector(), diff --git a/torchmorph/morphology/__init__.py b/torchmorph/morphology/__init__.py index 0efa28d..5819ec3 100644 --- a/torchmorph/morphology/__init__.py +++ b/torchmorph/morphology/__init__.py @@ -1,4 +1,5 @@ from .binary import binary_closing, binary_dilation, binary_erosion, binary_opening +from .grey import grey_erosion from .structure import generate_binary_structure __all__ = [ @@ -7,4 +8,5 @@ "binary_dilation", "binary_opening", "binary_closing", + "grey_erosion", ] diff --git a/torchmorph/morphology/grey.py b/torchmorph/morphology/grey.py new file mode 100644 index 0000000..bc1757a --- /dev/null +++ b/torchmorph/morphology/grey.py @@ -0,0 +1,142 @@ +import torch +from torch import Tensor + +from .. import _C + +_MODE_MAP = { + 'constant': 0, + 'reflect': 1, + 'nearest': 2, + 'mirror': 3, + 'wrap': 4, +} + + +def grey_erosion( + input: Tensor, + size: int | tuple[int, ...] | None = None, + footprint: Tensor | None = None, + structure: Tensor | None = None, + output: Tensor | None = None, + mode: str = 'reflect', + cval: float = 0.0, + origin: int | tuple[int, ...] = 0, +) -> Tensor: + """N-dimensional grey erosion for (B, C, Spatial...) tensors. + + For a single image or volume, add batch and channel dimensions first. + + The structuring element is resolved with the following priority: + 1. `structure` (non-flat, values are subtracted) + 2. `footprint` (flat, only positions where footprint > 0 are used) + 3. `size` (flat rectangle shaped by `size`) + Exactly one of these three parameters must be provided. + + Args: + input: Input tensor in (B, C, Spatial...) format. Must be on CUDA. + size: Shape of a flat rectangular structuring element. Mutually + exclusive with footprint and structure. + footprint: Binary structuring element. Non-zero positions define + the neighborhood over which the minimum is computed. Mutually + exclusive with size and structure. + structure: Value structuring element. Values are subtracted from + each pixel in the neighborhood. Mutually exclusive with size + and footprint. When given, `footprint` may optionally be + provided as a mask to restrict the neighborhood. + output: Pre-allocated output tensor. If provided, the result is + written in-place and returned. + mode: Border mode. One of 'constant', 'reflect', 'nearest', + 'mirror', 'wrap'. Default is 'reflect'. + cval: Constant fill value used when mode='constant'. Default is 0.0. + origin: Origin offset. Can be an int (same for all axes) or tuple. + + Returns: + Grey-eroded tensor with the same shape as input. + """ + if not input.is_cuda: + raise ValueError('Input tensor must be on CUDA device.') + if input.ndim < 3: + raise ValueError( + f'Input must be (B, C, Spatial...) with at least 3 dimensions, ' + f'got {input.shape}.' + ) + if input.numel() == 0: + raise ValueError( + f'Invalid input: empty tensor with shape {input.shape}.' + ) + + spatial_ndim = input.ndim - 2 + + footprint_cpu = torch.empty(0, dtype=torch.bool) + + # Resolve structuring element (priority: structure > footprint > size) + if structure is not None: + struct = structure.to(torch.float32) + if footprint is not None: + mask = footprint.to(torch.bool) + if tuple(mask.shape) != tuple(struct.shape): + raise ValueError( + f'Footprint shape {tuple(mask.shape)} must match structure ' + f'shape {tuple(struct.shape)}.' + ) + if not mask.any(): + raise ValueError('All-zero footprint is not supported.') + footprint_cpu = ( + mask.detach().to(device='cpu', dtype=torch.bool).contiguous() + ) + elif footprint is not None: + mask = footprint.to(torch.bool) + if not mask.any(): + raise ValueError('All-zero footprint is not supported.') + struct = torch.zeros(tuple(mask.shape), dtype=torch.float32) + footprint_cpu = ( + mask.detach().to(device='cpu', dtype=torch.bool).contiguous() + ) + elif size is not None: + if isinstance(size, int): + size = (size,) * spatial_ndim + struct = torch.zeros(size, dtype=torch.float32) + else: + raise ValueError( + 'At least one of size, footprint, or structure must be specified.' + ) + + if struct.ndim != spatial_ndim: + raise ValueError( + f'Structure dimension {struct.ndim} must match spatial dimension ' + f'{spatial_ndim}.' + ) + struct = struct.detach().to(device='cpu', dtype=torch.float32).contiguous() + + # Normalize origin + if isinstance(origin, int): + origin_list = [origin] * spatial_ndim + else: + origin_list = list(origin) + if len(origin_list) != spatial_ndim: + raise ValueError( + f'Origin length {len(origin_list)} must match spatial dimension ' + f'{spatial_ndim}.' + ) + + # Map mode string to int + if mode not in _MODE_MAP: + raise ValueError( + f"Unknown mode '{mode}'. Must be one of {list(_MODE_MAP.keys())}." + ) + mode_int = _MODE_MAP[mode] + + # CUDA kernel call + result = _C.grey_erosion_cuda( + input.contiguous().float(), + struct.contiguous(), + footprint_cpu, + origin_list, + mode_int, + cval, + ) + + if output is not None: + output.copy_(result) + return output + return result From 7d01c22aa7aeaf81989a1d020280e532e0d828a3 Mon Sep 17 00:00:00 2001 From: Yuhandeng Date: Wed, 24 Jun 2026 03:38:36 +0800 Subject: [PATCH 02/25] ENG --- torchmorph/csrc/greyero_kernel.cu | 102 ++++++++++++------------------ torchmorph/csrc/torchmorph.cpp | 15 ++--- 2 files changed, 47 insertions(+), 70 deletions(-) diff --git a/torchmorph/csrc/greyero_kernel.cu b/torchmorph/csrc/greyero_kernel.cu index 73abb4e..51c92a2 100644 --- a/torchmorph/csrc/greyero_kernel.cu +++ b/torchmorph/csrc/greyero_kernel.cu @@ -6,7 +6,6 @@ #include #include -// 支持的最大空间维度数 #define GREYERO_MAX_NDIM 8 struct GreyErosionGeometry { @@ -16,16 +15,15 @@ struct GreyErosionGeometry { int64_t max_deltas[GREYERO_MAX_NDIM]; }; -// 边界填充模式枚举，与 scipy.ndimage 的 mode 参数对应 +// Matches scipy.ndimage border modes. enum BorderMode { - MODE_CONSTANT = 0, // 常数填充 - MODE_REFLECT = 1, // 镜像反射（不含边界像素本身） - MODE_NEAREST = 2, // 最近邻重复 - MODE_MIRROR = 3, // 镜像翻转（含边界像素本身） - MODE_WRAP = 4, // 环形/周期填充 + MODE_CONSTANT = 0, // Constant padding + MODE_REFLECT = 1, // Half-sample symmetric + MODE_NEAREST = 2, // Repeat nearest edge value + MODE_MIRROR = 3, // Whole-sample symmetric + MODE_WRAP = 4, // Periodic wrap padding }; -// 将越界坐标映射回有效范围 __device__ __forceinline__ int64_t map_coord( int64_t coord, int64_t size, BorderMode mode, bool* use_cval ) { @@ -34,23 +32,19 @@ __device__ __forceinline__ int64_t map_coord( switch (mode) { case MODE_CONSTANT: - // 常数模式：标记使用 cval，返回值无效 *use_cval = true; return 0; case MODE_NEAREST: - // 最近邻：截断到边界 return coord < 0 ? 0 : size - 1; case MODE_REFLECT: { - // 反射模式：关于边界反射，不含边界像素 - // 例：[a b c d] -> [b a | a b c d | c d] + // Half-sample symmetric: [a b c d] -> [b a | a b c d | c d] int64_t period = 2 * size; int64_t c = ((coord % period) + period) % period; if (c >= size) c = period - 1 - c; return c; } case MODE_MIRROR: { - // 镜像模式：关于边界反射，含边界像素 - // 例：[a b c d] -> [c b | a b c d | b c] + // Whole-sample symmetric: [a b c d] -> [c b | a b c d | b c] int64_t period = 2 * (size - 1); if (period <= 0) return 0; int64_t c = ((coord % period) + period) % period; @@ -58,7 +52,6 @@ __device__ __forceinline__ int64_t map_coord( return c; } case MODE_WRAP: { - // 环形模式：取模运算实现周期填充 return ((coord % size) + size) % size; } default: @@ -67,30 +60,27 @@ __device__ __forceinline__ int64_t map_coord( } } -// 融合 kernel：边界处理 + 最小值计算在一次遍历中完成 -// 每个线程处理一个输出像素 -// struct_meta: 每个结构元素位置的各维度偏移量，末尾附加 flat 偏移量 +// Fused erosion kernel. struct_meta stores per-element dimension offsets +// followed by the precomputed flat offset used by the interior fast path. __global__ void grey_erosion_fused_kernel( const float* __restrict__ input, float* __restrict__ output, - const float* __restrict__ struct_vals, // 结构元素的值（非平坦时使用） - const int64_t* __restrict__ struct_meta, // 每个结构元素的维度偏移量 + flat 偏移量 - const int num_struct, // 结构元素的元素总数 - const int ndim_spatial, // 空间维度数 - const GreyErosionGeometry geom, // 空间尺寸、步长、内部区域边界 - const int64_t total_spatial, // 空间元素总数 - const int64_t batch_channel, // batch * channel 总数 - const BorderMode mode, // 边界模式 - const float cval // 常数模式的填充值 + const float* __restrict__ struct_vals, + const int64_t* __restrict__ struct_meta, + const int num_struct, + const int ndim_spatial, + const GreyErosionGeometry geom, + const int64_t total_spatial, + const int64_t batch_channel, + const BorderMode mode, + const float cval ) { int64_t idx = blockIdx.x * blockDim.x + threadIdx.x; if (idx >= batch_channel * total_spatial) return; - // 计算当前线程属于哪个 batch*channel 和空间位置 int64_t bc = idx / total_spatial; int64_t sp = idx % total_spatial; - // 将展平的空间索引分解为 N 维坐标 int64_t coords[GREYERO_MAX_NDIM]; int64_t rem = sp; for (int d = 0; d < ndim_spatial; d++) { @@ -113,18 +103,17 @@ __global__ void grey_erosion_fused_kernel( if (is_interior) { for (int i = 0; i < num_struct; i++) { - float val = in_ptr[sp + struct_meta[i * meta_stride + ndim_spatial]]; + int64_t offset = struct_meta[i * meta_stride + ndim_spatial]; + float val = in_ptr[sp + offset]; min_val = fminf(min_val, val - struct_vals[i]); } output[idx] = min_val; return; } - // 遍历结构元素的所有位置，取最小值 for (int i = 0; i < num_struct; i++) { const int64_t* meta = struct_meta + i * meta_stride; - // 对每个维度，将坐标加上偏移后通过边界模式映射回有效范围 bool use_cval = false; int64_t flat_idx = 0; for (int d = 0; d < ndim_spatial; d++) { @@ -136,7 +125,6 @@ __global__ void grey_erosion_fused_kernel( flat_idx += mapped * geom.spatial_stride[d]; } - // 取输入值（或 cval）减去结构元素值，更新最小值 float val = use_cval ? cval : in_ptr[flat_idx]; min_val = fminf(min_val, val - struct_vals[i]); } @@ -144,43 +132,40 @@ __global__ void grey_erosion_fused_kernel( output[idx] = min_val; } -// ============================================================================== -// 主机端入口函数 -// ============================================================================== torch::Tensor grey_erosion_cuda( - torch::Tensor input, // (B, C, *spatial) float32, 连续, CUDA - torch::Tensor structure, // (*spatial) float32, CPU - torch::Tensor footprint, // (*spatial) bool CPU；空 tensor 表示全有效 + torch::Tensor input, + torch::Tensor structure, + torch::Tensor footprint, std::vector origin_vec, int mode_int, float cval ) { - TORCH_CHECK(input.is_cuda(), "input 必须是 CUDA 张量"); - TORCH_CHECK(!structure.is_cuda(), "structure 必须是 CPU 张量"); - TORCH_CHECK(!footprint.is_cuda(), "footprint 必须是 CPU 张量"); - TORCH_CHECK(input.is_contiguous(), "input 必须是连续的"); - TORCH_CHECK(structure.is_contiguous(), "structure 必须是连续的"); - TORCH_CHECK(footprint.is_contiguous(), "footprint 必须是连续的"); - TORCH_CHECK(input.dtype() == torch::kFloat32, "input 必须是 float32"); - TORCH_CHECK(structure.dtype() == torch::kFloat32, "structure 必须是 float32"); - TORCH_CHECK(footprint.dtype() == torch::kBool, "footprint 必须是 bool"); + TORCH_CHECK(input.is_cuda(), "input must be a CUDA tensor"); + TORCH_CHECK(!structure.is_cuda(), "structure must be a CPU tensor"); + TORCH_CHECK(!footprint.is_cuda(), "footprint must be a CPU tensor"); + TORCH_CHECK(input.is_contiguous(), "input must be contiguous"); + TORCH_CHECK(structure.is_contiguous(), "structure must be contiguous"); + TORCH_CHECK(footprint.is_contiguous(), "footprint must be contiguous"); + TORCH_CHECK(input.dtype() == torch::kFloat32, "input must be float32"); + TORCH_CHECK(structure.dtype() == torch::kFloat32, "structure must be float32"); + TORCH_CHECK(footprint.dtype() == torch::kBool, "footprint must be bool"); int ndim_spatial = structure.dim(); TORCH_CHECK(ndim_spatial > 0 && ndim_spatial <= GREYERO_MAX_NDIM, - "structure 维度必须在 1-", GREYERO_MAX_NDIM, - " 之间，实际为 ", ndim_spatial); + "structure dimension must be in 1-", GREYERO_MAX_NDIM, + ", got ", ndim_spatial); TORCH_CHECK((int)origin_vec.size() == ndim_spatial, - "origin 长度必须与 structure 维度一致"); + "origin length must match structure dimensions"); TORCH_CHECK(input.dim() == ndim_spatial + 2, - "input 空间维度必须与 structure 维度一致"); + "input spatial dimensions must match structure dimensions"); const bool use_footprint = footprint.numel() != 0; if (use_footprint) { TORCH_CHECK(footprint.dim() == ndim_spatial, - "footprint 维度必须与 structure 维度一致"); + "footprint dimensions must match structure dimensions"); for (int d = 0; d < ndim_spatial; d++) { TORCH_CHECK(footprint.size(d) == structure.size(d), - "footprint 尺寸必须与 structure 尺寸一致"); + "footprint shape must match structure shape"); } } @@ -190,7 +175,6 @@ torch::Tensor grey_erosion_cuda( int C = input.size(1); int64_t batch_channel = (int64_t)B * C; - // 计算逻辑空间尺寸和步长 std::vector h_spatial_size(ndim_spatial); std::vector h_spatial_stride(ndim_spatial); for (int d = 0; d < ndim_spatial; d++) { @@ -220,7 +204,7 @@ torch::Tensor grey_erosion_cuda( h_geom.max_deltas[d] = INT64_MIN; } - // 展平结构元素，仅保留 footprint 中有效的位置。 + // Keep only active footprint positions and precompute offsets. auto struct_flat = structure.flatten().contiguous(); const float* struct_ptr = struct_flat.data_ptr(); const int64_t total_struct = struct_flat.numel(); @@ -257,7 +241,6 @@ torch::Tensor grey_erosion_cuda( int64_t coord_d = tmp / h_structure_stride[d]; tmp %= h_structure_stride[d]; - // 偏移量 = 坐标 - 中心（中心考虑了 origin 偏移） int64_t center_d = structure.size(d) / 2 + origin_vec[d]; int64_t delta = coord_d - center_d; h_struct_meta.push_back(delta); @@ -269,9 +252,9 @@ torch::Tensor grey_erosion_cuda( } int num_struct = static_cast(h_struct_vals.size()); - TORCH_CHECK(num_struct > 0, "footprint 至少需要包含一个有效位置"); + TORCH_CHECK(num_struct > 0, + "footprint must contain at least one active position"); - // 将数据传到 GPU：先创建 CPU tensor，再 .to(cuda) 做真正的拷贝 auto opts_f = torch::TensorOptions().dtype(torch::kFloat32); auto opts_i = torch::TensorOptions().dtype(torch::kInt64); @@ -283,7 +266,6 @@ torch::Tensor grey_erosion_cuda( BorderMode bmode = static_cast(mode_int); - // 启动 kernel int64_t total_threads = batch_channel * total_spatial; int threads = 256; int blocks = (int)((total_threads + threads - 1) / threads); diff --git a/torchmorph/csrc/torchmorph.cpp b/torchmorph/csrc/torchmorph.cpp index c397474..353d489 100644 --- a/torchmorph/csrc/torchmorph.cpp +++ b/torchmorph/csrc/torchmorph.cpp @@ -1,9 +1,7 @@ #include -// CUDA 实现声明 torch::Tensor add_cuda(torch::Tensor input, float scalar); -// 灰度腐蚀（融合 kernel，边界处理在 CUDA 内部完成） torch::Tensor grey_erosion_cuda( torch::Tensor input, torch::Tensor structure, @@ -13,7 +11,6 @@ torch::Tensor grey_erosion_cuda( float cval ); -// 距离变换函数 std::tuple edt_cuda( torch::Tensor input, std::vector sampling, @@ -39,11 +36,10 @@ std::tuple bfdt_cuda( PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { - m.def("add_cuda", &add_cuda, "标量加法"); + m.def("add_cuda", &add_cuda, "Add tensor with scalar"); - // 灰度形态学 m.def("grey_erosion_cuda", &grey_erosion_cuda, - "N维融合灰度腐蚀（CUDA，边界处理在内部完成）", + "N-dimensional fused grey erosion", py::arg("input"), py::arg("structure"), py::arg("footprint"), @@ -51,23 +47,22 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { py::arg("mode"), py::arg("cval")); - // 距离变换 m.def("edt_cuda", &edt_cuda, - "精确欧氏距离变换（Felzenszwalb 算法）", + "Exact Euclidean Distance Transform (Felzenszwalb algorithm)", py::arg("input"), py::arg("sampling"), py::arg("return_distances") = true, py::arg("return_indices") = false, py::arg("algorithm") = "exact"); m.def("cdt_cuda", &cdt_cuda, - "棋盘/曼哈顿距离变换", + "Chessboard/Manhattan distance transform", py::arg("input"), py::arg("metric") = "chessboard", py::arg("return_distances") = true, py::arg("return_indices") = false); m.def("bfdt_cuda", &bfdt_cuda, - "暴力距离变换", + "Brute-force distance transform", py::arg("input"), py::arg("metric") = "euclidean", py::arg("sampling") = std::vector(), From cdf76fe78daa37b10bdef2d3d0f8b1785bee6c61 Mon Sep 17 00:00:00 2001 From: Yuhandeng Date: Sat, 4 Jul 2026 22:00:02 +0800 Subject: [PATCH 03/25] dilation opening closing funciton --- CONTEXT.md | 53 +++++++ README.md | 17 +- benchmark/grey.py | 110 ++++++++----- pytest.ini | 3 +- test/test_grey.py | 250 +++++++++++++++++------------- torchmorph/__init__.py | 6 + torchmorph/csrc/greyero_kernel.cu | 112 ++++++++++--- torchmorph/csrc/torchmorph.cpp | 18 +++ torchmorph/morphology/__init__.py | 5 +- torchmorph/morphology/grey.py | 198 ++++++++++++++++------- 10 files changed, 543 insertions(+), 229 deletions(-) create mode 100644 CONTEXT.md diff --git a/CONTEXT.md b/CONTEXT.md new file mode 100644 index 0000000..d849f48 --- /dev/null +++ b/CONTEXT.md @@ -0,0 +1,53 @@ +# TorchMorph + +TorchMorph provides tensor-native morphological image operations and distance transforms for PyTorch workloads. Its public language follows SciPy ndimage where possible while making batch and channel dimensions explicit. + +## Language + +**Morphological Operation**: +An operation that transforms binary or grey-valued spatial samples by comparing each location with a neighborhood defined by a structuring element. +_Avoid_: filter, convolution + +**Binary Morphology**: +Morphological operations where every non-zero input value is treated as foreground and zero is treated as background. +_Avoid_: boolean convolution, mask filtering + +**Grey Morphology**: +Morphological operations over numeric intensities, where the result depends on ordered values rather than only foreground/background membership. +_Avoid_: grayscale filter + +**Grey Erosion**: +A grey morphology operation that selects the minimum response from a neighborhood defined by a structuring element. +_Avoid_: minimum filter + +**Grey Dilation**: +A grey morphology operation that selects the maximum response from a neighborhood defined by a structuring element. +_Avoid_: maximum filter + +**Grey Opening**: +A grey morphology operation formed by grey erosion followed by grey dilation with the same structuring element. +_Avoid_: erosion-dilation filter + +**Grey Closing**: +A grey morphology operation formed by grey dilation followed by grey erosion with the same structuring element. +_Avoid_: dilation-erosion filter + +**Structuring Element**: +The neighborhood shape or weighted neighborhood used by a morphological operation. +_Avoid_: kernel, filter + +**Footprint**: +A binary structuring element that selects which neighboring locations participate in a grey morphology operation. +_Avoid_: mask, stencil + +**Feature Transform**: +The nearest-background coordinate result returned alongside a distance transform. +_Avoid_: index map, nearest point map + +**Spatial Dimensions**: +The trailing dimensions of an input tensor that represent the image, volume, or higher-dimensional sample being transformed. +_Avoid_: data dimensions + +**Batch-Channel Tensor**: +A tensor shaped as `(B, C, Spatial...)`, where `B` and `C` are independent leading dimensions and every `(B, C)` slice is transformed independently. +_Avoid_: image tensor, BCHW-only tensor diff --git a/README.md b/README.md index e3b68fc..1fe1d87 100644 --- a/README.md +++ b/README.md @@ -2,18 +2,23 @@ > ⚡ CUDA-Accelerated, Batch-Parallel Morphological Transformations for PyTorch -TorchMorph is a lightweight, extensible library that brings **GPU-accelerated morphological operations** into the PyTorch ecosystem. +TorchMorph is a lightweight, extensible library that brings **GPU-accelerated morphological operations** into the PyTorch ecosystem. It provides a **clean Python API** backed by **custom CUDA kernels**, enabling highly efficient, **batch-parallel** transformations for real-time and large-scale vision tasks. --- ## 🚀 Key Features -- ⚡ **CUDA Acceleration** – All operators are implemented with native CUDA kernels for maximum throughput. -- 🧩 **Seamless PyTorch Integration** – Accepts and returns `torch.Tensor` objects, fully compatible with autograd and CUDA streams. -- 🧠 **Highly Batch-Parallel** – Optimized to process large batches and multi-dimensional inputs concurrently. -- 🧱 **Modular Design** – Each operation is isolated in its own kernel, making it easy to add or extend transformations. -- ✅ **Lightweight & Self-Contained** – No third-party dependencies beyond PyTorch and a working CUDA toolkit. +- ⚡ **CUDA Acceleration** – All operators are implemented with native CUDA kernels for maximum throughput. +- 🧩 **PyTorch Tensor Integration** – Accepts and returns CUDA `torch.Tensor` objects for direct use in tensor-processing pipelines. +- 🧠 **Highly Batch-Parallel** – Optimized to process large batches and multi-dimensional inputs concurrently. +- 🧱 **Modular Design** – Public operators reuse shared CUDA morphology primitives and compose them into higher-level transformations. +- ✅ **Lightweight & Self-Contained** – No third-party dependencies beyond PyTorch and a working CUDA toolkit. + + +> **NaN values:** NaN propagation is not guaranteed to match SciPy. Replace or +> otherwise handle NaN values before calling a morphological operator when +> deterministic missing-value behavior is required. --- diff --git a/benchmark/grey.py b/benchmark/grey.py index 1496ef6..a53bf75 100644 --- a/benchmark/grey.py +++ b/benchmark/grey.py @@ -1,3 +1,5 @@ +import argparse + import scipy.ndimage as ndi import torch import torch.utils.benchmark as benchmark @@ -5,24 +7,33 @@ import torchmorph as tm -image_size = [64, 128, 256, 1024] -batch_size = [1, 2, 4, 16] +IMAGE_SIZES = [64, 128, 256, 1024] +BATCH_SIZES = [1, 2, 4, 16] STRUCTURE_SIZE = 3 MIN_RUN_TIME = 1.0 +GREY_OPERATORS = { + "erosion": (ndi.grey_erosion, tm.grey_erosion), + "dilation": (ndi.grey_dilation, tm.grey_dilation), + "opening": (ndi.grey_opening, tm.grey_opening), + "closing": (ndi.grey_closing, tm.grey_closing), +} + -def _sync_grey_erosion(x): - out = tm.grey_erosion(x, size=STRUCTURE_SIZE) +def run_cuda(torch_op, x): + result = torch_op(x, size=STRUCTURE_SIZE) torch.cuda.synchronize() - return out + return result + +def bench_grey_operator(operation): + scipy_op, torch_op = GREY_OPERATORS[operation] -def bench_grey_erosion(): print("\n============================================") - print(" Benchmark: grey erosion ") + print(f" Benchmark: grey {operation} ") print("============================================") - for batch in batch_size: + for batch_size in BATCH_SIZES: table = PrettyTable() table.field_names = [ "Size", @@ -35,63 +46,84 @@ def bench_grey_erosion(): for column in table.field_names: table.align[column] = "r" - for size in image_size: - x = torch.randn(batch, 1, size, size, device="cuda", dtype=torch.float32) - # scipy data - x_np_list = [x[i, 0].detach().cpu().numpy() for i in range(batch)] - # torch 1x data - x_one = [x[i : i + 1] for i in range(batch)] + for image_size in IMAGE_SIZES: + x = torch.randn( + batch_size, + 1, + image_size, + image_size, + device="cuda", + dtype=torch.float32, + ) + x_np_list = [x[i, 0].detach().cpu().numpy() for i in range(batch_size)] + x_one = [x[i : i + 1] for i in range(batch_size)] - t_scipy = benchmark.Timer( - stmt="[scipy_erosion(data, size=structure_size) for data in x_np_list]", + scipy_time = benchmark.Timer( + stmt="[scipy_op(data, size=structure_size) for data in inputs]", globals={ - "scipy_erosion": ndi.grey_erosion, + "scipy_op": scipy_op, "structure_size": STRUCTURE_SIZE, - "x_np_list": x_np_list, + "inputs": x_np_list, }, ).blocked_autorange(min_run_time=MIN_RUN_TIME) - # warmup for data in x_one: - tm.grey_erosion(data, size=STRUCTURE_SIZE) + torch_op(data, size=STRUCTURE_SIZE) torch.cuda.synchronize() - t_torch_1x = benchmark.Timer( - stmt="[tm_erosion(data) for data in x_one]", + torch_single_time = benchmark.Timer( + stmt="[run_cuda(torch_op, data) for data in inputs]", globals={ - "tm_erosion": _sync_grey_erosion, - "x_one": x_one, + "run_cuda": run_cuda, + "torch_op": torch_op, + "inputs": x_one, }, ).blocked_autorange(min_run_time=MIN_RUN_TIME) - t_torch_batch = benchmark.Timer( - stmt="tm_erosion(x)", + torch_batch_time = benchmark.Timer( + stmt="run_cuda(torch_op, x)", globals={ - "tm_erosion": _sync_grey_erosion, + "run_cuda": run_cuda, + "torch_op": torch_op, "x": x, }, ).blocked_autorange(min_run_time=MIN_RUN_TIME) - scipy_ms = t_scipy.median * 1e3 / batch - torch_1x_ms = t_torch_1x.median * 1e3 / batch - torch_batch_ms = t_torch_batch.median * 1e3 / batch - - speedup1x = scipy_ms / torch_1x_ms - speedupbatch = scipy_ms / torch_batch_ms + scipy_ms = scipy_time.median * 1e3 / batch_size + torch_single_ms = torch_single_time.median * 1e3 / batch_size + torch_batch_ms = torch_batch_time.median * 1e3 / batch_size table.add_row( [ - size, + image_size, f"{scipy_ms:.3f}", - f"{torch_1x_ms:.3f}", + f"{torch_single_ms:.3f}", f"{torch_batch_ms:.3f}", - f"{speedup1x:.1f}x", - f"{speedupbatch:.1f}x", + f"{scipy_ms / torch_single_ms:.1f}x", + f"{scipy_ms / torch_batch_ms:.1f}x", ] ) - print(f"\n=========== Batch size : {batch} ===========") + + print(f"\n=========== Batch size : {batch_size} ===========") print(table) +def main(): + choices = (*GREY_OPERATORS, "all") + parser = argparse.ArgumentParser(description="Benchmark grey morphology operators.") + parser.add_argument( + "operation", + choices=choices, + nargs="?", + default="all", + help="Operator to benchmark (default: all).", + ) + args = parser.parse_args() + + operations = GREY_OPERATORS if args.operation == "all" else (args.operation,) + for operation in operations: + bench_grey_operator(operation) + + if __name__ == "__main__": - bench_grey_erosion() + main() diff --git a/pytest.ini b/pytest.ini index 8416814..86ff5f1 100644 --- a/pytest.ini +++ b/pytest.ini @@ -1,4 +1,5 @@ [pytest] markers = nd: ND tests - bfdt: brute force DT tests \ No newline at end of file + bfdt: brute force DT tests + cuda: CUDA tests \ No newline at end of file diff --git a/test/test_grey.py b/test/test_grey.py index ad0feda..03563f2 100644 --- a/test/test_grey.py +++ b/test/test_grey.py @@ -1,7 +1,11 @@ import numpy as np import pytest import torch -from scipy.ndimage import generate_binary_structure, grey_erosion +from scipy.ndimage import generate_binary_structure +from scipy.ndimage import grey_closing as scipy_grey_closing +from scipy.ndimage import grey_dilation as scipy_grey_dilation +from scipy.ndimage import grey_erosion as scipy_grey_erosion +from scipy.ndimage import grey_opening as scipy_grey_opening import torchmorph as tm @@ -9,153 +13,187 @@ not torch.cuda.is_available(), reason="CUDA is required for torchmorph tests" ) +GREY_OPERATORS = [ + pytest.param(tm.grey_erosion, scipy_grey_erosion, id="erosion"), + pytest.param(tm.grey_dilation, scipy_grey_dilation, id="dilation"), + pytest.param(tm.grey_opening, scipy_grey_opening, id="opening"), + pytest.param(tm.grey_closing, scipy_grey_closing, id="closing"), +] -def apply_scipy_to_batch(np_input, scipy_func, **kwargs): - input_shape = np_input.shape - batch_shape = input_shape[:2] - spatial_shape = input_shape[2:] - batch_size = int(np.prod(batch_shape)) - flatten_input = np_input.reshape(batch_size, *spatial_shape) +TORCH_OPERATORS = [ + pytest.param(tm.grey_erosion, id="erosion"), + pytest.param(tm.grey_dilation, id="dilation"), + pytest.param(tm.grey_opening, id="opening"), + pytest.param(tm.grey_closing, id="closing"), +] - output = kwargs.pop("output", None) - if output is not None and output.shape != spatial_shape: - raise ValueError(f"output shape must be {spatial_shape}, got {output.shape}") +CASE_2D = np.arange(1, 10, dtype=np.float32).reshape(1, 1, 3, 3) +CASE_3D = np.arange(2 * 1 * 4 * 4 * 4, dtype=np.float32).reshape(2, 1, 4, 4, 4) - results = [] - for sample in flatten_input: - sample_kwargs = dict(kwargs) - if output is not None: - sample_output = np.empty_like(output) - sample_kwargs["output"] = sample_output - result = scipy_func(sample, **sample_kwargs) - results.append(np.asarray(result).copy()) - return np.stack(results, axis=0).reshape(*batch_shape, *spatial_shape) +FOOTPRINT_CROSS = generate_binary_structure(rank=2, connectivity=1).astype(bool) +STRUCTURE_CROSS = FOOTPRINT_CROSS.astype(np.float32) +ASYMMETRIC_STRUCTURE = np.arange(9, dtype=np.float32).reshape(3, 3) -# test cases -case_2d = np.array( - [ - [ - [ - [1.0, 2.0, 3.0], - [4.0, 5.0, 6.0], - [7.0, 8.0, 9.0], - ] - ] - ], - dtype=np.float32, -) - -case_3d = np.arange(2 * 1 * 4 * 4 * 4, dtype=np.float32).reshape(2, 1, 4, 4, 4) +def apply_scipy_to_batch(np_input, scipy_op, **kwargs): + batch_shape = np_input.shape[:2] + spatial_shape = np_input.shape[2:] + samples = np_input.reshape(-1, *spatial_shape) + results = [scipy_op(sample, **kwargs) for sample in samples] + return np.stack(results).reshape(*batch_shape, *spatial_shape) -# structures -footprint_cross = generate_binary_structure(rank=2, connectivity=1).astype(bool) -structure_cross = footprint_cross.astype(np.float32) -footprint_3d = generate_binary_structure(rank=3, connectivity=2).astype(bool) - -# output tensors -output_2d_tm = torch.empty(1, 1, 3, 3) +def optional_cuda_tensor(value, dtype): + if value is None: + return None + return torch.as_tensor(value, dtype=dtype, device="cuda") +@pytest.mark.parametrize(("torch_op", "scipy_op"), GREY_OPERATORS) @pytest.mark.parametrize( ("np_input", "size", "footprint", "structure", "origin"), [ - pytest.param(case_2d, 3, None, None, 0, id="2d_size_int"), - pytest.param(case_2d, (3, 3), None, None, 0, id="2d_size_tuple"), - pytest.param(case_2d, None, footprint_cross, None, 0, id="2d_footprint"), - pytest.param(case_2d, None, None, structure_cross, 0, id="2d_structure"), + pytest.param(CASE_2D, 3, None, None, 0, id="2d_size_int"), + pytest.param(CASE_2D, (2, 3), None, None, 0, id="2d_even_size"), + pytest.param(CASE_2D, None, FOOTPRINT_CROSS, None, 0, id="2d_footprint"), + pytest.param( + CASE_2D, + None, + None, + ASYMMETRIC_STRUCTURE, + 1, + id="2d_asymmetric_structure_origin", + ), pytest.param( - case_2d, None, footprint_cross, structure_cross, 0, + CASE_2D, + None, + FOOTPRINT_CROSS, + STRUCTURE_CROSS, + 0, id="2d_structure_and_footprint", ), - pytest.param(case_2d, 3, None, None, 1, id="2d_size_origin_1"), - pytest.param(case_3d, 3, None, None, 0, id="3d_size"), - ] - ) - - -def test_grey_erosion_basic(np_input, size, footprint, structure, origin): - x = torch.as_tensor(np_input, dtype=torch.float32).cuda() - ft = ( - torch.as_tensor(footprint, dtype=torch.bool).cuda() - if footprint is not None - else None - ) - st = ( - torch.as_tensor(structure, dtype=torch.float32).cuda() - if structure is not None - else None - ) - result = tm.grey_erosion( - x, + pytest.param(CASE_3D, 3, None, None, 0, id="3d_size"), + ], +) +def test_grey_morphology_matches_scipy( + torch_op, scipy_op, np_input, size, footprint, structure, origin +): + result = torch_op( + torch.as_tensor(np_input, device="cuda"), size=size, - footprint=ft, - structure=st, + footprint=optional_cuda_tensor(footprint, torch.bool), + structure=optional_cuda_tensor(structure, torch.float32), origin=origin, ) - expected_np = apply_scipy_to_batch( + expected = apply_scipy_to_batch( np_input, - grey_erosion, + scipy_op, size=size, footprint=footprint, structure=structure, origin=origin, ) - expected = torch.as_tensor(expected_np) - torch.testing.assert_close(result.cpu(), expected) + torch.testing.assert_close(result.cpu(), torch.as_tensor(expected)) +@pytest.mark.parametrize(("torch_op", "scipy_op"), GREY_OPERATORS) @pytest.mark.parametrize( ("mode", "cval"), [ - pytest.param("reflect", 0.0, id="mode_reflect"), - pytest.param("constant", 0.0, id="mode_constant"), - pytest.param("constant", 99.0, id="mode_constant_cval99"), - pytest.param("nearest", 0.0, id="mode_nearest"), - pytest.param("mirror", 0.0, id="mode_mirror"), - pytest.param("wrap", 0.0, id="mode_wrap"), + pytest.param("reflect", 0.0, id="reflect"), + pytest.param("constant", 0.0, id="constant_zero"), + pytest.param("constant", 99.0, id="constant_positive"), + pytest.param("constant", -99.0, id="constant_negative"), + pytest.param("nearest", 0.0, id="nearest"), + pytest.param("mirror", 0.0, id="mirror"), + pytest.param("wrap", 0.0, id="wrap"), ], ) -def test_grey_erosion_modes(mode, cval): - x = torch.as_tensor(case_2d, dtype=torch.float32).cuda() - result = tm.grey_erosion(x, size=3, mode=mode, cval=cval) - expected_np = apply_scipy_to_batch( - case_2d, grey_erosion, size=3, mode=mode, cval=cval - ) - expected = torch.as_tensor(expected_np) - torch.testing.assert_close(result.cpu(), expected) +def test_grey_morphology_border_modes(torch_op, scipy_op, mode, cval): + x = torch.as_tensor(CASE_2D, device="cuda") + result = torch_op(x, size=3, mode=mode, cval=cval) + expected = apply_scipy_to_batch(CASE_2D, scipy_op, size=3, mode=mode, cval=cval) + torch.testing.assert_close(result.cpu(), torch.as_tensor(expected)) + +@pytest.mark.parametrize( + ("torch_op", "scipy_op", "value"), + [ + pytest.param(tm.grey_erosion, scipy_grey_erosion, np.inf, id="erosion_positive"), + pytest.param(tm.grey_dilation, scipy_grey_dilation, -np.inf, id="dilation_negative"), + ], +) +def test_grey_morphology_preserves_infinite_extrema(torch_op, scipy_op, value): + np_input = np.full((1, 1, 5, 5), value, dtype=np.float32) + result = torch_op(torch.as_tensor(np_input, device="cuda"), size=3) + expected = apply_scipy_to_batch(np_input, scipy_op, size=3) + torch.testing.assert_close(result.cpu(), torch.as_tensor(expected)) + + +@pytest.mark.parametrize(("torch_op", "scipy_op"), GREY_OPERATORS) +def test_grey_morphology_output(torch_op, scipy_op): + x = torch.as_tensor(CASE_2D, device="cuda") + output = torch.empty_like(x) + + result = torch_op(x, size=3, output=output) -def test_grey_erosion_output(): - x = torch.as_tensor(case_2d, dtype=torch.float32).cuda() - out = output_2d_tm.cuda() - result = tm.grey_erosion(x, size=3, output=out) - assert result is out - expected_np = apply_scipy_to_batch(case_2d, grey_erosion, size=3) - expected = torch.as_tensor(expected_np) - torch.testing.assert_close(result.cpu(), expected) + assert result is output + expected = apply_scipy_to_batch(CASE_2D, scipy_op, size=3) + torch.testing.assert_close(result.cpu(), torch.as_tensor(expected)) -def test_grey_erosion_no_params(): - x = torch.as_tensor(case_2d, dtype=torch.float32).cuda() +@pytest.mark.parametrize("torch_op", TORCH_OPERATORS) +def test_grey_morphology_requires_structuring_element(torch_op): + x = torch.as_tensor(CASE_2D, device="cuda") with pytest.raises(ValueError, match="At least one"): - tm.grey_erosion(x) + torch_op(x) -def test_grey_erosion_invalid_device(): - x = torch.tensor([[[1.0, 2.0], [3.0, 4.0]]]) +@pytest.mark.parametrize("torch_op", TORCH_OPERATORS) +def test_grey_morphology_requires_cuda(torch_op): + x = torch.as_tensor(CASE_2D) with pytest.raises(ValueError, match="CUDA"): - tm.grey_erosion(x, size=3) + torch_op(x, size=3) -def test_grey_erosion_invalid_ndim(): - x = torch.tensor([1.0, 2.0]).cuda() +@pytest.mark.parametrize("torch_op", TORCH_OPERATORS) +def test_grey_morphology_requires_spatial_dimensions(torch_op): + x = torch.tensor([1.0, 2.0], device="cuda") with pytest.raises(ValueError, match="at least 3 dimensions"): - tm.grey_erosion(x, size=3) + torch_op(x, size=3) -def test_grey_erosion_invalid_mode(): - x = torch.as_tensor(case_2d, dtype=torch.float32).cuda() +@pytest.mark.parametrize("torch_op", TORCH_OPERATORS) +def test_grey_morphology_rejects_unknown_mode(torch_op): + x = torch.as_tensor(CASE_2D, device="cuda") with pytest.raises(ValueError, match="Unknown mode"): - tm.grey_erosion(x, size=3, mode="invalid") + torch_op(x, size=3, mode="invalid") + + +@pytest.mark.parametrize("torch_op", TORCH_OPERATORS) +@pytest.mark.parametrize("origin", [-2, 2]) +def test_grey_morphology_rejects_invalid_origin(torch_op, origin): + x = torch.as_tensor(CASE_2D, device="cuda") + with pytest.raises(ValueError, match="invalid origin"): + torch_op(x, size=3, origin=origin) + + +@pytest.mark.parametrize(("torch_op", "scipy_op"), GREY_OPERATORS) +def test_grey_morphology_uses_current_cuda_stream(torch_op, scipy_op): + stream = torch.cuda.Stream() + with torch.cuda.stream(stream): + x = torch.arange(25, dtype=torch.float32, device="cuda").reshape(1, 1, 5, 5) + result = torch_op(x, size=3) + stream.synchronize() + + expected = apply_scipy_to_batch(x.cpu().numpy(), scipy_op, size=3) + torch.testing.assert_close(result.cpu(), torch.as_tensor(expected)) + + +@pytest.mark.skipif(torch.cuda.device_count() < 2, reason="requires two CUDA devices") +@pytest.mark.parametrize("torch_op", TORCH_OPERATORS) +def test_grey_morphology_uses_input_cuda_device(torch_op): + with torch.cuda.device(0): + x = torch.arange(25, dtype=torch.float32, device="cuda:1").reshape(1, 1, 5, 5) + result = torch_op(x, size=3) + assert result.device == x.device diff --git a/torchmorph/__init__.py b/torchmorph/__init__.py index 485b88e..e864ece 100644 --- a/torchmorph/__init__.py +++ b/torchmorph/__init__.py @@ -10,7 +10,10 @@ binary_erosion, binary_opening, generate_binary_structure, + grey_closing, + grey_dilation, grey_erosion, + grey_opening, ) __all__ = [ @@ -24,4 +27,7 @@ "binary_opening", "binary_closing", "grey_erosion", + "grey_dilation", + "grey_opening", + "grey_closing", ] diff --git a/torchmorph/csrc/greyero_kernel.cu b/torchmorph/csrc/greyero_kernel.cu index 51c92a2..5c2241e 100644 --- a/torchmorph/csrc/greyero_kernel.cu +++ b/torchmorph/csrc/greyero_kernel.cu @@ -1,18 +1,22 @@ #include +#include +#include +#include #include #include +#include #include #include #include #include -#define GREYERO_MAX_NDIM 8 +#define GREY_MORPH_MAX_NDIM 8 -struct GreyErosionGeometry { - int64_t spatial_size[GREYERO_MAX_NDIM]; - int64_t spatial_stride[GREYERO_MAX_NDIM]; - int64_t min_deltas[GREYERO_MAX_NDIM]; - int64_t max_deltas[GREYERO_MAX_NDIM]; +struct GreyMorphologyGeometry { + int64_t spatial_size[GREY_MORPH_MAX_NDIM]; + int64_t spatial_stride[GREY_MORPH_MAX_NDIM]; + int64_t min_deltas[GREY_MORPH_MAX_NDIM]; + int64_t max_deltas[GREY_MORPH_MAX_NDIM]; }; // Matches scipy.ndimage border modes. @@ -24,6 +28,38 @@ enum BorderMode { MODE_WRAP = 4, // Periodic wrap padding }; +struct GreyErosionOp { + __device__ __forceinline__ static float identity() { + return CUDART_INF_F; + } + + __device__ __forceinline__ static float combine( + float result, float value, float structure + ) { + return fminf(result, value - structure); + } + + __host__ __forceinline__ static int64_t offset(int64_t delta) { + return delta; + } +}; + +struct GreyDilationOp { + __device__ __forceinline__ static float identity() { + return -CUDART_INF_F; + } + + __device__ __forceinline__ static float combine( + float result, float value, float structure + ) { + return fmaxf(result, value + structure); + } + + __host__ __forceinline__ static int64_t offset(int64_t delta) { + return -delta; + } +}; + __device__ __forceinline__ int64_t map_coord( int64_t coord, int64_t size, BorderMode mode, bool* use_cval ) { @@ -60,16 +96,17 @@ __device__ __forceinline__ int64_t map_coord( } } -// Fused erosion kernel. struct_meta stores per-element dimension offsets +// Fused grey morphology kernel. struct_meta stores per-element dimension offsets // followed by the precomputed flat offset used by the interior fast path. -__global__ void grey_erosion_fused_kernel( +template +__global__ void grey_morphology_fused_kernel( const float* __restrict__ input, float* __restrict__ output, const float* __restrict__ struct_vals, const int64_t* __restrict__ struct_meta, const int num_struct, const int ndim_spatial, - const GreyErosionGeometry geom, + const GreyMorphologyGeometry geom, const int64_t total_spatial, const int64_t batch_channel, const BorderMode mode, @@ -81,14 +118,14 @@ __global__ void grey_erosion_fused_kernel( int64_t bc = idx / total_spatial; int64_t sp = idx % total_spatial; - int64_t coords[GREYERO_MAX_NDIM]; + int64_t coords[GREY_MORPH_MAX_NDIM]; int64_t rem = sp; for (int d = 0; d < ndim_spatial; d++) { coords[d] = rem / geom.spatial_stride[d]; rem %= geom.spatial_stride[d]; } - float min_val = FLT_MAX; + float result = MorphologyOp::identity(); const float* in_ptr = input + bc * total_spatial; const int meta_stride = ndim_spatial + 1; @@ -105,9 +142,9 @@ __global__ void grey_erosion_fused_kernel( for (int i = 0; i < num_struct; i++) { int64_t offset = struct_meta[i * meta_stride + ndim_spatial]; float val = in_ptr[sp + offset]; - min_val = fminf(min_val, val - struct_vals[i]); + result = MorphologyOp::combine(result, val, struct_vals[i]); } - output[idx] = min_val; + output[idx] = result; return; } @@ -126,13 +163,14 @@ __global__ void grey_erosion_fused_kernel( } float val = use_cval ? cval : in_ptr[flat_idx]; - min_val = fminf(min_val, val - struct_vals[i]); + result = MorphologyOp::combine(result, val, struct_vals[i]); } - output[idx] = min_val; + output[idx] = result; } -torch::Tensor grey_erosion_cuda( +template +static torch::Tensor grey_morphology_cuda( torch::Tensor input, torch::Tensor structure, torch::Tensor footprint, @@ -141,6 +179,8 @@ torch::Tensor grey_erosion_cuda( float cval ) { TORCH_CHECK(input.is_cuda(), "input must be a CUDA tensor"); + const c10::cuda::CUDAGuard device_guard(input.device()); + TORCH_CHECK(!structure.is_cuda(), "structure must be a CPU tensor"); TORCH_CHECK(!footprint.is_cuda(), "footprint must be a CPU tensor"); TORCH_CHECK(input.is_contiguous(), "input must be contiguous"); @@ -151,8 +191,8 @@ torch::Tensor grey_erosion_cuda( TORCH_CHECK(footprint.dtype() == torch::kBool, "footprint must be bool"); int ndim_spatial = structure.dim(); - TORCH_CHECK(ndim_spatial > 0 && ndim_spatial <= GREYERO_MAX_NDIM, - "structure dimension must be in 1-", GREYERO_MAX_NDIM, + TORCH_CHECK(ndim_spatial > 0 && ndim_spatial <= GREY_MORPH_MAX_NDIM, + "structure dimension must be in 1-", GREY_MORPH_MAX_NDIM, ", got ", ndim_spatial); TORCH_CHECK((int)origin_vec.size() == ndim_spatial, "origin length must match structure dimensions"); @@ -190,8 +230,8 @@ torch::Tensor grey_erosion_cuda( total_spatial *= h_spatial_size[d]; } - GreyErosionGeometry h_geom; - for (int d = 0; d < GREYERO_MAX_NDIM; d++) { + GreyMorphologyGeometry h_geom; + for (int d = 0; d < GREY_MORPH_MAX_NDIM; d++) { h_geom.spatial_size[d] = 0; h_geom.spatial_stride[d] = 0; h_geom.min_deltas[d] = 0; @@ -242,7 +282,7 @@ torch::Tensor grey_erosion_cuda( tmp %= h_structure_stride[d]; int64_t center_d = structure.size(d) / 2 + origin_vec[d]; - int64_t delta = coord_d - center_d; + int64_t delta = MorphologyOp::offset(coord_d - center_d); h_struct_meta.push_back(delta); flat_offset += delta * h_spatial_stride[d]; h_geom.min_deltas[d] = std::min(h_geom.min_deltas[d], delta); @@ -269,8 +309,9 @@ torch::Tensor grey_erosion_cuda( int64_t total_threads = batch_channel * total_spatial; int threads = 256; int blocks = (int)((total_threads + threads - 1) / threads); + const auto stream = c10::cuda::getCurrentCUDAStream(input.get_device()); - grey_erosion_fused_kernel<<>>( + grey_morphology_fused_kernel<<>>( input.data_ptr(), output.data_ptr(), d_struct_vals.data_ptr(), @@ -283,6 +324,33 @@ torch::Tensor grey_erosion_cuda( bmode, cval ); + C10_CUDA_KERNEL_LAUNCH_CHECK(); return output; } + +torch::Tensor grey_erosion_cuda( + torch::Tensor input, + torch::Tensor structure, + torch::Tensor footprint, + std::vector origin_vec, + int mode_int, + float cval +) { + return grey_morphology_cuda( + input, structure, footprint, origin_vec, mode_int, cval + ); +} + +torch::Tensor grey_dilation_cuda( + torch::Tensor input, + torch::Tensor structure, + torch::Tensor footprint, + std::vector origin_vec, + int mode_int, + float cval +) { + return grey_morphology_cuda( + input, structure, footprint, origin_vec, mode_int, cval + ); +} diff --git a/torchmorph/csrc/torchmorph.cpp b/torchmorph/csrc/torchmorph.cpp index 353d489..f87a1d6 100644 --- a/torchmorph/csrc/torchmorph.cpp +++ b/torchmorph/csrc/torchmorph.cpp @@ -11,6 +11,15 @@ torch::Tensor grey_erosion_cuda( float cval ); +torch::Tensor grey_dilation_cuda( + torch::Tensor input, + torch::Tensor structure, + torch::Tensor footprint, + std::vector origin, + int mode, + float cval +); + std::tuple edt_cuda( torch::Tensor input, std::vector sampling, @@ -47,6 +56,15 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { py::arg("mode"), py::arg("cval")); + m.def("grey_dilation_cuda", &grey_dilation_cuda, + "N-dimensional fused grey dilation", + py::arg("input"), + py::arg("structure"), + py::arg("footprint"), + py::arg("origin"), + py::arg("mode"), + py::arg("cval")); + m.def("edt_cuda", &edt_cuda, "Exact Euclidean Distance Transform (Felzenszwalb algorithm)", py::arg("input"), diff --git a/torchmorph/morphology/__init__.py b/torchmorph/morphology/__init__.py index 5819ec3..75bef93 100644 --- a/torchmorph/morphology/__init__.py +++ b/torchmorph/morphology/__init__.py @@ -1,5 +1,5 @@ from .binary import binary_closing, binary_dilation, binary_erosion, binary_opening -from .grey import grey_erosion +from .grey import grey_closing, grey_dilation, grey_erosion, grey_opening from .structure import generate_binary_structure __all__ = [ @@ -9,4 +9,7 @@ "binary_opening", "binary_closing", "grey_erosion", + "grey_dilation", + "grey_opening", + "grey_closing", ] diff --git a/torchmorph/morphology/grey.py b/torchmorph/morphology/grey.py index bc1757a..b6f1cda 100644 --- a/torchmorph/morphology/grey.py +++ b/torchmorph/morphology/grey.py @@ -12,7 +12,7 @@ } -def grey_erosion( +def _grey_morphology( input: Tensor, size: int | tuple[int, ...] | None = None, footprint: Tensor | None = None, @@ -21,49 +21,18 @@ def grey_erosion( mode: str = 'reflect', cval: float = 0.0, origin: int | tuple[int, ...] = 0, + *, + operation: str, ) -> Tensor: - """N-dimensional grey erosion for (B, C, Spatial...) tensors. - - For a single image or volume, add batch and channel dimensions first. - - The structuring element is resolved with the following priority: - 1. `structure` (non-flat, values are subtracted) - 2. `footprint` (flat, only positions where footprint > 0 are used) - 3. `size` (flat rectangle shaped by `size`) - Exactly one of these three parameters must be provided. - - Args: - input: Input tensor in (B, C, Spatial...) format. Must be on CUDA. - size: Shape of a flat rectangular structuring element. Mutually - exclusive with footprint and structure. - footprint: Binary structuring element. Non-zero positions define - the neighborhood over which the minimum is computed. Mutually - exclusive with size and structure. - structure: Value structuring element. Values are subtracted from - each pixel in the neighborhood. Mutually exclusive with size - and footprint. When given, `footprint` may optionally be - provided as a mask to restrict the neighborhood. - output: Pre-allocated output tensor. If provided, the result is - written in-place and returned. - mode: Border mode. One of 'constant', 'reflect', 'nearest', - 'mirror', 'wrap'. Default is 'reflect'. - cval: Constant fill value used when mode='constant'. Default is 0.0. - origin: Origin offset. Can be an int (same for all axes) or tuple. - - Returns: - Grey-eroded tensor with the same shape as input. - """ + """Shared argument normalization for grey erosion and dilation.""" if not input.is_cuda: raise ValueError('Input tensor must be on CUDA device.') if input.ndim < 3: raise ValueError( - f'Input must be (B, C, Spatial...) with at least 3 dimensions, ' - f'got {input.shape}.' + f'Input must be (B, C, Spatial...) with at least 3 dimensions, ' f'got {input.shape}.' ) if input.numel() == 0: - raise ValueError( - f'Invalid input: empty tensor with shape {input.shape}.' - ) + raise ValueError(f'Invalid input: empty tensor with shape {input.shape}.') spatial_ndim = input.ndim - 2 @@ -81,30 +50,23 @@ def grey_erosion( ) if not mask.any(): raise ValueError('All-zero footprint is not supported.') - footprint_cpu = ( - mask.detach().to(device='cpu', dtype=torch.bool).contiguous() - ) + footprint_cpu = mask.detach().to(device='cpu', dtype=torch.bool).contiguous() elif footprint is not None: mask = footprint.to(torch.bool) if not mask.any(): raise ValueError('All-zero footprint is not supported.') struct = torch.zeros(tuple(mask.shape), dtype=torch.float32) - footprint_cpu = ( - mask.detach().to(device='cpu', dtype=torch.bool).contiguous() - ) + footprint_cpu = mask.detach().to(device='cpu', dtype=torch.bool).contiguous() elif size is not None: if isinstance(size, int): size = (size,) * spatial_ndim struct = torch.zeros(size, dtype=torch.float32) else: - raise ValueError( - 'At least one of size, footprint, or structure must be specified.' - ) + raise ValueError('At least one of size, footprint, or structure must be specified.') if struct.ndim != spatial_ndim: raise ValueError( - f'Structure dimension {struct.ndim} must match spatial dimension ' - f'{spatial_ndim}.' + f'Structure dimension {struct.ndim} must match spatial dimension ' f'{spatial_ndim}.' ) struct = struct.detach().to(device='cpu', dtype=torch.float32).contiguous() @@ -115,19 +77,23 @@ def grey_erosion( origin_list = list(origin) if len(origin_list) != spatial_ndim: raise ValueError( - f'Origin length {len(origin_list)} must match spatial dimension ' - f'{spatial_ndim}.' + f'Origin length {len(origin_list)} must match spatial dimension ' f'{spatial_ndim}.' ) + for origin_value, structure_size in zip(origin_list, struct.shape): + min_origin = -(structure_size // 2) + max_origin = (structure_size - 1) // 2 + if not min_origin <= origin_value <= max_origin: + raise ValueError("invalid origin") + # Map mode string to int if mode not in _MODE_MAP: - raise ValueError( - f"Unknown mode '{mode}'. Must be one of {list(_MODE_MAP.keys())}." - ) + raise ValueError(f"Unknown mode '{mode}'. Must be one of {list(_MODE_MAP.keys())}.") mode_int = _MODE_MAP[mode] # CUDA kernel call - result = _C.grey_erosion_cuda( + kernel = _C.grey_erosion_cuda if operation == "erosion" else _C.grey_dilation_cuda + result = kernel( input.contiguous().float(), struct.contiguous(), footprint_cpu, @@ -140,3 +106,127 @@ def grey_erosion( output.copy_(result) return output return result + + +def grey_erosion( + input: Tensor, + size: int | tuple[int, ...] | None = None, + footprint: Tensor | None = None, + structure: Tensor | None = None, + output: Tensor | None = None, + mode: str = "reflect", + 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. + """ + return _grey_morphology( + input, + size, + footprint, + structure, + output, + mode, + cval, + origin, + operation="erosion", + ) + + +def grey_dilation( + input: Tensor, + size: int | tuple[int, ...] | None = None, + footprint: Tensor | None = None, + structure: Tensor | None = None, + output: Tensor | None = None, + mode: str = "reflect", + 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. + """ + return _grey_morphology( + input, + size, + footprint, + structure, + output, + mode, + cval, + origin, + operation="dilation", + ) + + +def grey_opening( + input: Tensor, + size: int | tuple[int, ...] | None = None, + footprint: Tensor | None = None, + structure: Tensor | None = None, + output: Tensor | None = None, + mode: str = "reflect", + 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. + """ + eroded = grey_erosion( + input, + size=size, + footprint=footprint, + structure=structure, + mode=mode, + cval=cval, + origin=origin, + ) + return grey_dilation( + eroded, + size=size, + footprint=footprint, + structure=structure, + output=output, + mode=mode, + cval=cval, + origin=origin, + ) + + +def grey_closing( + input: Tensor, + size: int | tuple[int, ...] | None = None, + footprint: Tensor | None = None, + structure: Tensor | None = None, + output: Tensor | None = None, + mode: str = "reflect", + 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. + """ + dilated = grey_dilation( + input, + size=size, + footprint=footprint, + structure=structure, + mode=mode, + cval=cval, + origin=origin, + ) + return grey_erosion( + dilated, + size=size, + footprint=footprint, + structure=structure, + output=output, + mode=mode, + cval=cval, + origin=origin, + ) From 3785a6277fb8fb3e46583da52bd76891a00489df Mon Sep 17 00:00:00 2001 From: Yuhandeng Date: Tue, 7 Jul 2026 14:21:45 +0800 Subject: [PATCH 04/25] rename --- torchmorph/csrc/{greyero_kernel.cu => grey_kernel.cu} | 0 1 file changed, 0 insertions(+), 0 deletions(-) rename torchmorph/csrc/{greyero_kernel.cu => grey_kernel.cu} (100%) diff --git a/torchmorph/csrc/greyero_kernel.cu b/torchmorph/csrc/grey_kernel.cu similarity index 100% rename from torchmorph/csrc/greyero_kernel.cu rename to torchmorph/csrc/grey_kernel.cu From 98964a5e0d7b03d3ad9f47bd62f4f0b78c510422 Mon Sep 17 00:00:00 2001 From: Yuhandeng Date: Tue, 7 Jul 2026 14:41:48 +0800 Subject: [PATCH 05/25] gradient black/white tophat --- benchmark/grey.py | 63 +++++++++++++++++----- test/test_grey.py | 9 ++++ torchmorph/__init__.py | 6 +++ torchmorph/morphology/__init__.py | 13 ++++- torchmorph/morphology/grey.py | 90 +++++++++++++++++++++++++++++++ 5 files changed, 167 insertions(+), 14 deletions(-) diff --git a/benchmark/grey.py b/benchmark/grey.py index a53bf75..89d1a67 100644 --- a/benchmark/grey.py +++ b/benchmark/grey.py @@ -17,23 +17,26 @@ "dilation": (ndi.grey_dilation, tm.grey_dilation), "opening": (ndi.grey_opening, tm.grey_opening), "closing": (ndi.grey_closing, tm.grey_closing), + "gradient": (ndi.morphological_gradient, tm.morphological_gradient), + "white_tophat": (ndi.white_tophat, tm.white_tophat), + "black_tophat": (ndi.black_tophat, tm.black_tophat), } -def run_cuda(torch_op, x): - result = torch_op(x, size=STRUCTURE_SIZE) +def run_cuda(torch_op, x, structure_size): + result = torch_op(x, size=structure_size) torch.cuda.synchronize() return result -def bench_grey_operator(operation): +def bench_grey_operator(operation, image_sizes, batch_sizes, structure_size, min_run_time): scipy_op, torch_op = GREY_OPERATORS[operation] print("\n============================================") print(f" Benchmark: grey {operation} ") print("============================================") - for batch_size in BATCH_SIZES: + for batch_size in batch_sizes: table = PrettyTable() table.field_names = [ "Size", @@ -46,7 +49,7 @@ def bench_grey_operator(operation): for column in table.field_names: table.align[column] = "r" - for image_size in IMAGE_SIZES: + for image_size in image_sizes: x = torch.randn( batch_size, 1, @@ -62,32 +65,34 @@ def bench_grey_operator(operation): stmt="[scipy_op(data, size=structure_size) for data in inputs]", globals={ "scipy_op": scipy_op, - "structure_size": STRUCTURE_SIZE, + "structure_size": structure_size, "inputs": x_np_list, }, - ).blocked_autorange(min_run_time=MIN_RUN_TIME) + ).blocked_autorange(min_run_time=min_run_time) for data in x_one: - torch_op(data, size=STRUCTURE_SIZE) + torch_op(data, size=structure_size) torch.cuda.synchronize() torch_single_time = benchmark.Timer( - stmt="[run_cuda(torch_op, data) for data in inputs]", + stmt="[run_cuda(torch_op, data, structure_size) for data in inputs]", globals={ "run_cuda": run_cuda, "torch_op": torch_op, "inputs": x_one, + "structure_size": structure_size, }, - ).blocked_autorange(min_run_time=MIN_RUN_TIME) + ).blocked_autorange(min_run_time=min_run_time) torch_batch_time = benchmark.Timer( - stmt="run_cuda(torch_op, x)", + stmt="run_cuda(torch_op, x, structure_size)", globals={ "run_cuda": run_cuda, "torch_op": torch_op, "x": x, + "structure_size": structure_size, }, - ).blocked_autorange(min_run_time=MIN_RUN_TIME) + ).blocked_autorange(min_run_time=min_run_time) scipy_ms = scipy_time.median * 1e3 / batch_size torch_single_ms = torch_single_time.median * 1e3 / batch_size @@ -118,11 +123,43 @@ def main(): default="all", help="Operator to benchmark (default: all).", ) + parser.add_argument( + "--sizes", + type=int, + nargs="+", + default=IMAGE_SIZES, + help="Image sizes to benchmark (default: 64 128 256 1024).", + ) + parser.add_argument( + "--batches", + type=int, + nargs="+", + default=BATCH_SIZES, + help="Batch sizes to benchmark (default: 1 2 4 16).", + ) + parser.add_argument( + "--structure-size", + type=int, + default=STRUCTURE_SIZE, + help="Square structuring element size (default: 3).", + ) + parser.add_argument( + "--min-run-time", + type=float, + default=MIN_RUN_TIME, + help="Minimum benchmark time per timer in seconds (default: 1.0).", + ) args = parser.parse_args() operations = GREY_OPERATORS if args.operation == "all" else (args.operation,) for operation in operations: - bench_grey_operator(operation) + bench_grey_operator( + operation, + image_sizes=args.sizes, + batch_sizes=args.batches, + structure_size=args.structure_size, + min_run_time=args.min_run_time, + ) if __name__ == "__main__": diff --git a/test/test_grey.py b/test/test_grey.py index 03563f2..6cadfd3 100644 --- a/test/test_grey.py +++ b/test/test_grey.py @@ -1,11 +1,14 @@ import numpy as np import pytest import torch +from scipy.ndimage import black_tophat as scipy_black_tophat from scipy.ndimage import generate_binary_structure from scipy.ndimage import grey_closing as scipy_grey_closing from scipy.ndimage import grey_dilation as scipy_grey_dilation from scipy.ndimage import grey_erosion as scipy_grey_erosion from scipy.ndimage import grey_opening as scipy_grey_opening +from scipy.ndimage import morphological_gradient as scipy_morphological_gradient +from scipy.ndimage import white_tophat as scipy_white_tophat import torchmorph as tm @@ -18,6 +21,9 @@ pytest.param(tm.grey_dilation, scipy_grey_dilation, id="dilation"), pytest.param(tm.grey_opening, scipy_grey_opening, id="opening"), pytest.param(tm.grey_closing, scipy_grey_closing, id="closing"), + pytest.param(tm.morphological_gradient, scipy_morphological_gradient, id="gradient"), + pytest.param(tm.white_tophat, scipy_white_tophat, id="white_tophat"), + pytest.param(tm.black_tophat, scipy_black_tophat, id="black_tophat"), ] TORCH_OPERATORS = [ @@ -25,6 +31,9 @@ pytest.param(tm.grey_dilation, id="dilation"), pytest.param(tm.grey_opening, id="opening"), pytest.param(tm.grey_closing, id="closing"), + pytest.param(tm.morphological_gradient, id="gradient"), + pytest.param(tm.white_tophat, id="white_tophat"), + pytest.param(tm.black_tophat, id="black_tophat"), ] CASE_2D = np.arange(1, 10, dtype=np.float32).reshape(1, 1, 3, 3) diff --git a/torchmorph/__init__.py b/torchmorph/__init__.py index e864ece..8463d17 100644 --- a/torchmorph/__init__.py +++ b/torchmorph/__init__.py @@ -9,11 +9,14 @@ binary_dilation, binary_erosion, binary_opening, + black_tophat, generate_binary_structure, grey_closing, grey_dilation, grey_erosion, grey_opening, + morphological_gradient, + white_tophat, ) __all__ = [ @@ -30,4 +33,7 @@ "grey_dilation", "grey_opening", "grey_closing", + "morphological_gradient", + "white_tophat", + "black_tophat", ] diff --git a/torchmorph/morphology/__init__.py b/torchmorph/morphology/__init__.py index 75bef93..dc2a5f5 100644 --- a/torchmorph/morphology/__init__.py +++ b/torchmorph/morphology/__init__.py @@ -1,5 +1,13 @@ from .binary import binary_closing, binary_dilation, binary_erosion, binary_opening -from .grey import grey_closing, grey_dilation, grey_erosion, grey_opening +from .grey import ( + black_tophat, + grey_closing, + grey_dilation, + grey_erosion, + grey_opening, + morphological_gradient, + white_tophat, +) from .structure import generate_binary_structure __all__ = [ @@ -12,4 +20,7 @@ "grey_dilation", "grey_opening", "grey_closing", + "morphological_gradient", + "white_tophat", + "black_tophat", ] diff --git a/torchmorph/morphology/grey.py b/torchmorph/morphology/grey.py index b6f1cda..f729ac3 100644 --- a/torchmorph/morphology/grey.py +++ b/torchmorph/morphology/grey.py @@ -230,3 +230,93 @@ def grey_closing( cval=cval, origin=origin, ) + + +def morphological_gradient( + input: Tensor, + size: int | tuple[int, ...] | None = None, + footprint: Tensor | None = None, + structure: Tensor | None = None, + output: Tensor | None = None, + mode: str = "reflect", + cval: float = 0.0, + origin: int | tuple[int, ...] = 0, +) -> Tensor: + """N-dimensional morphological gradient for ``(B, C, Spatial...)`` CUDA tensors.""" + dilated = grey_dilation( + input, + size=size, + footprint=footprint, + structure=structure, + mode=mode, + cval=cval, + origin=origin, + ) + eroded = grey_erosion( + input, + size=size, + footprint=footprint, + structure=structure, + mode=mode, + cval=cval, + origin=origin, + ) + result = dilated - eroded + if output is not None: + output.copy_(result) + return output + return result + + +def white_tophat( + input: Tensor, + size: int | tuple[int, ...] | None = None, + footprint: Tensor | None = None, + structure: Tensor | None = None, + output: Tensor | None = None, + mode: str = "reflect", + cval: float = 0.0, + origin: int | tuple[int, ...] = 0, +) -> Tensor: + """N-dimensional white top-hat filter for ``(B, C, Spatial...)`` CUDA tensors.""" + opened = grey_opening( + input, + size=size, + footprint=footprint, + structure=structure, + mode=mode, + cval=cval, + origin=origin, + ) + result = input.float() - opened + if output is not None: + output.copy_(result) + return output + return result + + +def black_tophat( + input: Tensor, + size: int | tuple[int, ...] | None = None, + footprint: Tensor | None = None, + structure: Tensor | None = None, + output: Tensor | None = None, + mode: str = "reflect", + cval: float = 0.0, + origin: int | tuple[int, ...] = 0, +) -> Tensor: + """N-dimensional black top-hat filter for ``(B, C, Spatial...)`` CUDA tensors.""" + closed = grey_closing( + input, + size=size, + footprint=footprint, + structure=structure, + mode=mode, + cval=cval, + origin=origin, + ) + result = closed - input.float() + if output is not None: + output.copy_(result) + return output + return result From 4f70eab93d18685cbb157765fe53859c6479f5b6 Mon Sep 17 00:00:00 2001 From: Yuhandeng Date: Tue, 7 Jul 2026 15:01:10 +0800 Subject: [PATCH 06/25] iterate structure --- benchmark/generate.py | 59 ----------- benchmark/structure.py | 154 +++++++++++++++++++++++++++++ test/test_structure.py | 49 +++++++++ torchmorph/__init__.py | 2 + torchmorph/morphology/__init__.py | 3 +- torchmorph/morphology/structure.py | 73 ++++++++++++++ 6 files changed, 280 insertions(+), 60 deletions(-) delete mode 100644 benchmark/generate.py create mode 100644 benchmark/structure.py diff --git a/benchmark/generate.py b/benchmark/generate.py deleted file mode 100644 index aebddd8..0000000 --- a/benchmark/generate.py +++ /dev/null @@ -1,59 +0,0 @@ -import scipy.ndimage as ndi -import torch.utils.benchmark as benchmark -from prettytable import PrettyTable - -import torchmorph as tm - -DATA_CASES = [(1, 1), (2, 1), (4, 1), (8, 1), (10, 1), (10, 4), (10, 10), (12, 1)] -MIN_RUN = 1.0 - - -def bench_generate_binary_structure(): - print("\n========================================") - print("\n Benchmark: generate binary structure ") - print("\n========================================") - - table = PrettyTable() - table.field_names = ["Rank", "Connectivity", "SciPy(ms)", "Torch(ms)", "Speedup"] - for column in table.field_names: - table.align[column] = "r" - - for rank, connectivity in DATA_CASES: - t_scipy = benchmark.Timer( - stmt="generate_binary_structure(rank, connectivity)", - globals={ - "generate_binary_structure": ndi.generate_binary_structure, - "rank": rank, - "connectivity": connectivity, - }, - ).blocked_autorange(min_run_time=MIN_RUN) - - t_torch = benchmark.Timer( - stmt="generate_binary_structure(rank, connectivity)", - globals={ - "generate_binary_structure": tm.generate_binary_structure, - "rank": rank, - "connectivity": connectivity, - }, - ).blocked_autorange(min_run_time=MIN_RUN) - - scipy_ms = t_scipy.median * 1e3 - torch_ms = t_torch.median * 1e3 - speedup = scipy_ms / torch_ms - - table.add_row( - [ - rank, - connectivity, - f"{scipy_ms:.3f}", - f"{torch_ms:.3f}", - f"{speedup:.1f}x", - ] - ) - - print("Load from:", tm.__file__) - print(table) - - -if __name__ == "__main__": - bench_generate_binary_structure() diff --git a/benchmark/structure.py b/benchmark/structure.py new file mode 100644 index 0000000..30d3d25 --- /dev/null +++ b/benchmark/structure.py @@ -0,0 +1,154 @@ +import argparse + +import scipy.ndimage as ndi +import torch +import torch.utils.benchmark as benchmark +from prettytable import PrettyTable + +import torchmorph as tm + +GENERATE_CASES = [(1, 1), (2, 1), (4, 1), (8, 1), (10, 1), (10, 4), (10, 10), (12, 1)] +ITERATE_CASES = [ + (1, 1, 3), + (2, 1, 2), + (2, 2, 3), + (3, 1, 2), + (3, 2, 3), + (4, 1, 2), + (7, 2, 2), +] +MIN_RUN_TIME = 1.0 + + +def bench_generate_binary_structure(min_run_time): + print("\n========================================") + print(" Benchmark: generate binary structure ") + print("========================================") + + table = PrettyTable() + table.field_names = ["Rank", "Connectivity", "SciPy(ms)", "Torch(ms)", "Speedup"] + for column in table.field_names: + table.align[column] = "r" + + for rank, connectivity in GENERATE_CASES: + t_scipy = benchmark.Timer( + stmt="generate_binary_structure(rank, connectivity)", + globals={ + "generate_binary_structure": ndi.generate_binary_structure, + "rank": rank, + "connectivity": connectivity, + }, + ).blocked_autorange(min_run_time=min_run_time) + + t_torch = benchmark.Timer( + stmt="generate_binary_structure(rank, connectivity)", + globals={ + "generate_binary_structure": tm.generate_binary_structure, + "rank": rank, + "connectivity": connectivity, + }, + ).blocked_autorange(min_run_time=min_run_time) + + scipy_ms = t_scipy.median * 1e3 + torch_ms = t_torch.median * 1e3 + + table.add_row( + [ + rank, + connectivity, + f"{scipy_ms:.3f}", + f"{torch_ms:.3f}", + f"{scipy_ms / torch_ms:.1f}x", + ] + ) + + print(table) + + +def bench_iterate_structure(min_run_time): + print("\n========================================") + print(" Benchmark: iterate structure ") + print("========================================") + + table = PrettyTable() + table.field_names = [ + "Rank", + "Connectivity", + "Iterations", + "Output", + "True", + "SciPy(ms)", + "Torch(ms)", + "Speedup", + ] + for column in table.field_names: + table.align[column] = "r" + + for rank, connectivity, iterations in ITERATE_CASES: + scipy_structure = ndi.generate_binary_structure(rank, connectivity) + torch_structure = torch.as_tensor(scipy_structure) + expected = ndi.iterate_structure(scipy_structure, iterations) + + t_scipy = benchmark.Timer( + stmt="iterate_structure(structure, iterations)", + globals={ + "iterate_structure": ndi.iterate_structure, + "structure": scipy_structure, + "iterations": iterations, + }, + ).blocked_autorange(min_run_time=min_run_time) + + t_torch = benchmark.Timer( + stmt="iterate_structure(structure, iterations)", + globals={ + "iterate_structure": tm.iterate_structure, + "structure": torch_structure, + "iterations": iterations, + }, + ).blocked_autorange(min_run_time=min_run_time) + + scipy_ms = t_scipy.median * 1e3 + torch_ms = t_torch.median * 1e3 + + table.add_row( + [ + rank, + connectivity, + iterations, + "x".join(str(size) for size in expected.shape), + int(expected.sum()), + f"{scipy_ms:.3f}", + f"{torch_ms:.3f}", + f"{scipy_ms / torch_ms:.1f}x", + ] + ) + + print(table) + + +def main(): + parser = argparse.ArgumentParser(description="Benchmark binary structure helpers.") + parser.add_argument( + "operation", + choices=("generate", "iterate", "all"), + nargs="?", + default="all", + help="Structure helper to benchmark (default: all).", + ) + parser.add_argument( + "--min-run-time", + type=float, + default=MIN_RUN_TIME, + help="Minimum benchmark time per timer in seconds (default: 1.0).", + ) + args = parser.parse_args() + + print("Load from:", tm.__file__) + if args.operation in ("generate", "all"): + bench_generate_binary_structure(args.min_run_time) + if args.operation in ("iterate", "all"): + bench_iterate_structure(args.min_run_time) + + +if __name__ == "__main__": + main() diff --git a/test/test_structure.py b/test/test_structure.py index 55726ce..442bc0a 100644 --- a/test/test_structure.py +++ b/test/test_structure.py @@ -2,6 +2,7 @@ import pytest import torch from scipy.ndimage import generate_binary_structure as scipy_generate_binary_structure +from scipy.ndimage import iterate_structure as scipy_iterate_structure import torchmorph as tm # noqa: F401 @@ -26,3 +27,51 @@ def test_generate_binary_structure(rank, connectivity): expected_tensor = torch.as_tensor(expected) torch.testing.assert_close(actual, expected_tensor) + + +@pytest.mark.parametrize( + "structure, iterations", + [ + pytest.param(scipy_generate_binary_structure(1, 1), 3, id="1d_cross_3"), + pytest.param(scipy_generate_binary_structure(2, 1), 1, id="2d_cross_1"), + pytest.param(scipy_generate_binary_structure(2, 1), 2, id="2d_cross_2"), + pytest.param(scipy_generate_binary_structure(2, 2), 3, id="2d_full_3"), + pytest.param(scipy_generate_binary_structure(3, 1), 2, id="3d_cross_2"), + pytest.param(np.array([[False, True], [True, True]], dtype=bool), 3, id="asymmetric"), + pytest.param(np.zeros((3, 3), dtype=bool), 2, id="empty"), + ], +) +def test_iterate_structure(structure, iterations): + expected = scipy_iterate_structure(structure, iterations) + actual = tm.iterate_structure(torch.as_tensor(structure), iterations) + + np.testing.assert_array_equal(actual, expected) + torch.testing.assert_close(actual, torch.as_tensor(expected)) + + +@pytest.mark.parametrize( + "origin", + [ + pytest.param(0, id="scalar_zero"), + pytest.param((1, 0), id="tuple_positive"), + pytest.param((-1, 1), id="tuple_mixed"), + ], +) +def test_iterate_structure_with_origin(origin): + structure = scipy_generate_binary_structure(2, 1) + + expected_structure, expected_origin = scipy_iterate_structure(structure, 3, origin=origin) + actual_structure, actual_origin = tm.iterate_structure( + torch.as_tensor(structure), 3, origin=origin + ) + + np.testing.assert_array_equal(actual_structure, expected_structure) + torch.testing.assert_close(actual_structure, torch.as_tensor(expected_structure)) + assert actual_origin == expected_origin + + +def test_iterate_structure_rejects_invalid_origin_dimension(): + structure = torch.ones((3, 3), dtype=torch.bool) + + with pytest.raises(ValueError, match="origin dimension"): + tm.iterate_structure(structure, 2, origin=(0, 0, 0)) diff --git a/torchmorph/__init__.py b/torchmorph/__init__.py index 8463d17..7a39565 100644 --- a/torchmorph/__init__.py +++ b/torchmorph/__init__.py @@ -15,6 +15,7 @@ grey_dilation, grey_erosion, grey_opening, + iterate_structure, morphological_gradient, white_tophat, ) @@ -25,6 +26,7 @@ "chamfer_distance_transform", "brute_force_distance_transform", "generate_binary_structure", + "iterate_structure", "binary_erosion", "binary_dilation", "binary_opening", diff --git a/torchmorph/morphology/__init__.py b/torchmorph/morphology/__init__.py index dc2a5f5..b3a9009 100644 --- a/torchmorph/morphology/__init__.py +++ b/torchmorph/morphology/__init__.py @@ -8,10 +8,11 @@ morphological_gradient, white_tophat, ) -from .structure import generate_binary_structure +from .structure import generate_binary_structure, iterate_structure __all__ = [ "generate_binary_structure", + "iterate_structure", "binary_erosion", "binary_dilation", "binary_opening", diff --git a/torchmorph/morphology/structure.py b/torchmorph/morphology/structure.py index 497c42a..573f29b 100644 --- a/torchmorph/morphology/structure.py +++ b/torchmorph/morphology/structure.py @@ -14,3 +14,76 @@ def generate_binary_structure(rank: int, connectivity: int) -> Tensor: grids = torch.meshgrid(*axes, indexing="ij") offsets = torch.stack(grids, dim=0) return (offsets != 0).sum(dim=0) <= connectivity + + +def _generated_structure_connectivity(structure: Tensor) -> int | None: + if any(size != 3 for size in structure.shape): + return None + + rank = structure.ndim + axes = [torch.tensor([-1, 0, 1], device=structure.device) for _ in range(rank)] + grids = torch.meshgrid(*axes, indexing="ij") + offsets = torch.stack(grids, dim=0) + distances = (offsets != 0).sum(dim=0) + + for connectivity in range(1, rank + 1): + if torch.equal(structure, distances <= connectivity): + return connectivity + return None + + +def _iterate_generated_structure(rank: int, connectivity: int, iterations: int, device) -> Tensor: + axes = [torch.arange(-iterations, iterations + 1, device=device) for _ in range(rank)] + grids = torch.meshgrid(*axes, indexing="ij") + offsets = torch.stack(grids, dim=0) + return offsets.abs().sum(dim=0) <= connectivity * iterations + + +def _prepare_origin(origin: int | tuple[int, ...], ndim: int) -> list[int]: + if isinstance(origin, int): + return [origin] * ndim + + origin_list = list(origin) + if len(origin_list) != ndim: + raise ValueError(f"origin dimension is not {ndim}, got {len(origin_list)}") + return origin_list + + +def iterate_structure( + structure: Tensor, + iterations: int, + origin: int | tuple[int, ...] | None = None, +) -> Tensor | tuple[Tensor, list[int]]: + """Iterate a binary structure by dilating it with itself.""" + structure = structure != 0 + + if iterations < 2: + result = structure.clone() + else: + connectivity = _generated_structure_connectivity(structure) + if connectivity is not None: + result = _iterate_generated_structure( + structure.ndim, connectivity, iterations, structure.device + ) + else: + coordinates = torch.nonzero(structure, as_tuple=False) + if coordinates.numel() == 0: + result_shape = tuple(iterations * (size - 1) + 1 for size in structure.shape) + result = torch.zeros(result_shape, dtype=torch.bool, device=structure.device) + else: + current = coordinates + for _ in range(iterations - 1): + current = (current[:, None, :] + coordinates[None, :, :]).reshape( + -1, structure.ndim + ) + current = torch.unique(current, dim=0) + + result_shape = tuple(iterations * (size - 1) + 1 for size in structure.shape) + result = torch.zeros(result_shape, dtype=torch.bool, device=structure.device) + result[tuple(current.unbind(dim=1))] = True + + if origin is None: + return result + + origin_list = _prepare_origin(origin, structure.ndim) + return result, [value * iterations for value in origin_list] From 01ae5f565f5922b614c434fad9ab4ad8ae4351c1 Mon Sep 17 00:00:00 2001 From: Yuhandeng Date: Tue, 7 Jul 2026 15:16:03 +0800 Subject: [PATCH 07/25] grey laplace --- benchmark/grey.py | 1 + test/test_grey.py | 3 +++ torchmorph/__init__.py | 2 ++ torchmorph/morphology/__init__.py | 2 ++ torchmorph/morphology/grey.py | 36 +++++++++++++++++++++++++++++++ 5 files changed, 44 insertions(+) diff --git a/benchmark/grey.py b/benchmark/grey.py index 89d1a67..f413631 100644 --- a/benchmark/grey.py +++ b/benchmark/grey.py @@ -18,6 +18,7 @@ "opening": (ndi.grey_opening, tm.grey_opening), "closing": (ndi.grey_closing, tm.grey_closing), "gradient": (ndi.morphological_gradient, tm.morphological_gradient), + "laplace": (ndi.morphological_laplace, tm.morphological_laplace), "white_tophat": (ndi.white_tophat, tm.white_tophat), "black_tophat": (ndi.black_tophat, tm.black_tophat), } diff --git a/test/test_grey.py b/test/test_grey.py index 6cadfd3..c9236fd 100644 --- a/test/test_grey.py +++ b/test/test_grey.py @@ -8,6 +8,7 @@ from scipy.ndimage import grey_erosion as scipy_grey_erosion from scipy.ndimage import grey_opening as scipy_grey_opening from scipy.ndimage import morphological_gradient as scipy_morphological_gradient +from scipy.ndimage import morphological_laplace as scipy_morphological_laplace from scipy.ndimage import white_tophat as scipy_white_tophat import torchmorph as tm @@ -22,6 +23,7 @@ pytest.param(tm.grey_opening, scipy_grey_opening, id="opening"), pytest.param(tm.grey_closing, scipy_grey_closing, id="closing"), pytest.param(tm.morphological_gradient, scipy_morphological_gradient, id="gradient"), + pytest.param(tm.morphological_laplace, scipy_morphological_laplace, id="laplace"), pytest.param(tm.white_tophat, scipy_white_tophat, id="white_tophat"), pytest.param(tm.black_tophat, scipy_black_tophat, id="black_tophat"), ] @@ -32,6 +34,7 @@ pytest.param(tm.grey_opening, id="opening"), pytest.param(tm.grey_closing, id="closing"), pytest.param(tm.morphological_gradient, id="gradient"), + pytest.param(tm.morphological_laplace, id="laplace"), pytest.param(tm.white_tophat, id="white_tophat"), pytest.param(tm.black_tophat, id="black_tophat"), ] diff --git a/torchmorph/__init__.py b/torchmorph/__init__.py index 7a39565..559cb0d 100644 --- a/torchmorph/__init__.py +++ b/torchmorph/__init__.py @@ -17,6 +17,7 @@ grey_opening, iterate_structure, morphological_gradient, + morphological_laplace, white_tophat, ) @@ -36,6 +37,7 @@ "grey_opening", "grey_closing", "morphological_gradient", + "morphological_laplace", "white_tophat", "black_tophat", ] diff --git a/torchmorph/morphology/__init__.py b/torchmorph/morphology/__init__.py index b3a9009..bf6dea0 100644 --- a/torchmorph/morphology/__init__.py +++ b/torchmorph/morphology/__init__.py @@ -6,6 +6,7 @@ grey_erosion, grey_opening, morphological_gradient, + morphological_laplace, white_tophat, ) from .structure import generate_binary_structure, iterate_structure @@ -22,6 +23,7 @@ "grey_opening", "grey_closing", "morphological_gradient", + "morphological_laplace", "white_tophat", "black_tophat", ] diff --git a/torchmorph/morphology/grey.py b/torchmorph/morphology/grey.py index f729ac3..b3a52e2 100644 --- a/torchmorph/morphology/grey.py +++ b/torchmorph/morphology/grey.py @@ -320,3 +320,39 @@ def black_tophat( output.copy_(result) return output return result + + +def morphological_laplace( + input: Tensor, + size: int | tuple[int, ...] | None = None, + footprint: Tensor | None = None, + structure: Tensor | None = None, + output: Tensor | None = None, + mode: str = "reflect", + cval: float = 0.0, + origin: int | tuple[int, ...] = 0, +) -> Tensor: + """N-dimensional morphological Laplace for ``(B, C, Spatial...)`` CUDA tensors.""" + dilated = grey_dilation( + input, + size=size, + footprint=footprint, + structure=structure, + mode=mode, + cval=cval, + origin=origin, + ) + eroded = grey_erosion( + input, + size=size, + footprint=footprint, + structure=structure, + mode=mode, + cval=cval, + origin=origin, + ) + result = dilated + eroded - 2 * input.float() + if output is not None: + output.copy_(result) + return output + return result From 853f98e1068c47b0728bd4ec367f8355406fb63f Mon Sep 17 00:00:00 2001 From: Yuhandeng Date: Tue, 7 Jul 2026 15:45:52 +0800 Subject: [PATCH 08/25] CUDA binary erosion dilation --- benchmark/binary.py | 350 ++++--------- ...alized-binary-morphology-cuda-primitive.md | 23 + test/test_binary.py | 402 +++++--------- torchmorph/csrc/binary_kernel.cu | 272 ++++++++++ torchmorph/csrc/torchmorph.cpp | 28 + torchmorph/morphology/binary.py | 494 ++++++++++-------- 6 files changed, 816 insertions(+), 753 deletions(-) create mode 100644 docs/adr/0001-generalized-binary-morphology-cuda-primitive.md create mode 100644 torchmorph/csrc/binary_kernel.cu diff --git a/benchmark/binary.py b/benchmark/binary.py index 319df3c..b519550 100644 --- a/benchmark/binary.py +++ b/benchmark/binary.py @@ -1,3 +1,5 @@ +import argparse + import scipy.ndimage as ndi import torch import torch.utils.benchmark as benchmark @@ -5,92 +7,32 @@ import torchmorph as tm -image_size = [64, 128, 256, 1024] -batch_size = [1, 2, 4, 16] +IMAGE_SIZES = [64, 128, 256, 1024] +BATCH_SIZES = [1, 2, 4, 16] MIN_RUN_TIME = 1.0 +BINARY_OPERATORS = { + "erosion": (ndi.binary_erosion, tm.binary_erosion), + "dilation": (ndi.binary_dilation, tm.binary_dilation), + "opening": (ndi.binary_opening, tm.binary_opening), + "closing": (ndi.binary_closing, tm.binary_closing), +} -def bench_binary_erosion(): - print("\n============================================") - print(" Benchmark: binary erosion ") - print("============================================") - - for batch in batch_size: - table = PrettyTable() - table.field_names = [ - "Size", - "SciPy(ms)", - "Torch 1x(ms)", - "Torch batch(ms)", - "Speedup 1x", - "Speedup batch", - ] - for column in table.field_names: - table.align[column] = "r" - - for size in image_size: - x = (torch.randn(batch, 1, size, size, device="cuda") > 0).to(torch.float32) - # scipy data - x_np_list = [x[i, 0].detach().cpu().numpy() for i in range(batch)] - # torch 1x data - x_one = [x[i : i + 1] for i in range(batch)] - - t_scipy = benchmark.Timer( - stmt="[scipy_erosion(data) for data in x_np_list]", - globals={ - "scipy_erosion": ndi.binary_erosion, - "x_np_list": x_np_list, - }, - ).blocked_autorange(min_run_time=MIN_RUN_TIME) - - # warmup - for data in x_one: - tm.binary_erosion(data) - torch.cuda.synchronize() - - t_torch_1x = benchmark.Timer( - stmt="[tm_erosion(data) for data in x_one]", - globals={ - "tm_erosion": tm.binary_erosion, - "x_one": x_one, - }, - ).blocked_autorange(min_run_time=MIN_RUN_TIME) - - t_torch_batch = benchmark.Timer( - stmt="tm_erosion(x)", - globals={ - "tm_erosion": tm.binary_erosion, - "x": x, - }, - ).blocked_autorange(min_run_time=MIN_RUN_TIME) - - scipy_ms = t_scipy.median * 1e3 / batch - torch_1x_ms = t_torch_1x.median * 1e3 / batch - torch_batch_ms = t_torch_batch.median * 1e3 / batch - speedup1x = scipy_ms / torch_1x_ms - speedupbatch = scipy_ms / torch_batch_ms +def run_cuda(torch_op, x): + result = torch_op(x) + torch.cuda.synchronize() + return result - table.add_row( - [ - size, - f"{scipy_ms:.3f}", - f"{torch_1x_ms:.3f}", - f"{torch_batch_ms:.3f}", - f"{speedup1x:.1f}x", - f"{speedupbatch:.1f}x", - ] - ) - print(f"\n=========== Batch size : {batch} ===========") - print(table) +def bench_binary_operator(operation, image_sizes, batch_sizes, min_run_time): + scipy_op, torch_op = BINARY_OPERATORS[operation] -def bench_binary_dilation(): print("\n============================================") - print(" Benchmark: binary dilation ") + print(f" Benchmark: binary {operation} ") print("============================================") - for batch in batch_size: + for batch_size in batch_sizes: table = PrettyTable() table.field_names = [ "Size", @@ -103,215 +45,103 @@ def bench_binary_dilation(): for column in table.field_names: table.align[column] = "r" - for size in image_size: - x = (torch.randn(batch, 1, size, size, device="cuda") > 0).to(torch.float32) - # scipy data - x_np_list = [x[i, 0].detach().cpu().numpy() for i in range(batch)] - # torch 1x data - x_one = [x[i : i + 1] for i in range(batch)] + for image_size in image_sizes: + x = (torch.randn(batch_size, 1, image_size, image_size, device="cuda") > 0).to( + torch.float32 + ) + x_np_list = [x[i, 0].detach().cpu().numpy() for i in range(batch_size)] + x_one = [x[i : i + 1] for i in range(batch_size)] - t_scipy = benchmark.Timer( - stmt="[scipy_dilation(data) for data in x_np_list]", + scipy_time = benchmark.Timer( + stmt="[scipy_op(data) for data in inputs]", globals={ - "scipy_dilation": ndi.binary_dilation, - "x_np_list": x_np_list, + "scipy_op": scipy_op, + "inputs": x_np_list, }, - ).blocked_autorange(min_run_time=MIN_RUN_TIME) + ).blocked_autorange(min_run_time=min_run_time) - # warmup for data in x_one: - tm.binary_dilation(data) + torch_op(data) torch.cuda.synchronize() - t_torch_1x = benchmark.Timer( - stmt="[tm_dilation(data) for data in x_one]", + torch_single_time = benchmark.Timer( + stmt="[run_cuda(torch_op, data) for data in inputs]", globals={ - "tm_dilation": tm.binary_dilation, - "x_one": x_one, + "run_cuda": run_cuda, + "torch_op": torch_op, + "inputs": x_one, }, - ).blocked_autorange(min_run_time=MIN_RUN_TIME) + ).blocked_autorange(min_run_time=min_run_time) - t_torch_batch = benchmark.Timer( - stmt="tm_dilation(x)", + torch_batch_time = benchmark.Timer( + stmt="run_cuda(torch_op, x)", globals={ - "tm_dilation": tm.binary_dilation, + "run_cuda": run_cuda, + "torch_op": torch_op, "x": x, }, - ).blocked_autorange(min_run_time=MIN_RUN_TIME) - - scipy_ms = t_scipy.median * 1e3 / batch - torch_1x_ms = t_torch_1x.median * 1e3 / batch - torch_batch_ms = t_torch_batch.median * 1e3 / batch + ).blocked_autorange(min_run_time=min_run_time) - speedup1x = scipy_ms / torch_1x_ms - speedupbatch = scipy_ms / torch_batch_ms + scipy_ms = scipy_time.median * 1e3 / batch_size + torch_single_ms = torch_single_time.median * 1e3 / batch_size + torch_batch_ms = torch_batch_time.median * 1e3 / batch_size table.add_row( [ - size, + image_size, f"{scipy_ms:.3f}", - f"{torch_1x_ms:.3f}", + f"{torch_single_ms:.3f}", f"{torch_batch_ms:.3f}", - f"{speedup1x:.1f}x", - f"{speedupbatch:.1f}x", + f"{scipy_ms / torch_single_ms:.1f}x", + f"{scipy_ms / torch_batch_ms:.1f}x", ] ) - print(f"\n=========== Batch size : {batch} ===========") - print(table) - - -def bench_binary_opening(): - print("\n============================================") - print(" Benchmark: binary opening ") - print("============================================") - - for batch in batch_size: - table = PrettyTable() - table.field_names = [ - "Size", - "SciPy(ms)", - "Torch 1x(ms)", - "Torch batch(ms)", - "Speedup 1x", - "Speedup batch", - ] - for column in table.field_names: - table.align[column] = "r" - - for size in image_size: - x = (torch.randn(batch, 1, size, size, device="cuda") > 0).to(torch.float32) - # scipy data - x_np_list = [x[i, 0].detach().cpu().numpy() for i in range(batch)] - # torch 1x data - x_one = [x[i : i + 1] for i in range(batch)] - - t_scipy = benchmark.Timer( - stmt="[scipy_opening(data) for data in x_np_list]", - globals={ - "scipy_opening": ndi.binary_opening, - "x_np_list": x_np_list, - }, - ).blocked_autorange(min_run_time=MIN_RUN_TIME) - - # warmup - for data in x_one: - tm.binary_opening(data) - - torch.cuda.synchronize() - t_torch_1x = benchmark.Timer( - stmt="[tm_opening(data) for data in x_one]", - globals={ - "tm_opening": tm.binary_opening, - "x_one": x_one, - }, - ).blocked_autorange(min_run_time=MIN_RUN_TIME) - t_torch_batch = benchmark.Timer( - stmt="tm_opening(x)", - globals={ - "tm_opening": tm.binary_opening, - "x": x, - }, - ).blocked_autorange(min_run_time=MIN_RUN_TIME) - - scipy_ms = t_scipy.median * 1e3 / batch - torch_1x_ms = t_torch_1x.median * 1e3 / batch - torch_batch_ms = t_torch_batch.median * 1e3 / batch - - speedup1x = scipy_ms / torch_1x_ms - speedupbatch = scipy_ms / torch_batch_ms - - table.add_row( - [ - size, - f"{scipy_ms:.3f}", - f"{torch_1x_ms:.3f}", - f"{torch_batch_ms:.3f}", - f"{speedup1x:.1f}x", - f"{speedupbatch:.1f}x", - ] - ) - print(f"\n=========== Batch size : {batch} ===========") + print(f"\n=========== Batch size : {batch_size} ===========") print(table) -def bench_binary_closing(): - print("\n============================================") - print(" Benchmark: binary closing ") - print("============================================") - - for batch in batch_size: - table = PrettyTable() - table.field_names = [ - "Size", - "SciPy(ms)", - "Torch 1x(ms)", - "Torch batch(ms)", - "Speedup 1x", - "Speedup batch", - ] - for column in table.field_names: - table.align[column] = "r" - - for size in image_size: - x = (torch.randn(batch, 1, size, size, device="cuda") > 0).to(torch.float32) - # scipy data - x_np_list = [x[i, 0].detach().cpu().numpy() for i in range(batch)] - # torch 1x data - x_one = [x[i : i + 1] for i in range(batch)] - - t_scipy = benchmark.Timer( - stmt="[scipy_closing(data) for data in x_np_list]", - globals={ - "scipy_closing": ndi.binary_closing, - "x_np_list": x_np_list, - }, - ).blocked_autorange(min_run_time=MIN_RUN_TIME) - - # warmup - for data in x_one: - tm.binary_closing(data) - - torch.cuda.synchronize() - t_torch_1x = benchmark.Timer( - stmt="[tm_closing(data) for data in x_one]", - globals={ - "tm_closing": tm.binary_closing, - "x_one": x_one, - }, - ).blocked_autorange(min_run_time=MIN_RUN_TIME) - - t_torch_batch = benchmark.Timer( - stmt="tm_closing(x)", - globals={ - "tm_closing": tm.binary_closing, - "x": x, - }, - ).blocked_autorange(min_run_time=MIN_RUN_TIME) - - scipy_ms = t_scipy.median * 1e3 / batch - torch_1x_ms = t_torch_1x.median * 1e3 / batch - torch_batch_ms = t_torch_batch.median * 1e3 / batch - - speedup1x = scipy_ms / torch_1x_ms - speedupbatch = scipy_ms / torch_batch_ms - - table.add_row( - [ - size, - f"{scipy_ms:.3f}", - f"{torch_1x_ms:.3f}", - f"{torch_batch_ms:.3f}", - f"{speedup1x:.1f}x", - f"{speedupbatch:.1f}x", - ] - ) - print(f"\n=========== Batch size : {batch} ===========") - print(table) +def main(): + choices = (*BINARY_OPERATORS, "all") + parser = argparse.ArgumentParser(description="Benchmark binary morphology operators.") + parser.add_argument( + "operation", + choices=choices, + nargs="?", + default="all", + help="Operator to benchmark (default: all).", + ) + parser.add_argument( + "--sizes", + type=int, + nargs="+", + default=IMAGE_SIZES, + help="Image sizes to benchmark (default: 64 128 256 1024).", + ) + parser.add_argument( + "--batches", + type=int, + nargs="+", + default=BATCH_SIZES, + help="Batch sizes to benchmark (default: 1 2 4 16).", + ) + parser.add_argument( + "--min-run-time", + type=float, + default=MIN_RUN_TIME, + help="Minimum benchmark time per timer in seconds (default: 1.0).", + ) + args = parser.parse_args() + + operations = BINARY_OPERATORS if args.operation == "all" else (args.operation,) + for operation in operations: + bench_binary_operator( + operation, + image_sizes=args.sizes, + batch_sizes=args.batches, + min_run_time=args.min_run_time, + ) if __name__ == "__main__": - bench_binary_erosion() - bench_binary_dilation() - bench_binary_opening() - bench_binary_closing() + main() diff --git a/docs/adr/0001-generalized-binary-morphology-cuda-primitive.md b/docs/adr/0001-generalized-binary-morphology-cuda-primitive.md new file mode 100644 index 0000000..435a58a --- /dev/null +++ b/docs/adr/0001-generalized-binary-morphology-cuda-primitive.md @@ -0,0 +1,23 @@ +# 0001 Generalized Binary Morphology CUDA Primitive + +## Status + +Accepted + +## Context + +TorchMorph already uses a generalized CUDA primitive for grey erosion and grey dilation. Binary erosion and binary dilation were implemented in Python with PyTorch padding and convolution helpers. That implementation kept SciPy-aligned behavior, but it made binary morphology less consistent with the project goal of CUDA-accelerated batch-parallel morphology. + +Binary erosion and binary dilation are also the base operations for binary opening, binary closing, binary propagation, and binary hole filling. Optimizing only the public erosion and dilation wrappers would leave later binary operators without a clear reusable CUDA foundation. + +## Decision + +Implement a generalized binary morphology CUDA primitive, modeled after the grey morphology CUDA primitive, and dispatch CUDA binary erosion and binary dilation through it. The primitive supports batch-channel tensors with arbitrary spatial rank up to the same practical maximum used by the grey implementation. Python remains responsible for iteration loops, mask application, output copying, and composite operators. + +Do not start with specialized 2D or 3D kernels. Keep those as later optimizations once the generalized primitive is correct and benchmarked. + +## Consequences + +The first CUDA implementation favors broad SciPy-aligned behavior and shared maintenance shape over peak 2D/3D performance. This gives binary propagation and binary hole filling a reusable base and keeps the public Python API stable. + +Specialized 2D and 3D kernels may still be added later as fast paths when benchmark data shows that fixed-rank kernels would materially improve common workloads. diff --git a/test/test_binary.py b/test/test_binary.py index c961b45..4d7d737 100644 --- a/test/test_binary.py +++ b/test/test_binary.py @@ -1,13 +1,11 @@ -import numpy as np # noqa: F401 +import numpy as np import pytest import torch -from scipy.ndimage import ( - binary_closing, - binary_dilation, - binary_erosion, - binary_opening, - generate_binary_structure, -) +from scipy.ndimage import binary_closing as scipy_binary_closing +from scipy.ndimage import binary_dilation as scipy_binary_dilation +from scipy.ndimage import binary_erosion as scipy_binary_erosion +from scipy.ndimage import binary_opening as scipy_binary_opening +from scipy.ndimage import generate_binary_structure import torchmorph as tm @@ -15,312 +13,174 @@ not torch.cuda.is_available(), reason="CUDA is required for torchmorph tests" ) +BINARY_OPERATORS = [ + pytest.param(tm.binary_erosion, scipy_binary_erosion, id="erosion"), + pytest.param(tm.binary_dilation, scipy_binary_dilation, id="dilation"), + pytest.param(tm.binary_opening, scipy_binary_opening, id="opening"), + pytest.param(tm.binary_closing, scipy_binary_closing, id="closing"), +] + +TORCH_OPERATORS = [ + pytest.param(tm.binary_erosion, id="erosion"), + pytest.param(tm.binary_dilation, id="dilation"), + pytest.param(tm.binary_opening, id="opening"), + pytest.param(tm.binary_closing, id="closing"), +] + +CASE_2D = np.array( + [[[[0, 1, 0], [1, 1, 1], [0, 1, 0]]]], + dtype=np.float32, +) +CASE_3D = np.zeros((2, 1, 5, 5, 5), dtype=bool) +CASE_3D[0, 0, 2:3, 2:4, 1:2] = True +CASE_4D = np.zeros((2, 1, 4, 4, 4, 4), dtype=bool) +CASE_4D[0, 0, 1:3, 1:2, 2:3] = True + +STRUCTURE_2D = generate_binary_structure(rank=2, connectivity=1) +STRUCTURE_3D_1 = generate_binary_structure(rank=3, connectivity=2) +STRUCTURE_3D_2 = generate_binary_structure(rank=3, connectivity=3) +STRUCTURE_4D = generate_binary_structure(rank=4, connectivity=4) + +MASK_2D = np.array([[1, 1, 0], [1, 1, 1], [1, 1, 0]], dtype=bool) +MASK_3D = np.zeros((5, 5, 5), dtype=bool) +MASK_3D[0:2, 0:1, 0:3] = True +MASK_4D = np.zeros((4, 4, 4, 4), dtype=bool) +MASK_4D[1:2, 0:1, 1:3] = True + -def apply_scipy_to_batch(np_input, scipy_func, **kwargs): - """ - Apply a SciPy operator to each (B, C) sample independently. - The input is expected to have shape (B, C, Spatial...). The batch and - channel dimensions are flattened, `scipy_func` is applied to each spatial - sample separately, and the results are reshaped back to the original - layout. - Args: - np_input: Input array in (B, C, Spatial...) format. - scipy_func: SciPy function applied to each spatial sample. - **kwargs: Additional keyword arguments forwarded to `scipy_func`. - Returns: - np.ndarray: Output array with the same shape as `np_input`. - """ - input_shape = np_input.shape - batch_shape = input_shape[:2] - spatial_shape = input_shape[2:] - batch_size = int(np.prod(batch_shape)) - flatten_input = np_input.reshape(batch_size, *spatial_shape) +def apply_scipy_to_batch(np_input, scipy_op, **kwargs): + batch_shape = np_input.shape[:2] + spatial_shape = np_input.shape[2:] + samples = np_input.reshape(-1, *spatial_shape) mask = kwargs.pop("mask", None) output = kwargs.pop("output", None) - if mask is not None and mask.shape != spatial_shape: - raise ValueError(f"mask shape must be {spatial_shape}, got {mask.shape}") - if output is not None and output.shape != spatial_shape: - raise ValueError(f"output shape must be {spatial_shape}, got {output.shape}") - results = [] - for sample in flatten_input: + for sample in samples: sample_kwargs = dict(kwargs) if mask is not None: sample_kwargs["mask"] = mask if output is not None: - sample_output = np.empty_like(output) - sample_kwargs["output"] = sample_output - result = scipy_func(sample, **sample_kwargs) + sample_kwargs["output"] = np.empty_like(output) + result = scipy_op(sample, **sample_kwargs) results.append(np.asarray(result).copy()) - return np.stack(results, axis=0).reshape(*batch_shape, *spatial_shape) - - -def batch_scipy( - np_input: np.ndarray, - scipy_func: callable, - structure: np.ndarray | None = None, - iterations: int = 1, - mask: np.ndarray | None = None, - output: np.ndarray | None = None, - border_value: bool = False, - origin: int | tuple[int, ...] = 0, -) -> np.ndarray: - return apply_scipy_to_batch( - np_input, - scipy_func, - structure=structure, - iterations=iterations, - mask=mask, - output=output, - border_value=border_value, - origin=origin, - ) - - -# case -case_2d = np.array( - [ - [ - [ - [0, 1, 0], - [1, 1, 1], - [0, 1, 0], - ] - ] - ], - dtype=np.float32, -) -case_3d = np.zeros((2, 1, 5, 5, 5), dtype=bool) -case_3d[0, 0, 2:3, 2:4, 1:2] = True -case_4d = np.zeros((2, 1, 4, 4, 4, 4), dtype=bool) -case_4d[0, 0, 1:3, 1:2, 2:3] = True + return np.stack(results).reshape(*batch_shape, *spatial_shape) -# structure -structure_2d = np.array([[0, 1, 0], [1, 1, 1], [0, 1, 0]]) -structure_3d_1 = generate_binary_structure(rank=3, connectivity=2) -structure_3d_2 = generate_binary_structure(rank=3, connectivity=3) +def optional_cuda_tensor(value, dtype): + if value is None: + return None + return torch.as_tensor(value, dtype=dtype, device="cuda") -structure_4d = generate_binary_structure(rank=4, connectivity=4) -# mask -mask_2d_np = np.array([[1, 1, 0], [1, 1, 1], [1, 1, 0]], dtype=bool) -mask_2d_tm = torch.tensor([[[[1, 1, 0], [1, 1, 1], [1, 1, 0]]]], dtype=bool) - -mask_3d_np = np.zeros((5, 5, 5), dtype=bool) -mask_3d_np[0:2, 0:1, 0:3] = True -mask_3d_tm = torch.zeros((2, 1, 5, 5, 5), dtype=bool) -mask_3d_tm[:, :, 0:2, 0:1, 0:3] = True - -mask_4d_np = np.zeros((4, 4, 4, 4), dtype=bool) -mask_4d_np[1:2, 0:1, 1:3] = True -mask_4d_tm = torch.zeros((2, 1, 4, 4, 4, 4), dtype=bool) -mask_4d_tm[:, :, 1:2, 0:1, 1:3] = True - -# output -output_2d_np = np.empty([3, 3]) -output_2d_tm = torch.empty([1, 1, 3, 3]) - -output_3d_np = np.empty([5, 5, 5]) -output_3d_tm = torch.empty([2, 1, 5, 5, 5]) - -output_4d_np = np.empty([4, 4, 4, 4]) -output_4d_tm = torch.empty([2, 1, 4, 4, 4, 4]) +def cuda_mask(mask, np_input): + return torch.as_tensor(mask, dtype=torch.bool, device="cuda").expand(np_input.shape) +@pytest.mark.parametrize(("torch_op", "scipy_op"), BINARY_OPERATORS) @pytest.mark.parametrize( - ("np_input, scipy_func, structure, iterations, origin, border_value"), + ("np_input", "structure", "iterations", "origin", "border_value"), [ - pytest.param(case_2d, binary_erosion, None, 1, 0, False, id="er_2D_basic"), - pytest.param(case_2d, binary_erosion, structure_2d, 1, 0, False, id="er_2D_structure"), - pytest.param(case_2d, binary_erosion, None, 2, 0, False, id="er_2D_2iterations"), - pytest.param(case_2d, binary_erosion, None, -1, 0, False, id="er_2D_-1iterations"), - pytest.param(case_2d, binary_erosion, None, 1, 1, False, id="er_2D_1origin_False"), - pytest.param(case_2d, binary_erosion, None, 1, 1, True, id="er_2D_1origin_True"), - pytest.param(case_3d, binary_erosion, None, 1, 0, False, id="er_3D_basic"), - pytest.param(case_3d, binary_erosion, structure_3d_1, 1, 0, False, id="er_3D_2Dstructure1"), - pytest.param(case_3d, binary_erosion, structure_3d_2, 1, 0, False, id="er_3D_3Dstructure2"), - pytest.param(case_3d, binary_erosion, None, 2, 0, False, id="er_3D_2iterations"), - pytest.param(case_3d, binary_erosion, None, -1, 0, False, id="er_3D_-1iterations"), - pytest.param(case_3d, binary_erosion, None, 1, 1, False, id="er_3D_1origin_False"), - pytest.param(case_3d, binary_erosion, None, 1, 1, True, id="er_3D_1origin_True"), - pytest.param(case_4d, binary_erosion, None, 1, 0, False, id="er_4D_basic"), - pytest.param(case_4d, binary_erosion, structure_4d, 1, 0, False, id="er_4D_4Dstructure"), - pytest.param(case_4d, binary_erosion, None, 2, 0, False, id="er_4D_2iterations"), - pytest.param(case_4d, binary_erosion, None, -1, 0, False, id="er_4D_-1iterations"), - pytest.param(case_4d, binary_erosion, None, 1, 1, False, id="er_4D_1origin_False"), - pytest.param(case_4d, binary_erosion, None, 1, 1, True, id="er_4D_1origin_True"), - pytest.param(case_2d, binary_dilation, None, 1, 0, False, id="di_2D_basic"), - pytest.param(case_2d, binary_dilation, structure_2d, 1, 0, False, id="di_2D_structure"), - pytest.param(case_2d, binary_dilation, None, 2, 0, False, id="di_2D_2iterations"), - pytest.param(case_2d, binary_dilation, None, -1, 0, False, id="di_2D_-1iterations"), - pytest.param(case_2d, binary_dilation, None, 1, 1, False, id="di_2D_1origin_False"), - pytest.param(case_2d, binary_dilation, None, 1, 1, True, id="di_2D_1origin_True"), - pytest.param(case_3d, binary_dilation, None, 1, 0, False, id="di_3D_basic"), - pytest.param( - case_3d, binary_dilation, structure_3d_1, 1, 0, False, id="di_3D_2Dstructure1" - ), - pytest.param( - case_3d, binary_dilation, structure_3d_2, 1, 0, False, id="di_3D_3Dstructure2" - ), - pytest.param(case_3d, binary_dilation, None, 2, 0, False, id="di_3D_2iterations"), - pytest.param(case_3d, binary_dilation, None, -1, 0, False, id="di_3D_-1iterations"), - pytest.param(case_3d, binary_dilation, None, 1, 1, False, id="di_3D_1origin_False"), - pytest.param(case_3d, binary_dilation, None, 1, 1, True, id="di_3D_1origin_True"), - pytest.param(case_4d, binary_dilation, None, 1, 0, False, id="di_4D_basic"), - pytest.param(case_4d, binary_dilation, structure_4d, 1, 0, False, id="di_4D_4Dstructure"), - pytest.param(case_4d, binary_dilation, None, 2, 0, False, id="di_4D_2iterations"), - pytest.param(case_4d, binary_dilation, None, -1, 0, False, id="di_4D_-1iterations"), - pytest.param(case_4d, binary_dilation, None, 1, 1, False, id="di_4D_1origin_False"), - pytest.param(case_4d, binary_dilation, None, 1, 1, True, id="di_4D_1origin_True"), - pytest.param(case_2d, binary_opening, None, 1, 0, False, id="op_2D_basic"), - pytest.param(case_2d, binary_opening, structure_2d, 1, 0, False, id="op_2D_structure"), - pytest.param(case_2d, binary_opening, None, 2, 0, False, id="op_2D_2iterations"), - pytest.param(case_2d, binary_opening, None, -1, 0, False, id="op_2D_-1iterations"), - pytest.param(case_2d, binary_opening, None, 1, 1, False, id="op_2D_1origin_False"), - pytest.param(case_2d, binary_opening, None, 1, 1, True, id="op_2D_1origin_True"), - pytest.param(case_3d, binary_opening, None, 1, 0, False, id="op_3D_basic"), - pytest.param(case_3d, binary_opening, structure_3d_1, 1, 0, False, id="op_3D_2Dstructure1"), - pytest.param(case_3d, binary_opening, structure_3d_2, 1, 0, False, id="op_3D_3Dstructure2"), - pytest.param(case_3d, binary_opening, None, 2, 0, False, id="op_3D_2iterations"), - pytest.param(case_3d, binary_opening, None, -1, 0, False, id="op_3D_-1iterations"), - pytest.param(case_3d, binary_opening, None, 1, 1, False, id="op_3D_1origin_False"), - pytest.param(case_3d, binary_opening, None, 1, 1, True, id="op_3D_1origin_True"), - pytest.param(case_4d, binary_opening, None, 1, 0, False, id="op_4D_basic"), - pytest.param(case_4d, binary_opening, structure_4d, 1, 0, False, id="op_4D_4Dstructure"), - pytest.param(case_4d, binary_opening, None, 2, 0, False, id="op_4D_2iterations"), - pytest.param(case_4d, binary_opening, None, -1, 0, False, id="op_4D_-1iterations"), - pytest.param(case_4d, binary_opening, None, 1, 1, False, id="op_4D_1origin_False"), - pytest.param(case_4d, binary_opening, None, 1, 1, True, id="op_4D_1origin_True"), - pytest.param(case_2d, binary_closing, None, 1, 0, False, id="cl_2D_basic"), - pytest.param(case_2d, binary_closing, structure_2d, 1, 0, False, id="cl_2D_structure"), - pytest.param(case_2d, binary_closing, None, 2, 0, False, id="cl_2D_2iterations"), - pytest.param(case_2d, binary_closing, None, -1, 0, False, id="cl_2D_-1iterations"), - pytest.param(case_2d, binary_closing, None, 1, 1, False, id="cl_2D_1origin_False"), - pytest.param(case_2d, binary_closing, None, 1, 1, True, id="cl_2D_1origin_True"), - pytest.param(case_3d, binary_closing, None, 1, 0, False, id="cl_3D_basic"), - pytest.param(case_3d, binary_closing, structure_3d_1, 1, 0, False, id="cl_3D_2Dstructure1"), - pytest.param(case_3d, binary_closing, structure_3d_2, 1, 0, False, id="cl_3D_3Dstructure2"), - pytest.param(case_3d, binary_closing, None, 2, 0, False, id="cl_3D_2iterations"), - pytest.param(case_3d, binary_closing, None, -1, 0, False, id="cl_3D_-1iterations"), - pytest.param(case_3d, binary_closing, None, 1, 1, False, id="cl_3D_1origin_False"), - pytest.param(case_3d, binary_closing, None, 1, 1, True, id="cl_3D_1origin_True"), - pytest.param(case_4d, binary_closing, None, 1, 0, False, id="cl_4D_basic"), - pytest.param(case_4d, binary_closing, structure_4d, 1, 0, False, id="cl_4D_4Dstructure"), - pytest.param(case_4d, binary_closing, None, 2, 0, False, id="cl_4D_2iterations"), - pytest.param(case_4d, binary_closing, None, -1, 0, False, id="cl_4D_-1iterations"), - pytest.param(case_4d, binary_closing, None, 1, 1, False, id="cl_4D_1origin_False"), - pytest.param(case_4d, binary_closing, None, 1, 1, True, id="cl_4D_1origin_True"), + pytest.param(CASE_2D, None, 1, 0, False, id="2d_basic"), + pytest.param(CASE_2D, STRUCTURE_2D, 1, 0, False, id="2d_structure"), + pytest.param(CASE_2D, None, 2, 0, False, id="2d_iterations"), + pytest.param(CASE_2D, None, -1, 0, False, id="2d_until_stable"), + pytest.param(CASE_2D, None, 1, 1, True, id="2d_origin_border"), + pytest.param(CASE_3D, None, 1, 0, False, id="3d_basic"), + pytest.param(CASE_3D, STRUCTURE_3D_1, 1, 0, False, id="3d_structure_1"), + pytest.param(CASE_3D, STRUCTURE_3D_2, 1, 1, True, id="3d_origin_border"), + pytest.param(CASE_4D, STRUCTURE_4D, 1, 0, False, id="4d_structure"), + pytest.param(CASE_4D, None, 2, 0, False, id="4d_iterations"), ], ) -def test_binary_basic( - np_input, - scipy_func, - structure, - iterations, - origin, - border_value, +def test_binary_morphology_matches_scipy( + torch_op, scipy_op, np_input, structure, iterations, origin, border_value ): - x = torch.as_tensor(np_input, dtype=torch.float32) - if structure is not None: - structure_cuda = torch.as_tensor(structure, dtype=torch.float32) - else: - structure_cuda = None - tm_func = getattr(tm, scipy_func.__name__) - actual = tm_func( - x, - structure=structure_cuda, + result = torch_op( + torch.as_tensor(np_input, dtype=torch.float32, device="cuda"), + structure=optional_cuda_tensor(structure, torch.bool), iterations=iterations, origin=origin, border_value=border_value, ) - expected_np = batch_scipy( + expected = apply_scipy_to_batch( np_input, - scipy_func, + scipy_op, structure=structure, iterations=iterations, origin=origin, border_value=border_value, ) - expected = torch.as_tensor(expected_np) - torch.testing.assert_close(actual, expected) + torch.testing.assert_close(result.cpu(), torch.as_tensor(expected)) +@pytest.mark.parametrize(("torch_op", "scipy_op"), BINARY_OPERATORS) @pytest.mark.parametrize( - ("input_np, scipy_func, mask_np, mask_tm"), + ("np_input", "mask"), [ - pytest.param(case_2d, binary_erosion, mask_2d_np, mask_2d_tm, id="er_2d_mask"), - pytest.param(case_2d, binary_dilation, mask_2d_np, mask_2d_tm, id="di_2d_mask"), - pytest.param(case_2d, binary_opening, mask_2d_np, mask_2d_tm, id="op_2d_mask"), - pytest.param(case_2d, binary_closing, mask_2d_np, mask_2d_tm, id="cl_2d_mask"), - pytest.param(case_3d, binary_erosion, mask_3d_np, mask_3d_tm, id="er_3d_mask"), - pytest.param(case_3d, binary_dilation, mask_3d_np, mask_3d_tm, id="di_3d_mask"), - pytest.param(case_3d, binary_opening, mask_3d_np, mask_3d_tm, id="op_3d_mask"), - pytest.param(case_3d, binary_closing, mask_3d_np, mask_3d_tm, id="cl_3d_mask"), - pytest.param(case_4d, binary_erosion, mask_4d_np, mask_4d_tm, id="er_4d_mask"), - pytest.param(case_4d, binary_dilation, mask_4d_np, mask_4d_tm, id="di_4d_mask"), - pytest.param(case_4d, binary_opening, mask_4d_np, mask_4d_tm, id="op_4d_mask"), - pytest.param(case_4d, binary_closing, mask_4d_np, mask_4d_tm, id="cl_4d_mask"), + pytest.param(CASE_2D, MASK_2D, id="2d_mask"), + pytest.param(CASE_3D, MASK_3D, id="3d_mask"), + pytest.param(CASE_4D, MASK_4D, id="4d_mask"), ], ) -def test_binary_mask( - input_np, - scipy_func, - mask_np, - mask_tm, -): - x = torch.as_tensor(input_np, dtype=torch.float32) - tm_func = getattr(tm, scipy_func.__name__) - actual = tm_func( - x, - mask=mask_tm, - ) - - expected_np = batch_scipy( - input_np, - scipy_func, - mask=mask_np, +def test_binary_morphology_mask(torch_op, scipy_op, np_input, mask): + result = torch_op( + torch.as_tensor(np_input, dtype=torch.float32, device="cuda"), + mask=cuda_mask(mask, np_input), ) - expected = torch.as_tensor(expected_np) - torch.testing.assert_close(actual, expected) + expected = apply_scipy_to_batch(np_input, scipy_op, mask=mask) + torch.testing.assert_close(result.cpu(), torch.as_tensor(expected)) +@pytest.mark.parametrize(("torch_op", "scipy_op"), BINARY_OPERATORS) @pytest.mark.parametrize( - "input_np,scipy_func, output_np, output_tm", + ("np_input", "output_shape"), [ - pytest.param(case_2d, binary_erosion, output_2d_np, output_2d_tm, id='er_output_2d'), - pytest.param(case_2d, binary_dilation, output_2d_np, output_2d_tm, id='di_output_2d'), - pytest.param(case_2d, binary_opening, output_2d_np, output_2d_tm, id='op_output_2d'), - pytest.param(case_2d, binary_closing, output_2d_np, output_2d_tm, id='cl_output_2d'), - pytest.param(case_3d, binary_erosion, output_3d_np, output_3d_tm, id='er_output_3d'), - pytest.param(case_3d, binary_dilation, output_3d_np, output_3d_tm, id='di_output_3d'), - pytest.param(case_3d, binary_opening, output_3d_np, output_3d_tm, id='op_output_3d'), - pytest.param(case_3d, binary_closing, output_3d_np, output_3d_tm, id='cl_output_3d'), - pytest.param(case_4d, binary_erosion, output_4d_np, output_4d_tm, id='er_output_4d'), - pytest.param(case_4d, binary_dilation, output_4d_np, output_4d_tm, id='di_output_4d'), - pytest.param(case_4d, binary_opening, output_4d_np, output_4d_tm, id='op_output_4d'), - pytest.param(case_4d, binary_closing, output_4d_np, output_4d_tm, id='cl_output_4d'), + pytest.param(CASE_2D, (3, 3), id="2d_output"), + pytest.param(CASE_3D, (5, 5, 5), id="3d_output"), + pytest.param(CASE_4D, (4, 4, 4, 4), id="4d_output"), ], ) -def test_binary_output( - input_np, - scipy_func, - output_np, - output_tm, -): - x = torch.as_tensor(input_np, dtype=torch.float32) - tm_func = getattr(tm, scipy_func.__name__) +def test_binary_morphology_output(torch_op, scipy_op, np_input, output_shape): + x = torch.as_tensor(np_input, dtype=torch.float32, device="cuda") + output = torch.empty_like(x, dtype=torch.bool) - actual = tm_func( - x, - output=output_tm, - ) + result = torch_op(x, output=output) - expected_np = batch_scipy( - input_np, - scipy_func, - output=output_np, - ) - expected = torch.as_tensor(expected_np, dtype=torch.float32) - torch.testing.assert_close(actual, expected) + assert result is output + expected = apply_scipy_to_batch(np_input, scipy_op, output=np.empty(output_shape, dtype=bool)) + torch.testing.assert_close(result.cpu(), torch.as_tensor(expected)) + + +@pytest.mark.parametrize(("torch_op", "scipy_op"), BINARY_OPERATORS) +def test_binary_morphology_cpu_fallback_matches_scipy(torch_op, scipy_op): + result = torch_op(torch.as_tensor(CASE_2D, dtype=torch.float32)) + expected = apply_scipy_to_batch(CASE_2D, scipy_op) + torch.testing.assert_close(result, torch.as_tensor(expected)) + + +@pytest.mark.parametrize("torch_op", TORCH_OPERATORS) +def test_binary_morphology_rejects_invalid_origin_dimension(torch_op): + x = torch.as_tensor(CASE_2D, dtype=torch.float32, device="cuda") + with pytest.raises(ValueError, match="origin dimension"): + torch_op(x, origin=(0, 0, 0)) + + +@pytest.mark.parametrize("torch_op", TORCH_OPERATORS) +def test_binary_morphology_uses_current_cuda_stream(torch_op): + stream = torch.cuda.Stream() + with torch.cuda.stream(stream): + x = torch.as_tensor(CASE_2D, dtype=torch.float32, device="cuda") + result = torch_op(x) + stream.synchronize() + + assert result.device == x.device + + +@pytest.mark.skipif(torch.cuda.device_count() < 2, reason="requires two CUDA devices") +@pytest.mark.parametrize("torch_op", TORCH_OPERATORS) +def test_binary_morphology_uses_input_cuda_device(torch_op): + with torch.cuda.device(0): + x = torch.as_tensor(CASE_2D, dtype=torch.float32, device="cuda:1") + result = torch_op(x) + assert result.device == x.device diff --git a/torchmorph/csrc/binary_kernel.cu b/torchmorph/csrc/binary_kernel.cu new file mode 100644 index 0000000..dae818a --- /dev/null +++ b/torchmorph/csrc/binary_kernel.cu @@ -0,0 +1,272 @@ +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#define BINARY_MORPH_MAX_NDIM 8 + +struct BinaryMorphologyGeometry { + int64_t spatial_size[BINARY_MORPH_MAX_NDIM]; + int64_t spatial_stride[BINARY_MORPH_MAX_NDIM]; + int64_t min_deltas[BINARY_MORPH_MAX_NDIM]; + int64_t max_deltas[BINARY_MORPH_MAX_NDIM]; +}; + +struct BinaryErosionOp { + __device__ __forceinline__ static bool identity() { + return true; + } + + __device__ __forceinline__ static bool done(bool result) { + return !result; + } + + __device__ __forceinline__ static bool combine(bool result, bool value) { + return result && value; + } + + __host__ __forceinline__ static int64_t offset(int64_t delta) { + return delta; + } +}; + +struct BinaryDilationOp { + __device__ __forceinline__ static bool identity() { + return false; + } + + __device__ __forceinline__ static bool done(bool result) { + return result; + } + + __device__ __forceinline__ static bool combine(bool result, bool value) { + return result || value; + } + + __host__ __forceinline__ static int64_t offset(int64_t delta) { + return -delta; + } +}; + +template +__global__ void binary_morphology_fused_kernel( + const bool* __restrict__ input, + bool* __restrict__ output, + const int64_t* __restrict__ struct_meta, + const int num_struct, + const int ndim_spatial, + const BinaryMorphologyGeometry geom, + const int64_t total_spatial, + const int64_t batch_channel, + const bool border_value +) { + int64_t idx = blockIdx.x * blockDim.x + threadIdx.x; + if (idx >= batch_channel * total_spatial) return; + + int64_t bc = idx / total_spatial; + int64_t sp = idx % total_spatial; + + int64_t coords[BINARY_MORPH_MAX_NDIM]; + int64_t rem = sp; + for (int d = 0; d < ndim_spatial; d++) { + coords[d] = rem / geom.spatial_stride[d]; + rem %= geom.spatial_stride[d]; + } + + bool result = MorphologyOp::identity(); + const bool* in_ptr = input + bc * total_spatial; + const int meta_stride = ndim_spatial + 1; + + bool is_interior = true; + for (int d = 0; d < ndim_spatial; d++) { + if (coords[d] + geom.min_deltas[d] < 0 || + coords[d] + geom.max_deltas[d] >= geom.spatial_size[d]) { + is_interior = false; + break; + } + } + + if (is_interior) { + for (int i = 0; i < num_struct; i++) { + int64_t offset = struct_meta[i * meta_stride + ndim_spatial]; + result = MorphologyOp::combine(result, in_ptr[sp + offset]); + if (MorphologyOp::done(result)) break; + } + output[idx] = result; + return; + } + + for (int i = 0; i < num_struct; i++) { + const int64_t* meta = struct_meta + i * meta_stride; + + bool value = border_value; + bool in_bounds = true; + int64_t flat_idx = 0; + for (int d = 0; d < ndim_spatial; d++) { + int64_t coord = coords[d] + meta[d]; + if (coord < 0 || coord >= geom.spatial_size[d]) { + in_bounds = false; + break; + } + flat_idx += coord * geom.spatial_stride[d]; + } + + if (in_bounds) { + value = in_ptr[flat_idx]; + } + result = MorphologyOp::combine(result, value); + if (MorphologyOp::done(result)) break; + } + + output[idx] = result; +} + +template +static torch::Tensor binary_morphology_cuda( + torch::Tensor input, + torch::Tensor structure, + std::vector origin_vec, + bool border_value +) { + TORCH_CHECK(input.is_cuda(), "input must be a CUDA tensor"); + const c10::cuda::CUDAGuard device_guard(input.device()); + + TORCH_CHECK(!structure.is_cuda(), "structure must be a CPU tensor"); + TORCH_CHECK(input.is_contiguous(), "input must be contiguous"); + TORCH_CHECK(structure.is_contiguous(), "structure must be contiguous"); + TORCH_CHECK(input.dtype() == torch::kBool, "input must be bool"); + TORCH_CHECK(structure.dtype() == torch::kBool, "structure must be bool"); + + int ndim_spatial = structure.dim(); + TORCH_CHECK(ndim_spatial > 0 && ndim_spatial <= BINARY_MORPH_MAX_NDIM, + "structure dimension must be in 1-", BINARY_MORPH_MAX_NDIM, + ", got ", ndim_spatial); + TORCH_CHECK((int)origin_vec.size() == ndim_spatial, + "origin length must match structure dimensions"); + TORCH_CHECK(input.dim() == ndim_spatial + 2, + "input spatial dimensions must match structure dimensions"); + + auto output = torch::empty_like(input); + + int B = input.size(0); + int C = input.size(1); + int64_t batch_channel = (int64_t)B * C; + + std::vector h_spatial_size(ndim_spatial); + std::vector h_spatial_stride(ndim_spatial); + for (int d = 0; d < ndim_spatial; d++) { + h_spatial_size[d] = input.size(d + 2); + } + h_spatial_stride[ndim_spatial - 1] = 1; + for (int d = ndim_spatial - 2; d >= 0; d--) { + h_spatial_stride[d] = h_spatial_stride[d + 1] * h_spatial_size[d + 1]; + } + + int64_t total_spatial = 1; + for (int d = 0; d < ndim_spatial; d++) { + total_spatial *= h_spatial_size[d]; + } + + BinaryMorphologyGeometry h_geom; + for (int d = 0; d < BINARY_MORPH_MAX_NDIM; d++) { + h_geom.spatial_size[d] = 0; + h_geom.spatial_stride[d] = 0; + h_geom.min_deltas[d] = 0; + h_geom.max_deltas[d] = 0; + } + for (int d = 0; d < ndim_spatial; d++) { + h_geom.spatial_size[d] = h_spatial_size[d]; + h_geom.spatial_stride[d] = h_spatial_stride[d]; + h_geom.min_deltas[d] = INT64_MAX; + h_geom.max_deltas[d] = INT64_MIN; + } + + auto struct_flat = structure.flatten().contiguous(); + const bool* struct_ptr = struct_flat.data_ptr(); + const int64_t total_struct = struct_flat.numel(); + + std::vector h_structure_stride(ndim_spatial); + h_structure_stride[ndim_spatial - 1] = 1; + for (int d = ndim_spatial - 2; d >= 0; d--) { + h_structure_stride[d] = h_structure_stride[d + 1] * structure.size(d + 1); + } + + std::vector h_struct_meta; + h_struct_meta.reserve(total_struct * (ndim_spatial + 1)); + + for (int64_t i = 0; i < total_struct; i++) { + if (!struct_ptr[i]) { + continue; + } + + int64_t tmp = i; + int64_t flat_offset = 0; + for (int d = 0; d < ndim_spatial; d++) { + int64_t coord_d = tmp / h_structure_stride[d]; + tmp %= h_structure_stride[d]; + + int64_t center_d = structure.size(d) / 2 + origin_vec[d]; + int64_t delta = MorphologyOp::offset(coord_d - center_d); + h_struct_meta.push_back(delta); + flat_offset += delta * h_spatial_stride[d]; + h_geom.min_deltas[d] = std::min(h_geom.min_deltas[d], delta); + h_geom.max_deltas[d] = std::max(h_geom.max_deltas[d], delta); + } + h_struct_meta.push_back(flat_offset); + } + + int num_struct = static_cast(h_struct_meta.size() / (ndim_spatial + 1)); + TORCH_CHECK(num_struct > 0, "structure must contain at least one active position"); + + auto opts_i = torch::TensorOptions().dtype(torch::kInt64); + auto d_struct_meta = torch::from_blob( + h_struct_meta.data(), {num_struct * (ndim_spatial + 1)}, opts_i + ).to(input.device()); + + int64_t total_threads = batch_channel * total_spatial; + int threads = 256; + int blocks = (int)((total_threads + threads - 1) / threads); + const auto stream = c10::cuda::getCurrentCUDAStream(input.get_device()); + + binary_morphology_fused_kernel<<>>( + input.data_ptr(), + output.data_ptr(), + d_struct_meta.data_ptr(), + num_struct, + ndim_spatial, + h_geom, + total_spatial, + batch_channel, + border_value + ); + C10_CUDA_KERNEL_LAUNCH_CHECK(); + + return output; +} + +torch::Tensor binary_erosion_cuda( + torch::Tensor input, + torch::Tensor structure, + std::vector origin_vec, + bool border_value +) { + return binary_morphology_cuda( + input, structure, origin_vec, border_value + ); +} + +torch::Tensor binary_dilation_cuda( + torch::Tensor input, + torch::Tensor structure, + std::vector origin_vec, + bool border_value +) { + return binary_morphology_cuda( + input, structure, origin_vec, border_value + ); +} diff --git a/torchmorph/csrc/torchmorph.cpp b/torchmorph/csrc/torchmorph.cpp index f87a1d6..2c7a361 100644 --- a/torchmorph/csrc/torchmorph.cpp +++ b/torchmorph/csrc/torchmorph.cpp @@ -20,6 +20,20 @@ torch::Tensor grey_dilation_cuda( float cval ); +torch::Tensor binary_erosion_cuda( + torch::Tensor input, + torch::Tensor structure, + std::vector origin, + bool border_value +); + +torch::Tensor binary_dilation_cuda( + torch::Tensor input, + torch::Tensor structure, + std::vector origin, + bool border_value +); + std::tuple edt_cuda( torch::Tensor input, std::vector sampling, @@ -65,6 +79,20 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { py::arg("mode"), py::arg("cval")); + m.def("binary_erosion_cuda", &binary_erosion_cuda, + "N-dimensional fused binary erosion", + py::arg("input"), + py::arg("structure"), + py::arg("origin"), + py::arg("border_value")); + + m.def("binary_dilation_cuda", &binary_dilation_cuda, + "N-dimensional fused binary dilation", + py::arg("input"), + py::arg("structure"), + py::arg("origin"), + py::arg("border_value")); + m.def("edt_cuda", &edt_cuda, "Exact Euclidean Distance Transform (Felzenszwalb algorithm)", py::arg("input"), diff --git a/torchmorph/morphology/binary.py b/torchmorph/morphology/binary.py index 49ec8b9..dfeeea7 100644 --- a/torchmorph/morphology/binary.py +++ b/torchmorph/morphology/binary.py @@ -1,222 +1,272 @@ -import torch -import torch.nn.functional as F -from torch import Tensor - -from ._convnd import conv_nd -from .structure import generate_binary_structure - - -def _prepare_origin(origin: int | tuple[int, ...], ndim=int) -> tuple[int, ...]: - """change the origin into tuple""" - if isinstance(origin, int): - return (origin,) * ndim - origin = tuple(origin) - - if (len(origin)) != ndim: - raise ValueError(f"origin dimension is not {ndim}, got {len(origin)}") - - return origin - - -def _extend_pad(kernel_shape: torch.Size, origin: tuple[int, ...]) -> list[int]: - """extend the padlist for kernel""" - pad = [] - for dim in range(len(kernel_shape) - 1, -1, -1): - center = kernel_shape[dim] // 2 - pad_before = center + origin[dim] - pad_after = kernel_shape[dim] - 1 - pad_before - pad.extend([pad_before, pad_after]) - return pad - - -def _flip_structure(structure: Tensor) -> Tensor: - dim = tuple(range(structure.ndim)) - return torch.flip(structure, dim) - - -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: - iterations_flag = iterations < 1 - - spatial_ndim = input.ndim - 2 - - if structure is None: - structure = generate_binary_structure(spatial_ndim, 1) - - batch, channels = input.shape[:2] - spatial_shape = input.shape[2:] - - x = (input != 0).to(dtype=torch.float32).reshape(batch * channels, 1, *spatial_shape) - structure = (structure != 0).to(device=input.device, dtype=torch.float32) - if mode == "dilation": - structure = _flip_structure(structure) - kernel = structure.unsqueeze(0).unsqueeze(0) - kernel_sum = kernel.sum() - - origin = _prepare_origin(origin, spatial_ndim) - if mode == "dilation": - origin = tuple(-value for value in origin) - pad = _extend_pad(structure.shape, origin) - pad_value = float(bool(border_value)) - - if mask is not None: - mask_flat = mask.to(dtype=torch.bool).reshape(batch * channels, 1, *spatial_shape) - input_flat = ( - (input != 0).to(dtype=torch.float32).reshape(batch * channels, 1, *spatial_shape) - ) - else: - mask_flat = None - input_flat = None - - if iterations_flag: - old = None - while True: - x_padded = F.pad(x, pad, value=pad_value) - conv = conv_nd(x_padded, kernel) - if mode == "erosion": - x = (conv == kernel_sum).to(dtype=torch.float32) - else: - x = (conv > 0).to(dtype=torch.float32) - if mask_flat is not None: - x = torch.where(mask_flat, x, input_flat) - - if old is not None and torch.equal(x, old): - break - - old = x.clone() - else: - for _ in range(iterations): - x_padded = F.pad(x, pad, value=pad_value) - conv = conv_nd(x_padded, kernel) - if mode == "erosion": - x = (conv == kernel_sum).to(dtype=torch.float32) - else: - x = (conv > 0).to(dtype=torch.float32) - - if mask_flat is not None: - x = torch.where(mask_flat, x, input_flat) - - result = x.reshape(batch, channels, *spatial_shape).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: - """ - N-dimensional binary erosion for `(B, C, Spatial...)` 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, -) -> Tensor: - """binary dilation for `(B, C, Spatial...)` tensors.""" - return _binary_morphology( - input, - structure, - iterations, - mask, - output, - border_value, - origin, - mode="dilation", - ) - - -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: - """binary opening for '(B, C, ...)' Tensors .""" - x = binary_erosion( - input, - structure, - iterations, - mask, - output, - 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: - """binary closing for (B, C, ...) Tensors.""" - x = binary_dilation( - input, - structure, - iterations, - mask, - output, - border_value, - origin, - ) - x = binary_erosion( - x, - structure, - iterations, - mask, - output, - border_value, - origin, - ) - return x +import torch +import torch.nn.functional as F +from torch import Tensor + +from .. import _C +from ._convnd import conv_nd +from .structure import generate_binary_structure + + +def _prepare_origin(origin: int | tuple[int, ...], ndim=int) -> tuple[int, ...]: + """change the origin into tuple""" + if isinstance(origin, int): + return (origin,) * ndim + origin = tuple(origin) + + if (len(origin)) != ndim: + raise ValueError(f"origin dimension is not {ndim}, got {len(origin)}") + + return origin + + +def _extend_pad(kernel_shape: torch.Size, origin: tuple[int, ...]) -> list[int]: + """extend the padlist for kernel""" + pad = [] + for dim in range(len(kernel_shape) - 1, -1, -1): + center = kernel_shape[dim] // 2 + pad_before = center + origin[dim] + pad_after = kernel_shape[dim] - 1 - pad_before + pad.extend([pad_before, pad_after]) + return pad + + +def _flip_structure(structure: Tensor) -> Tensor: + dim = tuple(range(structure.ndim)) + return torch.flip(structure, dim) + + +def _binary_morphology_cuda_step( + input: Tensor, + structure: Tensor, + border_value: bool, + origin: tuple[int, ...], + *, + mode: str, +) -> Tensor: + x = (input != 0).contiguous() + structure = (structure != 0).detach().to(device="cpu", dtype=torch.bool).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: + iterations_flag = iterations < 1 + + spatial_ndim = input.ndim - 2 + + if structure is None: + structure = generate_binary_structure(spatial_ndim, 1) + + batch, channels = input.shape[:2] + spatial_shape = input.shape[2:] + + origin = _prepare_origin(origin, spatial_ndim) + + if input.is_cuda: + x = input != 0 + input_bool = x + mask_bool = mask.to(device=input.device, dtype=torch.bool) if mask is not None else None + + def step(value: Tensor) -> Tensor: + return _binary_morphology_cuda_step( + value, + structure, + border_value, + origin, + mode=mode, + ) + + if iterations_flag: + 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.clone() + 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) + else: + x = (input != 0).to(dtype=torch.float32).reshape(batch * channels, 1, *spatial_shape) + structure = (structure != 0).to(device=input.device, dtype=torch.float32) + if mode == "dilation": + structure = _flip_structure(structure) + kernel = structure.unsqueeze(0).unsqueeze(0) + kernel_sum = kernel.sum() + + conv_origin = origin + if mode == "dilation": + conv_origin = tuple(-value for value in origin) + pad = _extend_pad(structure.shape, conv_origin) + pad_value = float(bool(border_value)) + + if mask is not None: + mask_flat = mask.to(dtype=torch.bool).reshape(batch * channels, 1, *spatial_shape) + input_flat = ( + (input != 0).to(dtype=torch.float32).reshape(batch * channels, 1, *spatial_shape) + ) + else: + mask_flat = None + input_flat = None + + if iterations_flag: + old = None + while True: + x_padded = F.pad(x, pad, value=pad_value) + conv = conv_nd(x_padded, kernel) + if mode == "erosion": + x = (conv == kernel_sum).to(dtype=torch.float32) + else: + x = (conv > 0).to(dtype=torch.float32) + if mask_flat is not None: + x = torch.where(mask_flat, x, input_flat) + + if old is not None and torch.equal(x, old): + break + + old = x.clone() + else: + for _ in range(iterations): + x_padded = F.pad(x, pad, value=pad_value) + conv = conv_nd(x_padded, kernel) + if mode == "erosion": + x = (conv == kernel_sum).to(dtype=torch.float32) + else: + x = (conv > 0).to(dtype=torch.float32) + + if mask_flat is not None: + x = torch.where(mask_flat, x, input_flat) + + result = x.reshape(batch, channels, *spatial_shape).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: + """ + N-dimensional binary erosion for `(B, C, Spatial...)` 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, +) -> Tensor: + """binary dilation for `(B, C, Spatial...)` tensors.""" + return _binary_morphology( + input, + structure, + iterations, + mask, + output, + border_value, + origin, + mode="dilation", + ) + + +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: + """binary opening for '(B, C, ...)' Tensors .""" + x = binary_erosion( + input, + structure, + iterations, + mask, + output, + 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: + """binary closing for (B, C, ...) Tensors.""" + x = binary_dilation( + input, + structure, + iterations, + mask, + output, + border_value, + origin, + ) + x = binary_erosion( + x, + structure, + iterations, + mask, + output, + border_value, + origin, + ) + return x From 5e6ba03152816643b60438576e1e5ddee87fe88e Mon Sep 17 00:00:00 2001 From: Yuhandeng Date: Tue, 7 Jul 2026 16:15:13 +0800 Subject: [PATCH 09/25] delete convnd --- ...alized-binary-morphology-cuda-primitive.md | 2 +- test/test_binary.py | 10 +- torchmorph/morphology/_convnd.py | 60 ------- torchmorph/morphology/binary.py | 151 +++++------------- 4 files changed, 48 insertions(+), 175 deletions(-) delete mode 100644 torchmorph/morphology/_convnd.py diff --git a/docs/adr/0001-generalized-binary-morphology-cuda-primitive.md b/docs/adr/0001-generalized-binary-morphology-cuda-primitive.md index 435a58a..776dd34 100644 --- a/docs/adr/0001-generalized-binary-morphology-cuda-primitive.md +++ b/docs/adr/0001-generalized-binary-morphology-cuda-primitive.md @@ -12,7 +12,7 @@ Binary erosion and binary dilation are also the base operations for binary openi ## Decision -Implement a generalized binary morphology CUDA primitive, modeled after the grey morphology CUDA primitive, and dispatch CUDA binary erosion and binary dilation through it. The primitive supports batch-channel tensors with arbitrary spatial rank up to the same practical maximum used by the grey implementation. Python remains responsible for iteration loops, mask application, output copying, and composite operators. +Implement a generalized binary morphology CUDA primitive, modeled after the grey morphology CUDA primitive, and dispatch CUDA binary erosion and binary dilation through it. The primitive supports batch-channel tensors with arbitrary spatial rank up to the same practical maximum used by the grey implementation. Python remains responsible for iteration loops, mask application, output copying, and composite operators. Binary morphology follows the CUDA-only public contract used by the other accelerated morphology and distance-transform operators. Do not start with specialized 2D or 3D kernels. Keep those as later optimizations once the generalized primitive is correct and benchmarked. diff --git a/test/test_binary.py b/test/test_binary.py index 4d7d737..b851783 100644 --- a/test/test_binary.py +++ b/test/test_binary.py @@ -152,11 +152,11 @@ def test_binary_morphology_output(torch_op, scipy_op, np_input, output_shape): torch.testing.assert_close(result.cpu(), torch.as_tensor(expected)) -@pytest.mark.parametrize(("torch_op", "scipy_op"), BINARY_OPERATORS) -def test_binary_morphology_cpu_fallback_matches_scipy(torch_op, scipy_op): - result = torch_op(torch.as_tensor(CASE_2D, dtype=torch.float32)) - expected = apply_scipy_to_batch(CASE_2D, scipy_op) - torch.testing.assert_close(result, torch.as_tensor(expected)) +@pytest.mark.parametrize("torch_op", TORCH_OPERATORS) +def test_binary_morphology_requires_cuda(torch_op): + x = torch.as_tensor(CASE_2D, dtype=torch.float32) + with pytest.raises(ValueError, match="CUDA"): + torch_op(x) @pytest.mark.parametrize("torch_op", TORCH_OPERATORS) diff --git a/torchmorph/morphology/_convnd.py b/torchmorph/morphology/_convnd.py deleted file mode 100644 index b1ee854..0000000 --- a/torchmorph/morphology/_convnd.py +++ /dev/null @@ -1,60 +0,0 @@ -import torch -from torch import Tensor - -_NATIVE_CONVS = { - 1: torch.nn.functional.conv1d, - 2: torch.nn.functional.conv2d, - 3: torch.nn.functional.conv3d, -} - - -def conv_nd(x: Tensor, weight: Tensor) -> Tensor: - """Apply correlation on tensors shaped as ``(N, C, Spatial...)``. - - This is a minimal private helper for ``torchmorph.morphology.binary``. - Inputs are expected to already be padded, use channel-first layout, and - convolve over the trailing spatial dimensions with stride=1, dilation=1, - and groups=1. - """ - - num_spatial = x.ndim - 2 - if num_spatial < 1: - raise ValueError(f"expected at least 1 spatial dim, got input ndim={x.ndim}") - if weight.ndim != x.ndim: - raise ValueError( - f"expected weight ndim {x.ndim} to match input ndim {x.ndim}, got {weight.ndim}" - ) - - return _conv_core(x, weight) - - -def _conv_core(x: Tensor, weight: Tensor) -> Tensor: - num_spatial = x.ndim - 2 - if num_spatial in _NATIVE_CONVS: - return _NATIVE_CONVS[num_spatial](x, weight, bias=None, stride=1, padding=0, dilation=1) - return _conv_recursive(x, weight) - - -def _conv_recursive(x: Tensor, weight: Tensor) -> Tensor: - kernel_size = weight.shape[2] - input_size = x.shape[2] - output_size = input_size - kernel_size + 1 - batch_size = x.shape[0] - - accumulated = None - for kernel_index in range(kernel_size): - x_slice = x[:, :, kernel_index : kernel_index + output_size] - flattened = x_slice.moveaxis(2, 1).reshape( - batch_size * output_size, - x.shape[1], - *x_slice.shape[3:], - ) - - partial = _conv_core(flattened, weight[:, :, kernel_index]) - partial = partial.reshape(batch_size, output_size, partial.shape[1], *partial.shape[2:]) - partial = partial.moveaxis(2, 1) - accumulated = partial if accumulated is None else accumulated + partial - - if accumulated is None: - raise RuntimeError("recursive convolution produced no output") - return accumulated diff --git a/torchmorph/morphology/binary.py b/torchmorph/morphology/binary.py index dfeeea7..9ecf0d5 100644 --- a/torchmorph/morphology/binary.py +++ b/torchmorph/morphology/binary.py @@ -1,9 +1,7 @@ import torch -import torch.nn.functional as F from torch import Tensor from .. import _C -from ._convnd import conv_nd from .structure import generate_binary_structure @@ -19,22 +17,6 @@ def _prepare_origin(origin: int | tuple[int, ...], ndim=int) -> tuple[int, ...]: return origin -def _extend_pad(kernel_shape: torch.Size, origin: tuple[int, ...]) -> list[int]: - """extend the padlist for kernel""" - pad = [] - for dim in range(len(kernel_shape) - 1, -1, -1): - center = kernel_shape[dim] // 2 - pad_before = center + origin[dim] - pad_after = kernel_shape[dim] - 1 - pad_before - pad.extend([pad_before, pad_after]) - return pad - - -def _flip_structure(structure: Tensor) -> Tensor: - dim = tuple(range(structure.ndim)) - return torch.flip(structure, dim) - - def _binary_morphology_cuda_step( input: Tensor, structure: Tensor, @@ -60,102 +42,53 @@ def _binary_morphology( *, mode: str, ) -> Tensor: - iterations_flag = iterations < 1 + if not input.is_cuda: + raise ValueError("Input tensor must be on CUDA device.") + if input.ndim < 3: + raise ValueError( + f"Input must be (B, C, Spatial...) with at least 3 dimensions, got {input.shape}." + ) + if input.numel() == 0: + raise ValueError(f"Invalid input: empty tensor with shape {input.shape}.") + iterations_flag = iterations < 1 spatial_ndim = input.ndim - 2 if structure is None: structure = generate_binary_structure(spatial_ndim, 1) - batch, channels = input.shape[:2] - spatial_shape = input.shape[2:] - origin = _prepare_origin(origin, spatial_ndim) - - if input.is_cuda: - x = input != 0 - input_bool = x - mask_bool = mask.to(device=input.device, dtype=torch.bool) if mask is not None else None - - def step(value: Tensor) -> Tensor: - return _binary_morphology_cuda_step( - value, - structure, - border_value, - origin, - mode=mode, - ) - - if iterations_flag: - 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.clone() - 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) + x = input != 0 + input_bool = x + mask_bool = mask.to(device=input.device, dtype=torch.bool) if mask is not None else None + + def step(value: Tensor) -> Tensor: + return _binary_morphology_cuda_step( + value, + structure, + border_value, + origin, + mode=mode, + ) + + if iterations_flag: + 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.clone() else: - x = (input != 0).to(dtype=torch.float32).reshape(batch * channels, 1, *spatial_shape) - structure = (structure != 0).to(device=input.device, dtype=torch.float32) - if mode == "dilation": - structure = _flip_structure(structure) - kernel = structure.unsqueeze(0).unsqueeze(0) - kernel_sum = kernel.sum() - - conv_origin = origin - if mode == "dilation": - conv_origin = tuple(-value for value in origin) - pad = _extend_pad(structure.shape, conv_origin) - pad_value = float(bool(border_value)) - - if mask is not None: - mask_flat = mask.to(dtype=torch.bool).reshape(batch * channels, 1, *spatial_shape) - input_flat = ( - (input != 0).to(dtype=torch.float32).reshape(batch * channels, 1, *spatial_shape) - ) - else: - mask_flat = None - input_flat = None - - if iterations_flag: - old = None - while True: - x_padded = F.pad(x, pad, value=pad_value) - conv = conv_nd(x_padded, kernel) - if mode == "erosion": - x = (conv == kernel_sum).to(dtype=torch.float32) - else: - x = (conv > 0).to(dtype=torch.float32) - if mask_flat is not None: - x = torch.where(mask_flat, x, input_flat) - - if old is not None and torch.equal(x, old): - break - - old = x.clone() - else: - for _ in range(iterations): - x_padded = F.pad(x, pad, value=pad_value) - conv = conv_nd(x_padded, kernel) - if mode == "erosion": - x = (conv == kernel_sum).to(dtype=torch.float32) - else: - x = (conv > 0).to(dtype=torch.float32) - - if mask_flat is not None: - x = torch.where(mask_flat, x, input_flat) + for _ in range(iterations): + x = step(x) + if mask_bool is not None: + x = torch.where(mask_bool, x, input_bool) - result = x.reshape(batch, channels, *spatial_shape).to(dtype=torch.bool) + result = x.to(dtype=torch.bool) if output is not None: output.copy_(result) return output @@ -172,7 +105,7 @@ def binary_erosion( origin: int | tuple[int, ...] = 0, ) -> Tensor: """ - N-dimensional binary erosion for `(B, C, Spatial...)` tensors. + N-dimensional binary erosion for `(B, C, Spatial...)` CUDA tensors. For a single image or volume, add batch and channel dimensions first. """ @@ -197,7 +130,7 @@ def binary_dilation( border_value: bool = False, origin: int | tuple[int, ...] = 0, ) -> Tensor: - """binary dilation for `(B, C, Spatial...)` tensors.""" + """N-dimensional binary dilation for `(B, C, Spatial...)` CUDA tensors.""" return _binary_morphology( input, structure, @@ -219,7 +152,7 @@ def binary_opening( border_value: bool = False, origin: int | tuple[int, ...] = 0, ) -> Tensor: - """binary opening for '(B, C, ...)' Tensors .""" + """N-dimensional binary opening for `(B, C, Spatial...)` CUDA tensors.""" x = binary_erosion( input, structure, @@ -250,7 +183,7 @@ def binary_closing( border_value: bool = False, origin: int | tuple[int, ...] = 0, ) -> Tensor: - """binary closing for (B, C, ...) Tensors.""" + """N-dimensional binary closing for `(B, C, Spatial...)` CUDA tensors.""" x = binary_dilation( input, structure, From 10b03ed7bfb89a1285af89104801baaeaf1d9192 Mon Sep 17 00:00:00 2001 From: Yuhandeng Date: Tue, 7 Jul 2026 16:41:25 +0800 Subject: [PATCH 10/25] fix origin bug --- test/test_binary.py | 18 ++++++++++++++++++ torchmorph/morphology/binary.py | 24 ++++++------------------ torchmorph/morphology/structure.py | 16 ++++++++-------- 3 files changed, 32 insertions(+), 26 deletions(-) diff --git a/test/test_binary.py b/test/test_binary.py index b851783..254a07d 100644 --- a/test/test_binary.py +++ b/test/test_binary.py @@ -152,6 +152,24 @@ def test_binary_morphology_output(torch_op, scipy_op, np_input, output_shape): torch.testing.assert_close(result.cpu(), torch.as_tensor(expected)) +@pytest.mark.parametrize( + ("torch_op", "scipy_op"), + [ + pytest.param(tm.binary_opening, scipy_binary_opening, id="opening"), + pytest.param(tm.binary_closing, scipy_binary_closing, id="closing"), + ], +) +def test_binary_composite_output_contains_only_final_result(torch_op, scipy_op): + x = torch.as_tensor(CASE_2D, dtype=torch.float32, device="cuda") + output = torch.ones_like(x, dtype=torch.bool) + + result = torch_op(x, output=output) + + assert result is output + expected = apply_scipy_to_batch(CASE_2D, scipy_op, output=np.empty((3, 3), dtype=bool)) + torch.testing.assert_close(output.cpu(), torch.as_tensor(expected)) + + @pytest.mark.parametrize("torch_op", TORCH_OPERATORS) def test_binary_morphology_requires_cuda(torch_op): x = torch.as_tensor(CASE_2D, dtype=torch.float32) diff --git a/torchmorph/morphology/binary.py b/torchmorph/morphology/binary.py index 9ecf0d5..b494c4c 100644 --- a/torchmorph/morphology/binary.py +++ b/torchmorph/morphology/binary.py @@ -2,19 +2,7 @@ from torch import Tensor from .. import _C -from .structure import generate_binary_structure - - -def _prepare_origin(origin: int | tuple[int, ...], ndim=int) -> tuple[int, ...]: - """change the origin into tuple""" - if isinstance(origin, int): - return (origin,) * ndim - origin = tuple(origin) - - if (len(origin)) != ndim: - raise ValueError(f"origin dimension is not {ndim}, got {len(origin)}") - - return origin +from .structure import _normalize_origin, generate_binary_structure def _binary_morphology_cuda_step( @@ -51,13 +39,13 @@ def _binary_morphology( if input.numel() == 0: raise ValueError(f"Invalid input: empty tensor with shape {input.shape}.") - iterations_flag = iterations < 1 + iterate_until_stable = iterations < 1 spatial_ndim = input.ndim - 2 if structure is None: structure = generate_binary_structure(spatial_ndim, 1) - origin = _prepare_origin(origin, spatial_ndim) + origin = _normalize_origin(origin, spatial_ndim) x = input != 0 input_bool = x mask_bool = mask.to(device=input.device, dtype=torch.bool) if mask is not None else None @@ -71,7 +59,7 @@ def step(value: Tensor) -> Tensor: mode=mode, ) - if iterations_flag: + if iterate_until_stable: old = None while True: x = step(x) @@ -158,7 +146,7 @@ def binary_opening( structure, iterations, mask, - output, + None, border_value, origin, ) @@ -189,7 +177,7 @@ def binary_closing( structure, iterations, mask, - output, + None, border_value, origin, ) diff --git a/torchmorph/morphology/structure.py b/torchmorph/morphology/structure.py index 573f29b..ac42a77 100644 --- a/torchmorph/morphology/structure.py +++ b/torchmorph/morphology/structure.py @@ -39,14 +39,14 @@ def _iterate_generated_structure(rank: int, connectivity: int, iterations: int, return offsets.abs().sum(dim=0) <= connectivity * iterations -def _prepare_origin(origin: int | tuple[int, ...], ndim: int) -> list[int]: +def _normalize_origin(origin: int | tuple[int, ...], ndim: int) -> tuple[int, ...]: if isinstance(origin, int): - return [origin] * ndim + return (origin,) * ndim - origin_list = list(origin) - if len(origin_list) != ndim: - raise ValueError(f"origin dimension is not {ndim}, got {len(origin_list)}") - return origin_list + origin_tuple = tuple(origin) + if len(origin_tuple) != ndim: + raise ValueError(f"origin dimension is not {ndim}, got {len(origin_tuple)}") + return origin_tuple def iterate_structure( @@ -85,5 +85,5 @@ def iterate_structure( if origin is None: return result - origin_list = _prepare_origin(origin, structure.ndim) - return result, [value * iterations for value in origin_list] + origin_tuple = _normalize_origin(origin, structure.ndim) + return result, [value * iterations for value in origin_tuple] From 6c55e71d93042c1fe34b7fe8cd16cdee76344d00 Mon Sep 17 00:00:00 2001 From: Yuhandeng Date: Wed, 8 Jul 2026 14:39:09 +0800 Subject: [PATCH 11/25] binary propagation --- benchmark/binary.py | 1 + test/test_binary.py | 74 +++++++++++++++++++++++++++++++ torchmorph/__init__.py | 2 + torchmorph/morphology/__init__.py | 9 +++- torchmorph/morphology/binary.py | 21 +++++++++ 5 files changed, 106 insertions(+), 1 deletion(-) diff --git a/benchmark/binary.py b/benchmark/binary.py index b519550..f95a15f 100644 --- a/benchmark/binary.py +++ b/benchmark/binary.py @@ -16,6 +16,7 @@ "dilation": (ndi.binary_dilation, tm.binary_dilation), "opening": (ndi.binary_opening, tm.binary_opening), "closing": (ndi.binary_closing, tm.binary_closing), + "propagation": (ndi.binary_propagation, tm.binary_propagation), } diff --git a/test/test_binary.py b/test/test_binary.py index 254a07d..7a66000 100644 --- a/test/test_binary.py +++ b/test/test_binary.py @@ -5,6 +5,7 @@ from scipy.ndimage import binary_dilation as scipy_binary_dilation from scipy.ndimage import binary_erosion as scipy_binary_erosion from scipy.ndimage import binary_opening as scipy_binary_opening +from scipy.ndimage import binary_propagation as scipy_binary_propagation from scipy.ndimage import generate_binary_structure import torchmorph as tm @@ -20,11 +21,16 @@ pytest.param(tm.binary_closing, scipy_binary_closing, id="closing"), ] +PROPAGATION_OPERATORS = [ + pytest.param(tm.binary_propagation, scipy_binary_propagation, id="propagation"), +] + TORCH_OPERATORS = [ pytest.param(tm.binary_erosion, id="erosion"), pytest.param(tm.binary_dilation, id="dilation"), pytest.param(tm.binary_opening, id="opening"), pytest.param(tm.binary_closing, id="closing"), + pytest.param(tm.binary_propagation, id="propagation"), ] CASE_2D = np.array( @@ -132,6 +138,54 @@ def test_binary_morphology_mask(torch_op, scipy_op, np_input, mask): torch.testing.assert_close(result.cpu(), torch.as_tensor(expected)) +@pytest.mark.parametrize(("torch_op", "scipy_op"), PROPAGATION_OPERATORS) +@pytest.mark.parametrize( + ("np_input", "structure", "origin", "border_value"), + [ + pytest.param(CASE_2D, None, 0, False, id="2d_basic"), + pytest.param(CASE_2D, STRUCTURE_2D, 1, True, id="2d_origin_border"), + pytest.param(CASE_3D, STRUCTURE_3D_1, 0, False, id="3d_structure_1"), + pytest.param(CASE_3D, STRUCTURE_3D_2, 1, True, id="3d_origin_border"), + pytest.param(CASE_4D, STRUCTURE_4D, 0, False, id="4d_structure"), + ], +) +def test_binary_propagation_matches_scipy( + torch_op, scipy_op, np_input, structure, origin, border_value +): + result = torch_op( + torch.as_tensor(np_input, dtype=torch.float32, device="cuda"), + structure=optional_cuda_tensor(structure, torch.bool), + origin=origin, + border_value=border_value, + ) + expected = apply_scipy_to_batch( + np_input, + scipy_op, + structure=structure, + origin=origin, + border_value=border_value, + ) + torch.testing.assert_close(result.cpu(), torch.as_tensor(expected)) + + +@pytest.mark.parametrize(("torch_op", "scipy_op"), PROPAGATION_OPERATORS) +@pytest.mark.parametrize( + ("np_input", "mask"), + [ + pytest.param(CASE_2D, MASK_2D, id="2d_mask"), + pytest.param(CASE_3D, MASK_3D, id="3d_mask"), + pytest.param(CASE_4D, MASK_4D, id="4d_mask"), + ], +) +def test_binary_propagation_mask(torch_op, scipy_op, np_input, mask): + result = torch_op( + torch.as_tensor(np_input, dtype=torch.float32, device="cuda"), + mask=cuda_mask(mask, np_input), + ) + expected = apply_scipy_to_batch(np_input, scipy_op, mask=mask) + torch.testing.assert_close(result.cpu(), torch.as_tensor(expected)) + + @pytest.mark.parametrize(("torch_op", "scipy_op"), BINARY_OPERATORS) @pytest.mark.parametrize( ("np_input", "output_shape"), @@ -152,6 +206,26 @@ def test_binary_morphology_output(torch_op, scipy_op, np_input, output_shape): torch.testing.assert_close(result.cpu(), torch.as_tensor(expected)) +@pytest.mark.parametrize(("torch_op", "scipy_op"), PROPAGATION_OPERATORS) +@pytest.mark.parametrize( + ("np_input", "output_shape"), + [ + pytest.param(CASE_2D, (3, 3), id="2d_output"), + pytest.param(CASE_3D, (5, 5, 5), id="3d_output"), + pytest.param(CASE_4D, (4, 4, 4, 4), id="4d_output"), + ], +) +def test_binary_propagation_output(torch_op, scipy_op, np_input, output_shape): + x = torch.as_tensor(np_input, dtype=torch.float32, device="cuda") + output = torch.empty_like(x, dtype=torch.bool) + + result = torch_op(x, output=output) + + assert result is output + expected = apply_scipy_to_batch(np_input, scipy_op, output=np.empty(output_shape, dtype=bool)) + torch.testing.assert_close(result.cpu(), torch.as_tensor(expected)) + + @pytest.mark.parametrize( ("torch_op", "scipy_op"), [ diff --git a/torchmorph/__init__.py b/torchmorph/__init__.py index 559cb0d..e72652d 100644 --- a/torchmorph/__init__.py +++ b/torchmorph/__init__.py @@ -9,6 +9,7 @@ binary_dilation, binary_erosion, binary_opening, + binary_propagation, black_tophat, generate_binary_structure, grey_closing, @@ -32,6 +33,7 @@ "binary_dilation", "binary_opening", "binary_closing", + "binary_propagation", "grey_erosion", "grey_dilation", "grey_opening", diff --git a/torchmorph/morphology/__init__.py b/torchmorph/morphology/__init__.py index bf6dea0..0fc61c9 100644 --- a/torchmorph/morphology/__init__.py +++ b/torchmorph/morphology/__init__.py @@ -1,4 +1,10 @@ -from .binary import binary_closing, binary_dilation, binary_erosion, binary_opening +from .binary import ( + binary_closing, + binary_dilation, + binary_erosion, + binary_opening, + binary_propagation, +) from .grey import ( black_tophat, grey_closing, @@ -18,6 +24,7 @@ "binary_dilation", "binary_opening", "binary_closing", + "binary_propagation", "grey_erosion", "grey_dilation", "grey_opening", diff --git a/torchmorph/morphology/binary.py b/torchmorph/morphology/binary.py index b494c4c..e0ce3f5 100644 --- a/torchmorph/morphology/binary.py +++ b/torchmorph/morphology/binary.py @@ -131,6 +131,27 @@ def binary_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_opening( input: Tensor, structure: Tensor | None = None, From 49937ae2e93c325a6b26e1ce539dff236cb157bf Mon Sep 17 00:00:00 2001 From: Yuhandeng Date: Wed, 8 Jul 2026 16:45:32 +0800 Subject: [PATCH 12/25] format --- torchmorph/csrc/torchmorph.cpp | 1 + 1 file changed, 1 insertion(+) diff --git a/torchmorph/csrc/torchmorph.cpp b/torchmorph/csrc/torchmorph.cpp index dcda804..1f71d8c 100644 --- a/torchmorph/csrc/torchmorph.cpp +++ b/torchmorph/csrc/torchmorph.cpp @@ -96,6 +96,7 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { py::arg("return_distances") = true, py::arg("return_indices") = false, py::arg("algorithm") = "exact"); + m.def("cdt_cuda", &cdt_cuda, "Chessboard/Manhattan distance transform", py::arg("input"), From 5ba4ae5272c5f0ff0475c54da166b8a844e7cf06 Mon Sep 17 00:00:00 2001 From: Yuhandeng Date: Wed, 8 Jul 2026 17:58:59 +0800 Subject: [PATCH 13/25] update readme --- README.md | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/README.md b/README.md index 4d4f69d..ae34b3f 100644 --- a/README.md +++ b/README.md @@ -30,7 +30,7 @@ The instructions in this section are for **building TorchMorph from this reposit 1. Create and activate a fresh conda environment: ```bash - conda create -n torchmorph python=3.11 -y + conda create -n torchmorph python=3.12 -y conda activate torchmorph ``` @@ -41,7 +41,8 @@ The instructions in this section are for **building TorchMorph from this reposit 3. Ensure you have a compatible `nvcc` compiler. ```bash - conda install -c nvidia cuda-nvcc=12.4 -y + conda install -c nvidia -c conda-forge \ + cuda-version=12.4 cuda-cudart=12.4 cuda-cudart-dev=12.4 cuda-nvcc=12.4 cuda-cccl=12.4 ``` 4. Install TorchMorph's dependencies and build the extension: From 83effc871fb5a9eac896b090b3c7ebaaf42d629c Mon Sep 17 00:00:00 2001 From: Yuhandeng Date: Wed, 8 Jul 2026 18:01:55 +0800 Subject: [PATCH 14/25] readme --- README.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/README.md b/README.md index ae34b3f..ede098a 100644 --- a/README.md +++ b/README.md @@ -42,7 +42,7 @@ The instructions in this section are for **building TorchMorph from this reposit 3. Ensure you have a compatible `nvcc` compiler. ```bash conda install -c nvidia -c conda-forge \ - cuda-version=12.4 cuda-cudart=12.4 cuda-cudart-dev=12.4 cuda-nvcc=12.4 cuda-cccl=12.4 + cuda-version=12.4 cuda-cudart=12.4 cuda-cudart-dev=12.4 cuda-nvcc=12.4 cuda-cccl=12.4 ``` 4. Install TorchMorph's dependencies and build the extension: From 2b67bb605ff9a3a947924ce7185a051084e2eef6 Mon Sep 17 00:00:00 2001 From: Yuhandeng Date: Wed, 8 Jul 2026 19:57:13 +0800 Subject: [PATCH 15/25] binary_fill_hole and binary_hit_or_miss --- CONTEXT.md | 8 ++ benchmark/binary.py | 2 + test/test_binary.py | 153 +++++++++++++++++++++++++++++- torchmorph/__init__.py | 4 + torchmorph/morphology/__init__.py | 4 + torchmorph/morphology/binary.py | 95 +++++++++++++++++-- 6 files changed, 252 insertions(+), 14 deletions(-) diff --git a/CONTEXT.md b/CONTEXT.md index d849f48..12bfdac 100644 --- a/CONTEXT.md +++ b/CONTEXT.md @@ -12,6 +12,14 @@ _Avoid_: filter, convolution Morphological operations where every non-zero input value is treated as foreground and zero is treated as background. _Avoid_: boolean convolution, mask filtering +**Binary Hole Filling**: +A binary morphology operation that turns background regions not connected to the sample boundary into foreground. +_Avoid_: hole, fill hole + +**Binary Hit-or-Miss**: +A binary morphology operation that detects locations where foreground and background match paired structuring elements. +_Avoid_: hit miss, pattern mask + **Grey Morphology**: Morphological operations over numeric intensities, where the result depends on ordered values rather than only foreground/background membership. _Avoid_: grayscale filter diff --git a/benchmark/binary.py b/benchmark/binary.py index f95a15f..f1bdbd7 100644 --- a/benchmark/binary.py +++ b/benchmark/binary.py @@ -14,6 +14,8 @@ BINARY_OPERATORS = { "erosion": (ndi.binary_erosion, tm.binary_erosion), "dilation": (ndi.binary_dilation, tm.binary_dilation), + "fill_holes": (ndi.binary_fill_holes, tm.binary_fill_holes), + "hit_or_miss": (ndi.binary_hit_or_miss, tm.binary_hit_or_miss), "opening": (ndi.binary_opening, tm.binary_opening), "closing": (ndi.binary_closing, tm.binary_closing), "propagation": (ndi.binary_propagation, tm.binary_propagation), diff --git a/test/test_binary.py b/test/test_binary.py index 7a66000..201fc7f 100644 --- a/test/test_binary.py +++ b/test/test_binary.py @@ -4,6 +4,8 @@ from scipy.ndimage import binary_closing as scipy_binary_closing from scipy.ndimage import binary_dilation as scipy_binary_dilation from scipy.ndimage import binary_erosion as scipy_binary_erosion +from scipy.ndimage import binary_fill_holes as scipy_binary_fill_holes +from scipy.ndimage import binary_hit_or_miss as scipy_binary_hit_or_miss from scipy.ndimage import binary_opening as scipy_binary_opening from scipy.ndimage import binary_propagation as scipy_binary_propagation from scipy.ndimage import generate_binary_structure @@ -25,12 +27,26 @@ pytest.param(tm.binary_propagation, scipy_binary_propagation, id="propagation"), ] +FILL_HOLES_OPERATORS = [ + pytest.param(tm.binary_fill_holes, scipy_binary_fill_holes, id="fill_holes"), +] + +HIT_OR_MISS_OPERATORS = [ + pytest.param(tm.binary_hit_or_miss, scipy_binary_hit_or_miss, id="hit_or_miss"), +] + TORCH_OPERATORS = [ pytest.param(tm.binary_erosion, id="erosion"), pytest.param(tm.binary_dilation, id="dilation"), pytest.param(tm.binary_opening, id="opening"), pytest.param(tm.binary_closing, id="closing"), pytest.param(tm.binary_propagation, id="propagation"), + pytest.param(tm.binary_fill_holes, id="fill_holes"), +] + +ALL_TORCH_OPERATORS = [ + *TORCH_OPERATORS, + pytest.param(tm.binary_hit_or_miss, id="hit_or_miss"), ] CASE_2D = np.array( @@ -42,6 +58,16 @@ CASE_4D = np.zeros((2, 1, 4, 4, 4, 4), dtype=bool) CASE_4D[0, 0, 1:3, 1:2, 2:3] = True +CASE_HOLES_2D = np.zeros((1, 1, 5, 5), dtype=bool) +CASE_HOLES_2D[0, 0, 1:4, 1:4] = True +CASE_HOLES_2D[0, 0, 2, 2] = False + +CASE_HOLES_3D = np.zeros((1, 1, 5, 5, 5), dtype=bool) +CASE_HOLES_3D[0, 0, 1:4, 1:4, 1:4] = True +CASE_HOLES_3D[0, 0, 2, 2, 2] = False + +CASE_HIT_3D = generate_binary_structure(rank=3, connectivity=1)[None, None, ...] + STRUCTURE_2D = generate_binary_structure(rank=2, connectivity=1) STRUCTURE_3D_1 = generate_binary_structure(rank=3, connectivity=2) STRUCTURE_3D_2 = generate_binary_structure(rank=3, connectivity=3) @@ -64,11 +90,15 @@ def apply_scipy_to_batch(np_input, scipy_op, **kwargs): results = [] for sample in samples: sample_kwargs = dict(kwargs) + sample_output = None if mask is not None: sample_kwargs["mask"] = mask if output is not None: - sample_kwargs["output"] = np.empty_like(output) + sample_output = np.empty_like(output) + sample_kwargs["output"] = sample_output result = scipy_op(sample, **sample_kwargs) + if result is None: + result = sample_output results.append(np.asarray(result).copy()) return np.stack(results).reshape(*batch_shape, *spatial_shape) @@ -226,6 +256,109 @@ def test_binary_propagation_output(torch_op, scipy_op, np_input, output_shape): torch.testing.assert_close(result.cpu(), torch.as_tensor(expected)) +@pytest.mark.parametrize(("torch_op", "scipy_op"), FILL_HOLES_OPERATORS) +@pytest.mark.parametrize( + ("np_input", "structure", "origin"), + [ + pytest.param(CASE_HOLES_2D, None, 0, id="2d_default"), + pytest.param(CASE_HOLES_2D, STRUCTURE_2D, 0, id="2d_structure"), + pytest.param(CASE_HOLES_3D, None, 0, id="3d_default"), + ], +) +def test_binary_fill_holes_matches_scipy(torch_op, scipy_op, np_input, structure, origin): + result = torch_op( + torch.as_tensor(np_input, dtype=torch.float32, device="cuda"), + structure=optional_cuda_tensor(structure, torch.bool), + origin=origin, + ) + expected = apply_scipy_to_batch( + np_input, + scipy_op, + structure=structure, + origin=origin, + ) + torch.testing.assert_close(result.cpu(), torch.as_tensor(expected)) + + +@pytest.mark.parametrize(("torch_op", "scipy_op"), FILL_HOLES_OPERATORS) +@pytest.mark.parametrize( + ("np_input", "output_shape"), + [ + pytest.param(CASE_HOLES_2D, (5, 5), id="2d_output"), + pytest.param(CASE_HOLES_3D, (5, 5, 5), id="3d_output"), + ], +) +def test_binary_fill_holes_output(torch_op, scipy_op, np_input, output_shape): + x = torch.as_tensor(np_input, dtype=torch.float32, device="cuda") + output = torch.empty_like(x, dtype=torch.bool) + + result = torch_op(x, output=output) + + assert result is output + expected = apply_scipy_to_batch(np_input, scipy_op, output=np.empty(output_shape, dtype=bool)) + torch.testing.assert_close(result.cpu(), torch.as_tensor(expected)) + + +@pytest.mark.parametrize(("torch_op", "scipy_op"), HIT_OR_MISS_OPERATORS) +@pytest.mark.parametrize( + ("np_input", "structure1", "structure2", "origin1", "origin2"), + [ + pytest.param(CASE_2D, None, None, 0, None, id="2d_default"), + pytest.param( + CASE_2D, + STRUCTURE_2D, + np.logical_not(STRUCTURE_2D), + 0, + None, + id="2d_structures", + ), + pytest.param(CASE_HIT_3D, None, None, 0, None, id="3d_default"), + ], +) +def test_binary_hit_or_miss_matches_scipy( + torch_op, scipy_op, np_input, structure1, structure2, origin1, origin2 +): + result = torch_op( + torch.as_tensor(np_input, dtype=torch.float32, device="cuda"), + structure1=optional_cuda_tensor(structure1, torch.bool), + structure2=optional_cuda_tensor(structure2, torch.bool), + origin1=origin1, + origin2=origin2, + ) + expected = apply_scipy_to_batch( + np_input, + scipy_op, + structure1=structure1, + structure2=structure2, + origin1=origin1, + origin2=origin2, + ) + torch.testing.assert_close(result.cpu(), torch.as_tensor(expected)) + + +@pytest.mark.parametrize(("torch_op", "scipy_op"), HIT_OR_MISS_OPERATORS) +def test_binary_hit_or_miss_output(torch_op, scipy_op): + x = torch.as_tensor(CASE_2D, dtype=torch.float32, device="cuda") + output = torch.empty_like(x, dtype=torch.bool) + + result = torch_op(x, output=output) + + assert result is output + expected = apply_scipy_to_batch(CASE_2D, scipy_op, output=np.empty((3, 3), dtype=bool)) + torch.testing.assert_close(result.cpu(), torch.as_tensor(expected)) + + +def test_binary_hit_or_miss_allows_empty_miss_structure(): + x = torch.ones((1, 1, 3, 3), dtype=torch.bool, device="cuda") + structure1 = torch.ones((3, 3), dtype=torch.bool, device="cuda") + structure2 = torch.zeros((3, 3), dtype=torch.bool, device="cuda") + + result = tm.binary_hit_or_miss(x, structure1=structure1, structure2=structure2) + expected = tm.binary_erosion(x, structure=structure1) + + torch.testing.assert_close(result, expected) + + @pytest.mark.parametrize( ("torch_op", "scipy_op"), [ @@ -244,7 +377,7 @@ def test_binary_composite_output_contains_only_final_result(torch_op, scipy_op): torch.testing.assert_close(output.cpu(), torch.as_tensor(expected)) -@pytest.mark.parametrize("torch_op", TORCH_OPERATORS) +@pytest.mark.parametrize("torch_op", ALL_TORCH_OPERATORS) def test_binary_morphology_requires_cuda(torch_op): x = torch.as_tensor(CASE_2D, dtype=torch.float32) with pytest.raises(ValueError, match="CUDA"): @@ -258,7 +391,19 @@ def test_binary_morphology_rejects_invalid_origin_dimension(torch_op): torch_op(x, origin=(0, 0, 0)) -@pytest.mark.parametrize("torch_op", TORCH_OPERATORS) +def test_binary_hit_or_miss_rejects_invalid_origin1_dimension(): + x = torch.as_tensor(CASE_2D, dtype=torch.float32, device="cuda") + with pytest.raises(ValueError, match="origin dimension"): + tm.binary_hit_or_miss(x, origin1=(0, 0, 0)) + + +def test_binary_hit_or_miss_rejects_invalid_origin2_dimension(): + x = torch.as_tensor(CASE_2D, dtype=torch.float32, device="cuda") + with pytest.raises(ValueError, match="origin dimension"): + tm.binary_hit_or_miss(x, origin2=(0, 0, 0)) + + +@pytest.mark.parametrize("torch_op", ALL_TORCH_OPERATORS) def test_binary_morphology_uses_current_cuda_stream(torch_op): stream = torch.cuda.Stream() with torch.cuda.stream(stream): @@ -270,7 +415,7 @@ def test_binary_morphology_uses_current_cuda_stream(torch_op): @pytest.mark.skipif(torch.cuda.device_count() < 2, reason="requires two CUDA devices") -@pytest.mark.parametrize("torch_op", TORCH_OPERATORS) +@pytest.mark.parametrize("torch_op", ALL_TORCH_OPERATORS) def test_binary_morphology_uses_input_cuda_device(torch_op): with torch.cuda.device(0): x = torch.as_tensor(CASE_2D, dtype=torch.float32, device="cuda:1") diff --git a/torchmorph/__init__.py b/torchmorph/__init__.py index 6c0ba38..4f263cf 100644 --- a/torchmorph/__init__.py +++ b/torchmorph/__init__.py @@ -7,6 +7,8 @@ binary_closing, binary_dilation, binary_erosion, + binary_fill_holes, + binary_hit_or_miss, binary_opening, binary_propagation, black_tophat, @@ -29,6 +31,8 @@ "iterate_structure", "binary_erosion", "binary_dilation", + "binary_fill_holes", + "binary_hit_or_miss", "binary_opening", "binary_closing", "binary_propagation", diff --git a/torchmorph/morphology/__init__.py b/torchmorph/morphology/__init__.py index 0fc61c9..4388484 100644 --- a/torchmorph/morphology/__init__.py +++ b/torchmorph/morphology/__init__.py @@ -2,6 +2,8 @@ binary_closing, binary_dilation, binary_erosion, + binary_fill_holes, + binary_hit_or_miss, binary_opening, binary_propagation, ) @@ -22,6 +24,8 @@ "iterate_structure", "binary_erosion", "binary_dilation", + "binary_fill_holes", + "binary_hit_or_miss", "binary_opening", "binary_closing", "binary_propagation", diff --git a/torchmorph/morphology/binary.py b/torchmorph/morphology/binary.py index e0ce3f5..5956537 100644 --- a/torchmorph/morphology/binary.py +++ b/torchmorph/morphology/binary.py @@ -5,6 +5,19 @@ from .structure import _normalize_origin, generate_binary_structure +def _validate_binary_input(input: Tensor) -> int: + if not input.is_cuda: + raise ValueError("Input tensor must be on CUDA device.") + if input.ndim < 3: + raise ValueError( + f"Input must be (B, C, Spatial...) with at least 3 dimensions, got {input.shape}." + ) + if input.numel() == 0: + raise ValueError(f"Invalid input: empty tensor with shape {input.shape}.") + + return input.ndim - 2 + + def _binary_morphology_cuda_step( input: Tensor, structure: Tensor, @@ -30,17 +43,8 @@ def _binary_morphology( *, mode: str, ) -> Tensor: - if not input.is_cuda: - raise ValueError("Input tensor must be on CUDA device.") - if input.ndim < 3: - raise ValueError( - f"Input must be (B, C, Spatial...) with at least 3 dimensions, got {input.shape}." - ) - if input.numel() == 0: - raise ValueError(f"Invalid input: empty tensor with shape {input.shape}.") - iterate_until_stable = iterations < 1 - spatial_ndim = input.ndim - 2 + spatial_ndim = _validate_binary_input(input) if structure is None: structure = generate_binary_structure(spatial_ndim, 1) @@ -152,6 +156,77 @@ def binary_propagation( ) +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_binary_input(input) + + 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_binary_input(input) + 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) + elif structure1.ndim != spatial_ndim: + raise ValueError(f"structure1 dimension is not {spatial_ndim}, got {structure1.ndim}") + + structure1 = structure1.detach().to(device="cpu", dtype=torch.bool).contiguous() + if structure2 is None: + structure2 = torch.logical_not(structure1).contiguous() + else: + if structure2.ndim != spatial_ndim: + raise ValueError(f"structure2 dimension is not {spatial_ndim}, got {structure2.ndim}") + structure2 = structure2.detach().to(device="cpu", dtype=torch.bool).contiguous() + + 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, From ca5566142578f31f18720234bc5c34236915eae6 Mon Sep 17 00:00:00 2001 From: Yu Han Deng Date: Mon, 13 Jul 2026 12:44:23 +0800 Subject: [PATCH 16/25] fix bugs and add validation tests --- test/test_binary.py | 152 ++++++++++++++++++++++++++------ torchmorph/csrc/torchmorph.cpp | 4 +- torchmorph/morphology/binary.py | 57 ++++++++++-- 3 files changed, 176 insertions(+), 37 deletions(-) diff --git a/test/test_binary.py b/test/test_binary.py index 201fc7f..e6102a4 100644 --- a/test/test_binary.py +++ b/test/test_binary.py @@ -1,3 +1,5 @@ +import importlib + import numpy as np import pytest import torch @@ -16,38 +18,44 @@ not torch.cuda.is_available(), reason="CUDA is required for torchmorph tests" ) -BINARY_OPERATORS = [ - pytest.param(tm.binary_erosion, scipy_binary_erosion, id="erosion"), - pytest.param(tm.binary_dilation, scipy_binary_dilation, id="dilation"), - pytest.param(tm.binary_opening, scipy_binary_opening, id="opening"), - pytest.param(tm.binary_closing, scipy_binary_closing, id="closing"), +OPERATOR_CASES = [ + ("erosion", tm.binary_erosion, scipy_binary_erosion, {"binary", "torch", "masked", "all"}), + ( + "dilation", + tm.binary_dilation, + scipy_binary_dilation, + {"binary", "torch", "masked", "all"}, + ), + ("opening", tm.binary_opening, scipy_binary_opening, {"binary", "torch", "masked", "all"}), + ("closing", tm.binary_closing, scipy_binary_closing, {"binary", "torch", "masked", "all"}), + ( + "propagation", + tm.binary_propagation, + scipy_binary_propagation, + {"propagation", "torch", "masked", "all"}, + ), + ("fill_holes", tm.binary_fill_holes, scipy_binary_fill_holes, {"fill_holes", "torch", "all"}), + ("hit_or_miss", tm.binary_hit_or_miss, scipy_binary_hit_or_miss, {"hit_or_miss", "all"}), ] -PROPAGATION_OPERATORS = [ - pytest.param(tm.binary_propagation, scipy_binary_propagation, id="propagation"), -] -FILL_HOLES_OPERATORS = [ - pytest.param(tm.binary_fill_holes, scipy_binary_fill_holes, id="fill_holes"), -] +def operator_params(group, *, include_scipy=False): + params = [] + for name, torch_op, scipy_op, groups in OPERATOR_CASES: + if group not in groups: + continue + args = (torch_op, scipy_op) if include_scipy else (torch_op,) + params.append(pytest.param(*args, id=name)) + return params -HIT_OR_MISS_OPERATORS = [ - pytest.param(tm.binary_hit_or_miss, scipy_binary_hit_or_miss, id="hit_or_miss"), -] -TORCH_OPERATORS = [ - pytest.param(tm.binary_erosion, id="erosion"), - pytest.param(tm.binary_dilation, id="dilation"), - pytest.param(tm.binary_opening, id="opening"), - pytest.param(tm.binary_closing, id="closing"), - pytest.param(tm.binary_propagation, id="propagation"), - pytest.param(tm.binary_fill_holes, id="fill_holes"), -] - -ALL_TORCH_OPERATORS = [ - *TORCH_OPERATORS, - pytest.param(tm.binary_hit_or_miss, id="hit_or_miss"), -] +BINARY_OPERATORS = operator_params("binary", include_scipy=True) +PROPAGATION_OPERATORS = operator_params("propagation", include_scipy=True) +FILL_HOLES_OPERATORS = operator_params("fill_holes", include_scipy=True) +HIT_OR_MISS_OPERATORS = operator_params("hit_or_miss", include_scipy=True) +TORCH_OPERATORS = operator_params("torch") +ALL_TORCH_OPERATORS = operator_params("all") +MASKED_TORCH_OPERATORS = operator_params("masked") CASE_2D = np.array( [[[[0, 1, 0], [1, 1, 1], [0, 1, 0]]]], @@ -359,6 +367,23 @@ def test_binary_hit_or_miss_allows_empty_miss_structure(): torch.testing.assert_close(result, expected) +@pytest.mark.parametrize( + ("torch_op", "expected_value"), + [ + pytest.param(tm.binary_erosion, True, id="erosion"), + pytest.param(tm.binary_dilation, False, id="dilation"), + ], +) +def test_binary_morphology_supports_empty_structure(torch_op, expected_value): + x = torch.as_tensor(CASE_2D, dtype=torch.float32, device="cuda") + structure = torch.zeros((3, 3), dtype=torch.bool, device="cuda") + + result = torch_op(x, structure=structure) + + expected = torch.full_like(x, expected_value, dtype=torch.bool) + torch.testing.assert_close(result, expected) + + @pytest.mark.parametrize( ("torch_op", "scipy_op"), [ @@ -421,3 +446,76 @@ def test_binary_morphology_uses_input_cuda_device(torch_op): x = torch.as_tensor(CASE_2D, dtype=torch.float32, device="cuda:1") result = torch_op(x) assert result.device == x.device + + +@pytest.mark.parametrize("torch_op", TORCH_OPERATORS) +@pytest.mark.parametrize("origin", [-2, 2]) +def test_binary_morphology_rejects_invalid_origin_value(torch_op, origin): + x = torch.as_tensor(CASE_2D, dtype=torch.float32, device="cuda") + with pytest.raises(ValueError, match="invalid origin"): + torch_op(x, origin=origin) + + +@pytest.mark.parametrize("torch_op", MASKED_TORCH_OPERATORS) +def test_binary_morphology_requires_mask_shape_to_match_input(torch_op): + x = torch.as_tensor(CASE_2D, dtype=torch.float32, device="cuda") + mask = torch.ones(x.shape[-2:], dtype=torch.bool, device="cuda") + + with pytest.raises(ValueError, match="mask shape"): + torch_op(x, mask=mask) + + +@pytest.mark.parametrize("torch_op", ALL_TORCH_OPERATORS) +def test_binary_morphology_requires_output_shape_to_match_input(torch_op): + x = torch.as_tensor(CASE_2D, dtype=torch.float32, device="cuda") + output = torch.empty((1, 1, 2, 2), dtype=torch.bool, device="cuda") + + with pytest.raises(ValueError, match="output shape"): + torch_op(x, output=output) + + +@pytest.mark.parametrize("torch_op", ALL_TORCH_OPERATORS) +def test_binary_morphology_requires_output_on_input_device(torch_op): + x = torch.as_tensor(CASE_2D, dtype=torch.float32, device="cuda") + output = torch.empty_like(x, dtype=torch.bool, device="cpu") + + with pytest.raises(ValueError, match="same device"): + torch_op(x, output=output) + + +def test_binary_morphology_normalizes_structure_before_iterations(monkeypatch): + binary_module = importlib.import_module("torchmorph.morphology.binary") + original_step = binary_module._binary_morphology_cuda_step + observed_devices = [] + + def recording_step(input, structure, border_value, origin, *, mode): + observed_devices.append(structure.device.type) + return original_step(input, structure, border_value, origin, mode=mode) + + monkeypatch.setattr(binary_module, "_binary_morphology_cuda_step", recording_step) + x = torch.as_tensor(CASE_2D, dtype=torch.float32, device="cuda") + structure = torch.ones((3, 3), dtype=torch.bool, device="cuda") + + tm.binary_dilation(x, structure=structure, iterations=2) + + assert observed_devices == ["cpu", "cpu"] + + +@pytest.mark.parametrize("spatial_ndim", [1, 8]) +def test_binary_morphology_supports_dimension_range(spatial_ndim): + spatial_shape = (2,) * spatial_ndim + np_input = np.zeros((1, 1, *spatial_shape), dtype=bool) + np_input[(0, 0, *((0,) * spatial_ndim))] = True + x = torch.as_tensor(np_input, device="cuda") + + result = tm.binary_dilation(x) + + expected = apply_scipy_to_batch(np_input, scipy_binary_dilation) + torch.testing.assert_close(result.cpu(), torch.as_tensor(expected)) + + +def test_binary_morphology_rejects_more_than_eight_spatial_dimensions(): + x = torch.zeros((1, 1, *((1,) * 9)), dtype=torch.bool, device="cuda") + + with pytest.raises(ValueError, match="1 to 8"): + tm.binary_dilation(x) diff --git a/torchmorph/csrc/torchmorph.cpp b/torchmorph/csrc/torchmorph.cpp index 1f71d8c..6321765 100644 --- a/torchmorph/csrc/torchmorph.cpp +++ b/torchmorph/csrc/torchmorph.cpp @@ -74,7 +74,7 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { py::arg("origin"), py::arg("mode"), py::arg("cval")); - + m.def("binary_erosion_cuda", &binary_erosion_cuda, "N-dimensional fused binary erosion", py::arg("input"), @@ -96,7 +96,7 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { py::arg("return_distances") = true, py::arg("return_indices") = false, py::arg("algorithm") = "exact"); - + m.def("cdt_cuda", &cdt_cuda, "Chessboard/Manhattan distance transform", py::arg("input"), diff --git a/torchmorph/morphology/binary.py b/torchmorph/morphology/binary.py index 5956537..98f6133 100644 --- a/torchmorph/morphology/binary.py +++ b/torchmorph/morphology/binary.py @@ -4,6 +4,8 @@ from .. import _C from .structure import _normalize_origin, generate_binary_structure +_BINARY_MORPH_MAX_NDIM = 8 + def _validate_binary_input(input: Tensor) -> int: if not input.is_cuda: @@ -12,10 +14,42 @@ def _validate_binary_input(input: Tensor) -> int: raise ValueError( f"Input must be (B, C, Spatial...) with at least 3 dimensions, got {input.shape}." ) + + spatial_ndim = input.ndim - 2 + if spatial_ndim > _BINARY_MORPH_MAX_NDIM: + raise ValueError( + f"Input spatial dimensions must be in range 1 to {_BINARY_MORPH_MAX_NDIM}, " + f"got {spatial_ndim}." + ) if input.numel() == 0: raise ValueError(f"Invalid input: empty tensor with shape {input.shape}.") - return input.ndim - 2 + return spatial_ndim + + +def _validate_output(input: Tensor, output: Tensor | None) -> None: + if output is None: + return + if output.shape != input.shape: + raise ValueError(f"output shape {output.shape} must match input shape {input.shape}") + if output.device != input.device: + raise ValueError( + f"output must be on the same device as input, got {output.device} and {input.device}" + ) + + +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 _validate_origin(origin: tuple[int, ...], structure: Tensor, name: str = "origin") -> None: + for origin_value, structure_size in zip(origin, structure.shape): + min_origin = -(structure_size // 2) + max_origin = (structure_size - 1) // 2 + if not min_origin <= origin_value <= max_origin: + raise ValueError(f"invalid {name}") def _binary_morphology_cuda_step( @@ -27,7 +61,6 @@ def _binary_morphology_cuda_step( mode: str, ) -> Tensor: x = (input != 0).contiguous() - structure = (structure != 0).detach().to(device="cpu", dtype=torch.bool).contiguous() kernel = _C.binary_erosion_cuda if mode == "erosion" else _C.binary_dilation_cuda return kernel(x, structure, list(origin), bool(border_value)) @@ -45,16 +78,24 @@ def _binary_morphology( ) -> Tensor: iterate_until_stable = iterations < 1 spatial_ndim = _validate_binary_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, @@ -164,6 +205,7 @@ def binary_fill_holes( ) -> Tensor: """Fill holes in binary objects for `(B, C, Spatial...)` CUDA tensors.""" _validate_binary_input(input) + _validate_output(input, output) mask = input == 0 seed = torch.zeros_like(mask, dtype=torch.bool) @@ -193,21 +235,20 @@ def binary_hit_or_miss( ) -> Tensor: """N-dimensional binary hit-or-miss transform for `(B, C, Spatial...)` CUDA tensors.""" spatial_ndim = _validate_binary_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) - elif structure1.ndim != spatial_ndim: - raise ValueError(f"structure1 dimension is not {spatial_ndim}, got {structure1.ndim}") + structure1 = _normalize_structure(structure1, spatial_ndim, "structure1") + _validate_origin(origin1, structure1, "origin1") - structure1 = structure1.detach().to(device="cpu", dtype=torch.bool).contiguous() if structure2 is None: structure2 = torch.logical_not(structure1).contiguous() else: - if structure2.ndim != spatial_ndim: - raise ValueError(f"structure2 dimension is not {spatial_ndim}, got {structure2.ndim}") - structure2 = structure2.detach().to(device="cpu", dtype=torch.bool).contiguous() + structure2 = _normalize_structure(structure2, spatial_ndim, "structure2") + _validate_origin(origin2, structure2, "origin2") input_bool = input != 0 if structure1.any().item(): From f6563ec0c98bc35828951af879bcb6412dc39875 Mon Sep 17 00:00:00 2001 From: Yu Han Deng Date: Mon, 13 Jul 2026 13:06:56 +0800 Subject: [PATCH 17/25] update .gitignore --- .gitignore | 3 + CONTEXT.md | 61 ------------------- ...alized-binary-morphology-cuda-primitive.md | 23 ------- 3 files changed, 3 insertions(+), 84 deletions(-) delete mode 100644 CONTEXT.md delete mode 100644 docs/adr/0001-generalized-binary-morphology-cuda-primitive.md diff --git a/.gitignore b/.gitignore index 4ce6066..bb5af7e 100644 --- a/.gitignore +++ b/.gitignore @@ -55,3 +55,6 @@ Thumbs.db .coverage htmlcov/ +# Local development documentation +/CONTEXT.md +/docs/ diff --git a/CONTEXT.md b/CONTEXT.md deleted file mode 100644 index 12bfdac..0000000 --- a/CONTEXT.md +++ /dev/null @@ -1,61 +0,0 @@ -# TorchMorph - -TorchMorph provides tensor-native morphological image operations and distance transforms for PyTorch workloads. Its public language follows SciPy ndimage where possible while making batch and channel dimensions explicit. - -## Language - -**Morphological Operation**: -An operation that transforms binary or grey-valued spatial samples by comparing each location with a neighborhood defined by a structuring element. -_Avoid_: filter, convolution - -**Binary Morphology**: -Morphological operations where every non-zero input value is treated as foreground and zero is treated as background. -_Avoid_: boolean convolution, mask filtering - -**Binary Hole Filling**: -A binary morphology operation that turns background regions not connected to the sample boundary into foreground. -_Avoid_: hole, fill hole - -**Binary Hit-or-Miss**: -A binary morphology operation that detects locations where foreground and background match paired structuring elements. -_Avoid_: hit miss, pattern mask - -**Grey Morphology**: -Morphological operations over numeric intensities, where the result depends on ordered values rather than only foreground/background membership. -_Avoid_: grayscale filter - -**Grey Erosion**: -A grey morphology operation that selects the minimum response from a neighborhood defined by a structuring element. -_Avoid_: minimum filter - -**Grey Dilation**: -A grey morphology operation that selects the maximum response from a neighborhood defined by a structuring element. -_Avoid_: maximum filter - -**Grey Opening**: -A grey morphology operation formed by grey erosion followed by grey dilation with the same structuring element. -_Avoid_: erosion-dilation filter - -**Grey Closing**: -A grey morphology operation formed by grey dilation followed by grey erosion with the same structuring element. -_Avoid_: dilation-erosion filter - -**Structuring Element**: -The neighborhood shape or weighted neighborhood used by a morphological operation. -_Avoid_: kernel, filter - -**Footprint**: -A binary structuring element that selects which neighboring locations participate in a grey morphology operation. -_Avoid_: mask, stencil - -**Feature Transform**: -The nearest-background coordinate result returned alongside a distance transform. -_Avoid_: index map, nearest point map - -**Spatial Dimensions**: -The trailing dimensions of an input tensor that represent the image, volume, or higher-dimensional sample being transformed. -_Avoid_: data dimensions - -**Batch-Channel Tensor**: -A tensor shaped as `(B, C, Spatial...)`, where `B` and `C` are independent leading dimensions and every `(B, C)` slice is transformed independently. -_Avoid_: image tensor, BCHW-only tensor diff --git a/docs/adr/0001-generalized-binary-morphology-cuda-primitive.md b/docs/adr/0001-generalized-binary-morphology-cuda-primitive.md deleted file mode 100644 index 776dd34..0000000 --- a/docs/adr/0001-generalized-binary-morphology-cuda-primitive.md +++ /dev/null @@ -1,23 +0,0 @@ -# 0001 Generalized Binary Morphology CUDA Primitive - -## Status - -Accepted - -## Context - -TorchMorph already uses a generalized CUDA primitive for grey erosion and grey dilation. Binary erosion and binary dilation were implemented in Python with PyTorch padding and convolution helpers. That implementation kept SciPy-aligned behavior, but it made binary morphology less consistent with the project goal of CUDA-accelerated batch-parallel morphology. - -Binary erosion and binary dilation are also the base operations for binary opening, binary closing, binary propagation, and binary hole filling. Optimizing only the public erosion and dilation wrappers would leave later binary operators without a clear reusable CUDA foundation. - -## Decision - -Implement a generalized binary morphology CUDA primitive, modeled after the grey morphology CUDA primitive, and dispatch CUDA binary erosion and binary dilation through it. The primitive supports batch-channel tensors with arbitrary spatial rank up to the same practical maximum used by the grey implementation. Python remains responsible for iteration loops, mask application, output copying, and composite operators. Binary morphology follows the CUDA-only public contract used by the other accelerated morphology and distance-transform operators. - -Do not start with specialized 2D or 3D kernels. Keep those as later optimizations once the generalized primitive is correct and benchmarked. - -## Consequences - -The first CUDA implementation favors broad SciPy-aligned behavior and shared maintenance shape over peak 2D/3D performance. This gives binary propagation and binary hole filling a reusable base and keeps the public Python API stable. - -Specialized 2D and 3D kernels may still be added later as fast paths when benchmark data shows that fixed-rank kernels would materially improve common workloads. From 556d7a1df70e76eb89d4eba3ce470044547636f6 Mon Sep 17 00:00:00 2001 From: Yu Han Deng Date: Mon, 13 Jul 2026 13:51:20 +0800 Subject: [PATCH 18/25] back validation and refactor distance transform --- benchmark/distance_transform.py | 535 +++++---------- test/test_distance_transform.py | 1069 +++++++----------------------- test/test_grey.py | 17 + torchmorph/distance_transform.py | 492 +++++--------- torchmorph/morphology/binary.py | 19 +- torchmorph/morphology/grey.py | 51 +- 6 files changed, 652 insertions(+), 1531 deletions(-) diff --git a/benchmark/distance_transform.py b/benchmark/distance_transform.py index 465f011..a5927ed 100644 --- a/benchmark/distance_transform.py +++ b/benchmark/distance_transform.py @@ -1,398 +1,219 @@ import argparse +from functools import partial -import scipy.ndimage as ndi # noqa: F401 +import scipy.ndimage as ndi import torch import torch.utils.benchmark as benchmark from prettytable import PrettyTable -import torchmorph as tm # noqa: F401 +import torchmorph as tm -# Config -sizes_2d = [64, 128, 256, 512, 1024] -sizes_3d = [32, 64, 128, 256] -batches_2d = [1, 4, 8, 16] -batches_3d = [1, 2, 4, 8] -dtype = torch.float32 -device = "cuda" -MIN_RUN = 1.0 # seconds per measurement +SIZES_2D = [64, 128, 256, 512, 1024] +SIZES_3D = [32, 64, 128, 256] +SIZES_BFDT = [32, 64, 128, 256] +BATCHES_2D = [1, 4, 8, 16] +BATCHES_3D = [1, 2, 4, 8] +MIN_RUN_TIME = 1.0 -torch.set_num_threads(torch.get_num_threads()) +def run_scipy(operation, inputs): + return [operation(input) for input in inputs] -# ====================================================================== -# Section 1: Euclidean Distance Transform (EDT) — 2D -# ====================================================================== +def run_cuda(operation, input): + result = operation(input) + torch.cuda.synchronize() + return result -def bench_edt_2d(): - for B in batches_2d: - table = PrettyTable() - table.field_names = [ - "Size", - "SciPy (ms/img)", - "Exact 1× (ms/img)", - "Exact batch (ms/img)", - "JFA 1× (ms/img)", - "JFA batch (ms/img)", - "Speedup Exact", - "Speedup JFA", - ] - for c in table.field_names: - table.align[c] = "r" - - for s in sizes_2d: - # Inputs: (B, C, H, W) format - C=1 for single channel - x = (torch.randn(B, 1, s, s, device=device) > 0).to(dtype) - # For scipy, we need (H, W) arrays - x_np_list = [x[i, 0].detach().cpu().numpy() for i in range(B)] - # For torch single image processing: each is (1, 1, H, W) - x_imgs = [x[i : i + 1] for i in range(B)] - - # Expose locals to __main__ for benchmark.Timer - globals().update(x=x, x_np_list=x_np_list, x_imgs=x_imgs) - - # SciPy (CPU, one-by-one) - stmt_scipy = "out = [ndi.distance_transform_edt(arr) for arr in x_np_list]" - t_scipy = benchmark.Timer( - stmt=stmt_scipy, - setup="from __main__ import x_np_list, ndi", - num_threads=torch.get_num_threads(), - ).blocked_autorange(min_run_time=MIN_RUN) - scipy_per_img_ms = (t_scipy.median * 1e3) / B - - # Torch Exact (CUDA, one-by-one) - stmt_exact1 = """ -for xi in x_imgs: - tm.euclidean_distance_transform(xi, algorithm="exact") -""" - t_exact1 = benchmark.Timer( - stmt=stmt_exact1, - setup="from __main__ import x_imgs, tm", - num_threads=torch.get_num_threads(), - ).blocked_autorange(min_run_time=MIN_RUN) - exact1_per_img_ms = (t_exact1.median * 1e3) / B - - # Torch Exact (CUDA, batched) - t_exact_batch = benchmark.Timer( - stmt='tm.euclidean_distance_transform(x, algorithm="exact")', - setup="from __main__ import x, tm", - num_threads=torch.get_num_threads(), - ).blocked_autorange(min_run_time=MIN_RUN) - exactB_per_img_ms = (t_exact_batch.median * 1e3) / B - - # Torch JFA (CUDA, one-by-one) - stmt_jfa1 = """ -for xi in x_imgs: - tm.euclidean_distance_transform(xi, algorithm="jfa") -""" - t_jfa1 = benchmark.Timer( - stmt=stmt_jfa1, - setup="from __main__ import x_imgs, tm", - num_threads=torch.get_num_threads(), - ).blocked_autorange(min_run_time=MIN_RUN) - jfa1_per_img_ms = (t_jfa1.median * 1e3) / B - - # Torch JFA (CUDA, batched) - t_jfa_batch = benchmark.Timer( - stmt='tm.euclidean_distance_transform(x, algorithm="jfa")', - setup="from __main__ import x, tm", - num_threads=torch.get_num_threads(), - ).blocked_autorange(min_run_time=MIN_RUN) - jfaB_per_img_ms = (t_jfa_batch.median * 1e3) / B - - # Speedups (batch mode vs scipy) - speed_exact = scipy_per_img_ms / exactB_per_img_ms - speed_jfa = scipy_per_img_ms / jfaB_per_img_ms - - table.add_row( - [ - s, - f"{scipy_per_img_ms:.3f}", - f"{exact1_per_img_ms:.3f}", - f"{exactB_per_img_ms:.3f}", - f"{jfa1_per_img_ms:.3f}", - f"{jfaB_per_img_ms:.3f}", - f"{speed_exact:.1f}×", - f"{speed_jfa:.1f}×", - ] - ) - print(f"\n=== EDT 2D | Batch Size: {B} ===") - print(table) +def run_cuda_singles(operation, inputs): + result = [operation(input) for input in inputs] + torch.cuda.synchronize() + return result -# ====================================================================== -# Section 2: Euclidean Distance Transform (EDT) — 3D -# ====================================================================== +def measure_ms(statement, globals, batch_size, min_run_time): + measurement = benchmark.Timer( + stmt=statement, + globals=globals, + num_threads=torch.get_num_threads(), + ).blocked_autorange(min_run_time=min_run_time) + return measurement.median * 1e3 / batch_size -def bench_edt_3d(): - for B in batches_3d: - table = PrettyTable() - table.field_names = [ - "Size (D×H×W)", - "SciPy (ms/vol)", - "Exact 1× (ms/vol)", - "Exact batch (ms/vol)", - "JFA 1× (ms/vol)", - "JFA batch (ms/vol)", - "Speedup Exact", - "Speedup JFA", - ] - for c in table.field_names: - table.align[c] = "r" - - for s in sizes_3d: - # Skip large sizes with large batches to avoid OOM - if s >= 256 and B >= 4: - table.add_row([f"{s}³", "OOM", "OOM", "OOM", "OOM", "OOM", "-", "-"]) - continue +def make_inputs(batch_size, image_size, spatial_ndim): + shape = (batch_size, 1, *([image_size] * spatial_ndim)) + input = (torch.randn(shape, device="cuda") > 0).float() + scipy_inputs = [input[index, 0].cpu().numpy() for index in range(batch_size)] + torch_inputs = [input[index : index + 1] for index in range(batch_size)] + return input, scipy_inputs, torch_inputs - # Inputs: (B, D, H, W) format for 3D - no channel dimension for JFA 3D - x = (torch.randn(B, s, s, s, device=device) > 0).to(dtype) - # For scipy, we need (D, H, W) arrays - x_np_list = [x[i].detach().cpu().numpy() for i in range(B)] - # For torch single volume processing: each is (1, D, H, W) - x_vols = [x[i : i + 1] for i in range(B)] - - # Expose locals to __main__ for benchmark.Timer - globals().update(x=x, x_np_list=x_np_list, x_vols=x_vols) - - # SciPy (CPU, one-by-one) - stmt_scipy = "out = [ndi.distance_transform_edt(arr) for arr in x_np_list]" - t_scipy = benchmark.Timer( - stmt=stmt_scipy, - setup="from __main__ import x_np_list, ndi", - num_threads=torch.get_num_threads(), - ).blocked_autorange(min_run_time=MIN_RUN) - scipy_per_vol_ms = (t_scipy.median * 1e3) / B - - # Torch Exact (CUDA, one-by-one) - stmt_exact1 = """ -for xi in x_vols: - tm.euclidean_distance_transform(xi, algorithm="exact") -""" - t_exact1 = benchmark.Timer( - stmt=stmt_exact1, - setup="from __main__ import x_vols, tm", - num_threads=torch.get_num_threads(), - ).blocked_autorange(min_run_time=MIN_RUN) - exact1_per_vol_ms = (t_exact1.median * 1e3) / B - - # Torch Exact (CUDA, batched) - t_exact_batch = benchmark.Timer( - stmt='tm.euclidean_distance_transform(x, algorithm="exact")', - setup="from __main__ import x, tm", - num_threads=torch.get_num_threads(), - ).blocked_autorange(min_run_time=MIN_RUN) - exactB_per_vol_ms = (t_exact_batch.median * 1e3) / B - - # Torch JFA (CUDA, one-by-one) - stmt_jfa1 = """ -for xi in x_vols: - tm.euclidean_distance_transform(xi, algorithm="jfa") -""" - t_jfa1 = benchmark.Timer( - stmt=stmt_jfa1, - setup="from __main__ import x_vols, tm", - num_threads=torch.get_num_threads(), - ).blocked_autorange(min_run_time=MIN_RUN) - jfa1_per_vol_ms = (t_jfa1.median * 1e3) / B - - # Torch JFA (CUDA, batched) - t_jfa_batch = benchmark.Timer( - stmt='tm.euclidean_distance_transform(x, algorithm="jfa")', - setup="from __main__ import x, tm", - num_threads=torch.get_num_threads(), - ).blocked_autorange(min_run_time=MIN_RUN) - jfaB_per_vol_ms = (t_jfa_batch.median * 1e3) / B - - # Speedups (batch mode vs scipy) - speed_exact = scipy_per_vol_ms / exactB_per_vol_ms - speed_jfa = scipy_per_vol_ms / jfaB_per_vol_ms - - table.add_row( + +def benchmark_transform( + title, + scipy_operation, + torch_operations, + spatial_ndim, + sizes, + batch_sizes, + min_run_time, +): + print(f"\n=== {title} ===") + for batch_size in batch_sizes: + table = PrettyTable() + fields = ["Size", "SciPy (ms/item)"] + for name in torch_operations: + fields.extend( [ - f"{s}³", - f"{scipy_per_vol_ms:.3f}", - f"{exact1_per_vol_ms:.3f}", - f"{exactB_per_vol_ms:.3f}", - f"{jfa1_per_vol_ms:.3f}", - f"{jfaB_per_vol_ms:.3f}", - f"{speed_exact:.1f}×", - f"{speed_jfa:.1f}×", + f"{name} 1x (ms/item)", + f"{name} batch (ms/item)", + f"{name} speedup", ] ) + table.field_names = fields + for field in fields: + table.align[field] = "r" - print(f"\n=== EDT 3D | Batch Size: {B} ===") - print(table) + for image_size in sizes: + if spatial_ndim == 3 and image_size >= 256 and batch_size >= 4: + table.add_row([f"{image_size}^3", *(["OOM"] * (len(fields) - 1))]) + continue + input, scipy_inputs, torch_inputs = make_inputs(batch_size, image_size, spatial_ndim) + scipy_ms = measure_ms( + "run_scipy(operation, inputs)", + { + "run_scipy": run_scipy, + "operation": scipy_operation, + "inputs": scipy_inputs, + }, + batch_size, + min_run_time, + ) -# ====================================================================== -# Section 3: Chamfer Distance Transform (CDT) — 2D -# ====================================================================== - - -def bench_cdt_2d(): - for metric in ["chessboard", "taxicab"]: - print(f"\n{'=' * 60}") - print(f" CDT Benchmark - Metric: {metric}") - print(f"{'=' * 60}") - - for B in batches_2d: - table = PrettyTable() - table.field_names = [ - "Size", - "SciPy (ms/img)", - "Torch 1× (ms/img)", - "Torch batch (ms/img)", - "Speedup 1×", - "Speedup batch", - ] - for c in table.field_names: - table.align[c] = "r" - - for s in sizes_2d: - # Inputs: (B, C, H, W) format - C=1 for single channel - x = (torch.randn(B, 1, s, s, device=device) > 0).to(dtype) - # For scipy, we need (H, W) arrays - x_np_list = [x[i, 0].detach().cpu().numpy() for i in range(B)] - # For torch single image processing: each is (1, 1, H, W) - x_imgs = [x[i : i + 1] for i in range(B)] - - # Expose locals to __main__ for benchmark.Timer - globals().update(x=x, x_np_list=x_np_list, x_imgs=x_imgs) - - # SciPy (CPU, one-by-one) - stmt_scipy = ( - f"out = [ndi.distance_transform_cdt(arr,metric='{metric}')for arr in x_np_list]" + row = [f"{image_size}^{spatial_ndim}", f"{scipy_ms:.3f}"] + for operation in torch_operations.values(): + single_ms = measure_ms( + "run_cuda_singles(operation, inputs)", + { + "run_cuda_singles": run_cuda_singles, + "operation": operation, + "inputs": torch_inputs, + }, + batch_size, + min_run_time, + ) + batch_ms = measure_ms( + "run_cuda(operation, input)", + { + "run_cuda": run_cuda, + "operation": operation, + "input": input, + }, + batch_size, + min_run_time, ) - t_scipy = benchmark.Timer( - stmt=stmt_scipy, - setup="from __main__ import x_np_list, ndi", - num_threads=torch.get_num_threads(), - ).blocked_autorange(min_run_time=MIN_RUN) - scipy_per_img_ms = (t_scipy.median * 1e3) / B - - # Torch (CUDA, one-by-one) - stmt_torch1 = f""" -for xi in x_imgs: - tm.chamfer_distance_transform(xi, metric='{metric}') -""" - t_torch1 = benchmark.Timer( - stmt=stmt_torch1, - setup="from __main__ import x_imgs, tm", - num_threads=torch.get_num_threads(), - ).blocked_autorange(min_run_time=MIN_RUN) - torch1_per_img_ms = (t_torch1.median * 1e3) / B - - # Torch (CUDA, batched) - t_batch = benchmark.Timer( - stmt=f"tm.chamfer_distance_transform(x, metric='{metric}')", - setup="from __main__ import x, tm", - num_threads=torch.get_num_threads(), - ).blocked_autorange(min_run_time=MIN_RUN) - torchB_per_img_ms = (t_batch.median * 1e3) / B - - # Speedups - speed1 = scipy_per_img_ms / torch1_per_img_ms - speedB = scipy_per_img_ms / torchB_per_img_ms - - table.add_row( + row.extend( [ - s, - f"{scipy_per_img_ms:.3f}", - f"{torch1_per_img_ms:.3f}", - f"{torchB_per_img_ms:.3f}", - f"{speed1:.1f}×", - f"{speedB:.1f}×", + f"{single_ms:.3f}", + f"{batch_ms:.3f}", + f"{scipy_ms / batch_ms:.1f}x", ] ) + table.add_row(row) - print(f"\n=== CDT 2D | Metric: {metric}, Batch Size: {B} ===") - print(table) + print(f"\nBatch size: {batch_size}") + print(table) -# ====================================================================== -# Section 4: Brute-Force Distance Transform (BFDT) -# ====================================================================== +def bench_edt(args): + operations = { + "exact": partial(tm.euclidean_distance_transform, algorithm="exact"), + "jfa": partial(tm.euclidean_distance_transform, algorithm="jfa"), + } + benchmark_transform( + "EDT 2D", + ndi.distance_transform_edt, + operations, + 2, + args.sizes_2d, + args.batches_2d, + args.min_run_time, + ) + benchmark_transform( + "EDT 3D", + ndi.distance_transform_edt, + operations, + 3, + args.sizes_3d, + args.batches_3d, + args.min_run_time, + ) -def bench_bfdt(): - # BFDT is slow, so we use smaller sizes - bf_sizes = [32, 64, 128, 256] - for metric in ["euclidean", "taxicab", "chessboard"]: - print(f"\n{'=' * 60}") - print(f" BFDT Benchmark - Metric: {metric}") - print(f"{'=' * 60}") +def bench_cdt(args): + for metric in ("chessboard", "taxicab"): + benchmark_transform( + f"CDT 2D ({metric})", + partial(ndi.distance_transform_cdt, metric=metric), + { + metric: partial( + tm.chamfer_distance_transform, + metric=metric, + ) + }, + 2, + args.sizes_2d, + args.batches_2d, + args.min_run_time, + ) + + +def bench_bfdt(args): + for metric in ("euclidean", "taxicab", "chessboard"): + benchmark_transform( + f"BFDT 2D ({metric})", + partial(ndi.distance_transform_bf, metric=metric), + { + metric: partial( + tm.brute_force_distance_transform, + metric=metric, + ) + }, + 2, + args.sizes_bfdt, + [1], + min(args.min_run_time, 0.5), + ) - table = PrettyTable() - table.field_names = [ - "Size", - "SciPy (ms/img)", - "Torch (ms/img)", - "Speedup", - ] - for c in table.field_names: - table.align[c] = "r" - - for s in bf_sizes: - # Single batch for BFDT to avoid excessive wait - B = 1 - x = (torch.randn(B, 1, s, s, device=device) > 0).to(dtype) - x_np_list = [x[i, 0].detach().cpu().numpy() for i in range(B)] - - globals().update(x=x, x_np_list=x_np_list) - - # SciPy - stmt_scipy = ( - f"out = [ndi.distance_transform_bf(arr, metric='{metric}') for arr in x_np_list]" - ) - t_scipy = benchmark.Timer( - stmt=stmt_scipy, - setup="from __main__ import x_np_list, ndi", - num_threads=torch.get_num_threads(), - ).blocked_autorange(min_run_time=0.5) - scipy_ms = (t_scipy.median * 1e3) / B - - # Torch - stmt_torch = f"tm.brute_force_distance_transform(x, metric='{metric}')" - t_torch = benchmark.Timer( - stmt=stmt_torch, - setup="from __main__ import x, tm", - num_threads=torch.get_num_threads(), - ).blocked_autorange(min_run_time=0.5) - torch_ms = (t_torch.median * 1e3) / B - - table.add_row( - [ - f"{s}x{s}", - f"{scipy_ms:.3f}", - f"{torch_ms:.3f}", - f"{scipy_ms / torch_ms:.1f}×", - ] - ) - print(table) +BENCHMARKS = { + "edt": bench_edt, + "cdt": bench_cdt, + "bfdt": bench_bfdt, +} -if __name__ == "__main__": - parser = argparse.ArgumentParser(description="Distance Transform Benchmarks") +def main(): + parser = argparse.ArgumentParser(description="Distance transform benchmarks.") parser.add_argument( - "--section", - nargs="*", - choices=["edt-2d", "edt-3d", "cdt", "bfdt"], - default=None, - help="Sections to run (default: all)", + "operation", + nargs="?", + choices=(*BENCHMARKS, "all"), + default="all", ) + parser.add_argument("--sizes-2d", type=int, nargs="+", default=SIZES_2D) + parser.add_argument("--sizes-3d", type=int, nargs="+", default=SIZES_3D) + parser.add_argument("--sizes-bfdt", type=int, nargs="+", default=SIZES_BFDT) + parser.add_argument("--batches-2d", type=int, nargs="+", default=BATCHES_2D) + parser.add_argument("--batches-3d", type=int, nargs="+", default=BATCHES_3D) + parser.add_argument("--min-run-time", type=float, default=MIN_RUN_TIME) args = parser.parse_args() - sections = args.section if args.section else ["edt-2d", "edt-3d", "cdt", "bfdt"] - if "edt-2d" in sections: - bench_edt_2d() - if "edt-3d" in sections: - bench_edt_3d() - if "cdt" in sections: - bench_cdt_2d() - if "bfdt" in sections: - bench_bfdt() + operations = BENCHMARKS if args.operation == "all" else (args.operation,) + for operation in operations: + BENCHMARKS[operation](args) + + +if __name__ == "__main__": + main() diff --git a/test/test_distance_transform.py b/test/test_distance_transform.py index bf232b3..75b461a 100644 --- a/test/test_distance_transform.py +++ b/test/test_distance_transform.py @@ -1,865 +1,306 @@ -import numpy as np # noqa: F401 +import numpy as np import pytest import torch -from scipy.ndimage import distance_transform_bf as scipy_bfdt -from scipy.ndimage import distance_transform_cdt as scipy_cdt -from scipy.ndimage import distance_transform_edt as scipy_edt # noqa: F401 +from scipy import ndimage as ndi -import torchmorph as tm # noqa: F401 +import torchmorph as tm -# Global pytestmark to skip tests if CUDA is not available pytestmark = pytest.mark.skipif( not torch.cuda.is_available(), reason="CUDA is required for torchmorph tests" ) -def apply_scipy_to_batch(batch_numpy, spatial_ndim, scipy_func, **kwargs): - """Unified SciPy batch wrapper to follow DRY principle.""" - original_shape = batch_numpy.shape - spatial_shape = original_shape[-spatial_ndim:] if spatial_ndim > 0 else () - batch_shape = original_shape[:-spatial_ndim] if spatial_ndim > 0 else () - - batch_size = int(np.prod(batch_shape)) if batch_shape else 1 - flat_input = batch_numpy.reshape(batch_size, *spatial_shape) - - dist_results, indices_results = [], [] - - for sample in flat_input: - result = scipy_func(sample, **kwargs) - if isinstance(result, tuple): - dist_results.append(result[0]) - indices_results.append(result[1]) - else: - dist_results.append(result) - - output_dist = ( - np.stack(dist_results, axis=0).reshape(*batch_shape, *spatial_shape) - if batch_shape - else dist_results[0] - ) - - if indices_results: - # SciPy EDT/BFDT returns indices with shape (ndim, *spatial_shape) - # after stacking we get (batch_size, ndim, *spatial_shape) - # reshape to (*batch_shape, spatial_ndim, *spatial_shape) then move ndim axis to front - # to match CUDA convention: (spatial_ndim, *batch_shape, *spatial_shape) - output_indices = ( - np.moveaxis( - np.stack(indices_results, axis=0).reshape( - *batch_shape, spatial_ndim, *spatial_shape - ), - len(batch_shape), - 0, +def make_case(shape): + values = np.ones(shape, dtype=np.float32) + slices = [slice(None, None, 2) for _ in shape[2:]] + values[(slice(None), slice(None), *slices)] = 0 + return values + + +CASES = { + 1: make_case((2, 1, 7)), + 2: make_case((2, 1, 5, 6)), + 3: make_case((1, 2, 4, 5, 6)), +} + +DISTANCE_OPERATORS = [ + pytest.param(tm.euclidean_distance_transform, id="edt"), + pytest.param(tm.chamfer_distance_transform, id="cdt"), + pytest.param(tm.brute_force_distance_transform, id="bfdt"), +] + +SCIPY_CASES = [ + pytest.param(tm.euclidean_distance_transform, ndi.distance_transform_edt, {}, {}, id="edt"), + pytest.param( + tm.chamfer_distance_transform, + ndi.distance_transform_cdt, + {"metric": "chessboard"}, + {"metric": "chessboard"}, + id="cdt-chessboard", + ), + pytest.param( + tm.chamfer_distance_transform, + ndi.distance_transform_cdt, + {"metric": "taxicab"}, + {"metric": "taxicab"}, + id="cdt-taxicab", + ), + pytest.param( + tm.brute_force_distance_transform, + ndi.distance_transform_bf, + {"metric": "euclidean"}, + {"metric": "euclidean"}, + id="bfdt-euclidean", + ), + pytest.param( + tm.brute_force_distance_transform, + ndi.distance_transform_bf, + {"metric": "taxicab"}, + {"metric": "taxicab"}, + id="bfdt-taxicab", + ), + pytest.param( + tm.brute_force_distance_transform, + ndi.distance_transform_bf, + {"metric": "chessboard"}, + {"metric": "chessboard"}, + id="bfdt-chessboard", + ), +] + + +def scipy_batch(input_array, scipy_op, **kwargs): + spatial_shape = input_array.shape[2:] + samples = input_array.reshape(-1, *spatial_shape) + results = [scipy_op(sample, **kwargs) for sample in samples] + return np.stack(results).reshape(input_array.shape) + + +def cuda(array): + return torch.as_tensor(np.ascontiguousarray(array), device="cuda") + + +def assert_indices_describe_distances(input, distances, indices, metric, sampling=None): + spatial_ndim = input.ndim - 2 + assert indices.shape == (spatial_ndim, *input.shape) + sampling = sampling or [1.0] * spatial_ndim + + for batch in range(input.shape[0]): + for channel in range(input.shape[1]): + sample = input[batch, channel] + sample_indices = indices[:, batch, channel].long() + assert torch.all(sample[tuple(sample_indices)] == 0) + axes = torch.meshgrid( + *[torch.arange(size, device=input.device) for size in input.shape[2:]], + indexing="ij", ) - if batch_shape - else indices_results[0] - ) - return output_dist, output_indices - - return output_dist, None - - -# ====================================================================== -# EDT Helper functions -# ====================================================================== -def batch_scipy_edt_with_indices( - batch_numpy: np.ndarray, - spatial_ndim: int, -) -> tuple[np.ndarray, np.ndarray]: - """Compute SciPy EDT and indices for a batch of arrays.""" - return apply_scipy_to_batch( - batch_numpy, spatial_ndim, scipy_edt, return_indices=True, return_distances=True - ) - - -def batch_scipy_edt_with_sampling( - batch_numpy: np.ndarray, - spatial_ndim: int, - sampling: list[float], -) -> tuple[np.ndarray, np.ndarray]: - """Compute SciPy EDT with sampling for a batch of arrays.""" - return apply_scipy_to_batch( - batch_numpy, - spatial_ndim, - scipy_edt, - sampling=sampling, - return_indices=True, - return_distances=True, - ) - - -# ====================================================================== -# CDT Helper functions -# ====================================================================== -def batch_scipy_cdt( - batch_numpy: np.ndarray, - metric: str = "chessboard", - return_indices: bool = False, - spatial_ndim: int = 2, -) -> tuple[np.ndarray, np.ndarray | None]: - """Compute SciPy CDT for a batch of arrays.""" - return apply_scipy_to_batch( - batch_numpy, - spatial_ndim, - scipy_cdt, - metric=metric, - return_distances=True, - return_indices=return_indices, - ) - - -# ====================================================================== -# BFDT Helper functions -# ====================================================================== -def batch_scipy_bfdt( - batch_numpy: np.ndarray, - metric: str = "euclidean", - sampling: list[float] | None = None, - return_indices: bool = False, - spatial_ndim: int = 2, -) -> tuple[np.ndarray, np.ndarray | None]: - """Compute SciPy BFDT for a batch of arrays.""" - return apply_scipy_to_batch( - batch_numpy, - spatial_ndim, - scipy_bfdt, - metric=metric, - sampling=sampling, - return_distances=True, - return_indices=return_indices, - ) - - -# ====================================================================== -# EDT Test data: (B, C, Spatial...) format -# ====================================================================== -# 1D spatial: (B=2, C=1, W=6) -edt_case_1d = np.array( - [[[1, 1, 0, 1, 0, 1]], [[0, 1, 1, 1, 1, 0]]], - dtype=np.float32, -) - -# 2D spatial: (B=2, C=1, H=3, W=4) -edt_case_2d = np.array( - [ - [[[0.0, 1, 1, 1], [0, 0, 1, 1], [0, 1, 1, 0]]], - [[[0, 1, 0, 1], [1, 0, 1, 0], [0, 1, 0, 1]]], - ], - dtype=np.float32, -) - -# 2D spatial single batch: (B=1, C=1, H=4, W=4) -edt_case_2d_single = np.array( - [ - [ - [ - [0, 1, 0, 1], - [1, 0, 1, 0], - [0, 1, 0, 1], - [1, 0, 1, 0], + deltas = [ + (axis - sample_indices[dim]).abs() * sampling[dim] for dim, axis in enumerate(axes) ] - ] - ], - dtype=np.float32, -) - -# 3D spatial: (B=2, C=1, D=4, H=5, W=6) -_edt_case_3d_s1 = np.ones((1, 4, 5, 6), dtype=np.float32) -_edt_case_3d_s1[0, 1, 1, 1] = 0.0 -_edt_case_3d_s1[0, 2, 3, 4] = 0.0 - -_edt_case_3d_s2 = np.ones((1, 4, 5, 6), dtype=np.float32) -_edt_case_3d_s2[0, 0, 0, 0] = 0.0 + if metric == "euclidean": + expected = torch.sqrt(sum(delta.square() for delta in deltas)) + elif metric == "taxicab": + expected = sum(deltas) + else: + expected = torch.stack(deltas).amax(dim=0) + torch.testing.assert_close(distances[batch, channel], expected.float()) + + +@pytest.mark.parametrize(("torch_op", "scipy_op", "torch_kwargs", "scipy_kwargs"), SCIPY_CASES) +@pytest.mark.parametrize("spatial_ndim", [1, 2, 3]) +def test_distance_transform_matches_scipy( + torch_op, scipy_op, torch_kwargs, scipy_kwargs, spatial_ndim +): + input_array = CASES[spatial_ndim] + result = torch_op(cuda(input_array), **torch_kwargs) + expected = scipy_batch(input_array, scipy_op, **scipy_kwargs) + torch.testing.assert_close(result.cpu(), torch.as_tensor(expected).float()) -edt_case_3d = np.stack([_edt_case_3d_s1, _edt_case_3d_s2], axis=0) # (B=2, C=1, D=4, H=5, W=6) -# 2D with unit dimension: (B=2, C=1, H=5, W=1) -edt_case_2d_unit = np.ones((2, 1, 5, 1), dtype=np.float32) -edt_case_2d_unit[0, 0, 2, 0] = 0.0 -edt_case_2d_unit[1, 0, 4, 0] = 0.0 - - -# ====================================================================== -# CDT Test data: (B, C, Spatial...) format -# ====================================================================== -# 1D spatial: (B=2, C=1, W=9) -cdt_case_1d = np.array( - [[[0, 1, 1, 1, 1, 0, 1, 1, 0]], [[1, 1, 0, 1, 1, 1, 1, 0, 1]]], - dtype=np.float32, -) - -# 2D spatial: (B=1, C=1, H=5, W=6) -cdt_case_2d_simple = np.array( - [ - [ - [ - [0, 1, 1, 1, 1, 0], - [1, 1, 1, 1, 1, 1], - [1, 1, 1, 1, 1, 1], - [1, 1, 1, 1, 1, 1], - [0, 1, 1, 1, 1, 0], - ] - ] - ], - dtype=np.float32, -) - -# 2D spatial batch: (B=2, C=1, H=4, W=5) -cdt_case_2d_batch = np.array( - [ - [[[0, 1, 1, 1, 0], [1, 1, 1, 1, 1], [1, 1, 1, 1, 1], [0, 1, 1, 1, 0]]], - [[[1, 1, 0, 1, 1], [1, 1, 1, 1, 1], [1, 1, 1, 1, 1], [1, 1, 0, 1, 1]]], - ], - dtype=np.float32, -) - -# 2D checkerboard: (B=1, C=1, H=4, W=4) -cdt_case_checkerboard = np.array( - [ - [ - [ - [0, 1, 0, 1], - [1, 0, 1, 0], - [0, 1, 0, 1], - [1, 0, 1, 0], - ] - ] - ], - dtype=np.float32, -) - -# 3D spatial: (B=1, C=1, D=5, H=5, W=5) -_cdt_case_3d_simple = np.zeros((1, 1, 5, 5, 5), dtype=np.float32) -_cdt_case_3d_simple[0, 0, 1:4, 1:4, 1:4] = 1 # 3x3x3 cube of foreground -cdt_case_3d_simple = _cdt_case_3d_simple - -# 3D sphere: (B=1, C=1, D=7, H=7, W=7) -_cdt_case_3d_sphere = np.zeros((1, 1, 7, 7, 7), dtype=np.float32) -for z in range(7): - for y in range(7): - for x in range(7): - if (z - 3) ** 2 + (y - 3) ** 2 + (x - 3) ** 2 <= 4: - _cdt_case_3d_sphere[0, 0, z, y, x] = 1 -cdt_case_3d_sphere = _cdt_case_3d_sphere - -# 3D batch: (B=2, C=1, D=4, H=5, W=6) -_cdt_case_3d_batch_s1 = np.ones((1, 4, 5, 6), dtype=np.float32) -_cdt_case_3d_batch_s1[0, 1, 1, 1] = 0.0 -_cdt_case_3d_batch_s1[0, 2, 3, 4] = 0.0 - -_cdt_case_3d_batch_s2 = np.ones((1, 4, 5, 6), dtype=np.float32) -_cdt_case_3d_batch_s2[0, 0, 0, 0] = 0.0 - -cdt_case_3d_batch = np.stack( - [_cdt_case_3d_batch_s1, _cdt_case_3d_batch_s2], axis=0 -) # (B=2, C=1, D=4, H=5, W=6) - - -# ====================================================================== -# EDT Tests -# ====================================================================== @pytest.mark.parametrize( - "input_numpy, spatial_ndim", + ("torch_op", "kwargs", "metric"), [ - pytest.param(edt_case_1d, 1, id="1D_B2C1"), - pytest.param(edt_case_2d, 2, id="2D_B2C1"), - pytest.param(edt_case_2d_single, 2, id="2D_B1C1"), - pytest.param(edt_case_3d, 3, id="3D_B2C1"), - pytest.param(edt_case_2d_unit, 2, id="2D_UnitDim_B2C1"), + pytest.param(tm.euclidean_distance_transform, {}, "euclidean", id="edt"), + pytest.param( + tm.chamfer_distance_transform, + {"metric": "chessboard"}, + "chessboard", + id="cdt-chessboard", + ), + pytest.param( + tm.chamfer_distance_transform, + {"metric": "taxicab"}, + "taxicab", + id="cdt-taxicab", + ), + pytest.param( + tm.brute_force_distance_transform, + {"metric": "euclidean"}, + "euclidean", + id="bfdt-euclidean", + ), ], ) -def test_edt_distance_and_indices( - input_numpy: np.ndarray, - spatial_ndim: int, - request: pytest.FixtureRequest, -) -> None: - # 1. Prepare data - x_numpy_contiguous = np.ascontiguousarray(input_numpy) - x_cuda = torch.as_tensor(x_numpy_contiguous, device="cuda") - - # 2. Create sampling list to specify spatial dimensions - sampling = [1.0] * spatial_ndim - - # 3. Run CUDA EDT - dist_cuda, idx_cuda = tm.euclidean_distance_transform( - x_cuda.clone(), sampling=sampling, return_indices=True +@pytest.mark.parametrize("spatial_ndim", [1, 2, 3]) +def test_indices_point_to_nearest_background(torch_op, kwargs, metric, spatial_ndim): + input = cuda(CASES[spatial_ndim]) + distances, indices = torch_op(input, return_indices=True, **kwargs) + assert_indices_describe_distances(input, distances, indices, metric) + + +@pytest.mark.parametrize("distance_op", DISTANCE_OPERATORS) +def test_return_flags(distance_op): + input = cuda(CASES[2]) + distances = distance_op(input) + indices = distance_op(input, return_distances=False, return_indices=True) + both = distance_op(input, return_indices=True) + + assert distances.shape == input.shape + assert indices.shape == (2, *input.shape) + assert both[0].shape == input.shape + assert both[1].shape == (2, *input.shape) + with pytest.raises(ValueError, match="At least one"): + distance_op(input, return_distances=False, return_indices=False) + + +@pytest.mark.parametrize("distance_op", DISTANCE_OPERATORS) +def test_preallocated_outputs_are_filled_and_not_returned(distance_op): + input = cuda(CASES[2]) + expected_distances, expected_indices = distance_op(input, return_indices=True) + distances = torch.empty_like(input, dtype=torch.float64) + indices = torch.empty((2, *input.shape), device="cuda", dtype=torch.int32) + + result = distance_op( + input, + return_distances=False, + return_indices=False, + distances=distances, + indices=indices, ) - # 4. Run SciPy (ground truth) - dist_ref_numpy, _ = batch_scipy_edt_with_indices(x_numpy_contiguous, spatial_ndim) - dist_ref = torch.as_tensor(dist_ref_numpy, device="cuda", dtype=torch.float32) - - torch.testing.assert_close(dist_cuda, dist_ref, atol=1e-5, rtol=1e-5) + assert result is None + torch.testing.assert_close(distances, expected_distances.double()) + torch.testing.assert_close(indices, expected_indices.int()) -@pytest.mark.parametrize( - "input_numpy, spatial_ndim, sampling", - [ - # 2D with non-uniform sampling - pytest.param(edt_case_2d_single, 2, [0.5, 1.0], id="2D_Sampling_0.5_1.0"), - pytest.param(edt_case_2d_single, 2, [2.0, 0.5], id="2D_Sampling_2.0_0.5"), - pytest.param(edt_case_2d_single, 2, [0.25, 0.25], id="2D_Sampling_0.25_0.25"), - # 2D batch with sampling - pytest.param(edt_case_2d, 2, [1.5, 0.75], id="2D_Batch_Sampling"), - # 3D with sampling - pytest.param(edt_case_3d, 3, [1.0, 2.0, 0.5], id="3D_Batch_Sampling"), - # 1D with sampling - pytest.param(edt_case_1d, 1, [0.5], id="1D_Batch_Sampling"), - # Test single-element list broadcast - pytest.param(edt_case_2d_single, 2, [0.5], id="2D_SingleElementList_Broadcast"), - pytest.param(edt_case_3d, 3, [2.0], id="3D_SingleElementList_Broadcast"), - ], -) -def test_edt_with_sampling( - input_numpy: np.ndarray, - spatial_ndim: int, - sampling: list[float], - request: pytest.FixtureRequest, -) -> None: - """Test EDT with non-unit sampling (pixel spacing).""" - x_numpy_contiguous = np.ascontiguousarray(input_numpy) - x_cuda = torch.as_tensor(x_numpy_contiguous, device="cuda") - - # Run CUDA EDT with sampling - dist_cuda, idx_cuda = tm.euclidean_distance_transform( - x_cuda.clone(), sampling=sampling, return_indices=True - ) - - # Expand single-element list for SciPy (it doesn't support broadcast) - scipy_sampling = sampling if len(sampling) == spatial_ndim else sampling * spatial_ndim - - # Run SciPy with sampling (ground truth) - dist_ref_numpy, _ = batch_scipy_edt_with_sampling( - x_numpy_contiguous, spatial_ndim, scipy_sampling +@pytest.mark.parametrize("distance_op", DISTANCE_OPERATORS) +def test_provided_output_auto_enables_only_that_output(distance_op): + input = cuda(CASES[2]) + distances = torch.empty_like(input) + result = distance_op( + input, + return_distances=False, + return_indices=True, + distances=distances, ) - dist_ref = torch.as_tensor(dist_ref_numpy, device="cuda", dtype=torch.float32) + assert isinstance(result, torch.Tensor) + assert result.shape == (2, *input.shape) - torch.testing.assert_close(dist_cuda, dist_ref, atol=1e-5, rtol=1e-5) +@pytest.mark.parametrize("distance_op", DISTANCE_OPERATORS) +def test_distance_transform_validates_input(distance_op): + with pytest.raises(ValueError, match="CUDA"): + distance_op(torch.zeros((1, 1, 3))) + with pytest.raises(ValueError, match="at least 3 dimensions"): + distance_op(torch.zeros((1, 3), device="cuda")) + with pytest.raises(ValueError, match="empty tensor"): + distance_op(torch.zeros((1, 1, 0), device="cuda")) + with pytest.raises(ValueError, match="spatial dimensions"): + distance_op(torch.zeros((1, 1, *([1] * 9)), device="cuda")) + + +@pytest.mark.parametrize("distance_op", DISTANCE_OPERATORS) +def test_distance_transform_validates_outputs(distance_op): + input = cuda(CASES[2]) + with pytest.raises(ValueError, match="distances shape"): + distance_op(input, distances=torch.empty((1, 1, 2, 2), device="cuda")) + with pytest.raises(ValueError, match="same device"): + distance_op(input, distances=torch.empty_like(input, device="cpu")) + with pytest.raises(ValueError, match="indices shape"): + distance_op(input, indices=torch.empty((2, 1, 1, 2, 2), device="cuda")) + with pytest.raises(ValueError, match="same device"): + distance_op( + input, + indices=torch.empty((2, *input.shape), device="cpu", dtype=torch.int64), + ) -def test_edt_return_flags() -> None: - """Test return_distances and return_indices flags.""" - # (B=1, C=1, H=2, W=3) - x = torch.as_tensor([[[[1, 1, 0], [1, 0, 0]]]], device="cuda", dtype=torch.float32) - - # Only distances - result = tm.euclidean_distance_transform(x, return_distances=True, return_indices=False) - assert isinstance( - result, torch.Tensor - ), "Should return single tensor when only distances requested" - assert result.shape == x.shape - - # Only indices - result = tm.euclidean_distance_transform(x, return_distances=False, return_indices=True) - assert isinstance( - result, torch.Tensor - ), "Should return single tensor when only indices requested" - assert result.shape == (2, *x.shape) # (spatial_ndim, B, C, H, W) - - # Both - dist, idx = tm.euclidean_distance_transform(x, return_distances=True, return_indices=True) - assert dist.shape == x.shape - assert idx.shape == (2, *x.shape) - - print(">> Return flags test passed.") - - -def test_edt_single_float_sampling() -> None: - """Test that a single float sampling value applies to all dimensions.""" - # Use edt_case_2d_single which is (B=1, C=1, H=4, W=4) format - x_numpy = edt_case_2d_single - x_cuda = torch.as_tensor(x_numpy, device="cuda") - - # Single float should apply to all spatial dimensions - dist_cuda = tm.euclidean_distance_transform(x_cuda, sampling=0.5) - - # Compare with scipy using [0.5, 0.5] - use batch helper for BCHW format - spatial_ndim = 2 - dist_ref_numpy, _ = batch_scipy_edt_with_sampling(x_numpy, spatial_ndim, [0.5, 0.5]) - dist_ref = torch.as_tensor(dist_ref_numpy, device="cuda", dtype=torch.float32) - - torch.testing.assert_close(dist_cuda, dist_ref, atol=1e-5, rtol=1e-5) - - -@pytest.mark.nd -@pytest.mark.parametrize( - "shape, spatial_ndim, algorithm", - [ - # 1D-3D cases - pytest.param((2, 1, 32), 1, "exact", id="1D_exact"), - pytest.param((2, 1, 32, 32), 2, "exact", id="2D_exact"), - pytest.param((1, 1, 16, 16, 16), 3, "exact", id="3D_exact"), - # 4D-6D cases - pytest.param((1, 1, 8, 8, 8, 8), 4, "exact", id="4D_exact"), - pytest.param((1, 1, 5, 5, 5, 5, 5), 5, "exact", id="5D_exact"), - pytest.param((1, 1, 4, 4, 4, 4, 4, 4), 6, "exact", id="6D_exact"), - ], -) -def test_edt_nd(shape: tuple, spatial_ndim: int, algorithm: str) -> None: - """Test EDT with ND data (up to 6D).""" - np.random.seed(42) - input_numpy = (np.random.rand(*shape) > 0.5).astype(np.float32) - x_cuda = torch.as_tensor(input_numpy, device="cuda") - - dist_cuda = tm.euclidean_distance_transform(x_cuda, algorithm=algorithm) - dist_scipy, _ = batch_scipy_edt_with_indices(input_numpy, spatial_ndim) - dist_ref = torch.as_tensor(dist_scipy, device="cuda", dtype=torch.float32) - torch.testing.assert_close(dist_cuda, dist_ref, atol=1e-5, rtol=1e-5) +SAMPLING_OPERATORS = [ + pytest.param(tm.euclidean_distance_transform, ndi.distance_transform_edt, {}, id="edt"), + pytest.param( + tm.brute_force_distance_transform, + ndi.distance_transform_bf, + {"metric": "euclidean"}, + id="bfdt", + ), +] +@pytest.mark.parametrize(("torch_op", "scipy_op", "kwargs"), SAMPLING_OPERATORS) @pytest.mark.parametrize( - "input_numpy, spatial_ndim, algorithm", + ("sampling", "normalized"), [ - # 2D tests with different algorithms - pytest.param(edt_case_2d, 2, "exact", id="2D_exact"), - pytest.param(edt_case_2d, 2, "jfa", id="2D_jfa"), - pytest.param(edt_case_2d, 2, "auto", id="2D_auto"), - pytest.param(edt_case_2d_single, 2, "exact", id="2D_single_exact"), - pytest.param(edt_case_2d_single, 2, "jfa", id="2D_single_jfa"), - pytest.param(edt_case_2d_single, 2, "auto", id="2D_single_auto"), - # 3D tests with different algorithms - pytest.param(edt_case_3d, 3, "exact", id="3D_exact"), - pytest.param(edt_case_3d, 3, "jfa", id="3D_jfa"), - pytest.param(edt_case_3d, 3, "auto", id="3D_auto"), + pytest.param(0.5, [0.5, 0.5], id="scalar"), + pytest.param([0.5], [0.5, 0.5], id="singleton"), + pytest.param([0.5, 2.0], [0.5, 2.0], id="per-axis"), ], ) -def test_edt_algorithm( - input_numpy: np.ndarray, - spatial_ndim: int, - algorithm: str, - request: pytest.FixtureRequest, -) -> None: - """Test EDT with different algorithm options (exact, jfa, auto).""" - x_numpy_contiguous = np.ascontiguousarray(input_numpy) - x_cuda = torch.as_tensor(x_numpy_contiguous, device="cuda") - - # Run CUDA EDT with specified algorithm - dist_cuda = tm.euclidean_distance_transform(x_cuda.clone(), algorithm=algorithm) - - # Run SciPy (ground truth) - dist_ref_numpy, _ = batch_scipy_edt_with_indices(x_numpy_contiguous, spatial_ndim) - dist_ref = torch.as_tensor(dist_ref_numpy, device="cuda", dtype=torch.float32) - - torch.testing.assert_close(dist_cuda, dist_ref, rtol=1e-5, atol=1e-5) - - -def test_edt_algorithm_fallback_with_sampling() -> None: - """Test that JFA falls back to exact when sampling is provided.""" - x_numpy = edt_case_2d_single - x_cuda = torch.as_tensor(x_numpy, device="cuda") - - # With non-unit sampling, JFA should fall back to exact algorithm - # Both should give same result - dist_jfa = tm.euclidean_distance_transform(x_cuda.clone(), sampling=[0.5, 1.0], algorithm="jfa") - dist_exact = tm.euclidean_distance_transform( - x_cuda.clone(), sampling=[0.5, 1.0], algorithm="exact" - ) +def test_sampling_matches_scipy(torch_op, scipy_op, kwargs, sampling, normalized): + input_array = CASES[2] + result = torch_op(cuda(input_array), sampling=sampling, **kwargs) + expected = scipy_batch(input_array, scipy_op, sampling=normalized, **kwargs) + torch.testing.assert_close(result.cpu(), torch.as_tensor(expected).float()) - # Compare with scipy - spatial_ndim = 2 - dist_ref_numpy, _ = batch_scipy_edt_with_sampling(x_numpy, spatial_ndim, [0.5, 1.0]) - dist_ref = torch.as_tensor(dist_ref_numpy, device="cuda", dtype=torch.float32) - torch.testing.assert_close(dist_jfa, dist_ref, atol=1e-5, rtol=1e-5) - torch.testing.assert_close(dist_exact, dist_ref, atol=1e-5, rtol=1e-5) +@pytest.mark.parametrize(("distance_op", "_", "kwargs"), SAMPLING_OPERATORS) +@pytest.mark.parametrize("sampling", [[1.0, 2.0, 3.0], 0.0, -1.0, float("inf"), float("nan")]) +def test_sampling_validation(distance_op, _, kwargs, sampling): + with pytest.raises(ValueError, match="sampling"): + distance_op(cuda(CASES[2]), sampling=sampling, **kwargs) -def test_edt_jfa_vs_exact_consistency() -> None: - """Test that JFA and exact produce similar results for unit sampling.""" - # Create a larger random test case - torch.manual_seed(42) - x = (torch.randn(2, 1, 64, 64, device="cuda") > 0).float() +@pytest.mark.parametrize("algorithm", ["exact", "jfa", "auto"]) +@pytest.mark.parametrize("spatial_ndim", [2, 3]) +def test_edt_algorithms_match_scipy(algorithm, spatial_ndim): + input_array = CASES[spatial_ndim] + result = tm.euclidean_distance_transform(cuda(input_array), algorithm=algorithm) + expected = scipy_batch(input_array, ndi.distance_transform_edt) + torch.testing.assert_close(result.cpu(), torch.as_tensor(expected).float()) - dist_exact = tm.euclidean_distance_transform(x, algorithm="exact") - dist_jfa = tm.euclidean_distance_transform(x, algorithm="jfa") - # JFA should be very close to exact for most pixels - # Allow for small differences due to JFA's approximate nature - diff = torch.abs(dist_exact - dist_jfa) - mean_diff = diff.mean().item() +def test_edt_jfa_falls_back_for_nonunit_sampling(): + input = cuda(CASES[2]) + kwargs = {"sampling": [0.5, 2.0]} + exact = tm.euclidean_distance_transform(input, algorithm="exact", **kwargs) + jfa = tm.euclidean_distance_transform(input, algorithm="jfa", **kwargs) + torch.testing.assert_close(jfa, exact) - # Most pixels should be exact or very close - assert mean_diff < 0.1, f"Mean difference too large: {mean_diff}" - -# ====================================================================== -# CDT Tests -# ====================================================================== @pytest.mark.parametrize( - "input_numpy, spatial_ndim, metric", - [ - pytest.param(cdt_case_1d, 1, "chessboard", id="1D_B2C1_chessboard"), - pytest.param(cdt_case_1d, 1, "taxicab", id="1D_B2C1_taxicab"), - pytest.param(cdt_case_2d_simple, 2, "chessboard", id="2D_B1C1_chessboard"), - pytest.param(cdt_case_2d_simple, 2, "taxicab", id="2D_B1C1_taxicab"), - pytest.param(cdt_case_2d_batch, 2, "chessboard", id="2D_B2C1_chessboard"), - pytest.param(cdt_case_2d_batch, 2, "taxicab", id="2D_B2C1_taxicab"), - pytest.param(cdt_case_checkerboard, 2, "chessboard", id="2D_checkerboard_chessboard"), - pytest.param(cdt_case_checkerboard, 2, "taxicab", id="2D_checkerboard_taxicab"), - pytest.param(cdt_case_3d_simple, 3, "chessboard", id="3D_B1C1_simple_chessboard"), - pytest.param(cdt_case_3d_simple, 3, "taxicab", id="3D_B1C1_simple_taxicab"), - pytest.param(cdt_case_3d_sphere, 3, "chessboard", id="3D_B1C1_sphere_chessboard"), - pytest.param(cdt_case_3d_sphere, 3, "taxicab", id="3D_B1C1_sphere_taxicab"), - pytest.param(cdt_case_3d_batch, 3, "chessboard", id="3D_B2C1_batch_chessboard"), - pytest.param(cdt_case_3d_batch, 3, "taxicab", id="3D_B2C1_batch_taxicab"), - ], + ("alias", "canonical"), [("cityblock", "taxicab"), ("manhattan", "taxicab")] ) -def test_cdt_basic( - input_numpy: np.ndarray, - spatial_ndim: int, - metric: str, - request: pytest.FixtureRequest, -) -> None: - """Test CDT distance computation against scipy with BCHW format.""" - x_numpy_contiguous = np.ascontiguousarray(input_numpy) - x_cuda = torch.as_tensor(x_numpy_contiguous, device="cuda") - - # Run torchmorph CDT - dist_cuda = tm.chamfer_distance_transform(x_cuda, metric=metric) - - # Run scipy CDT (ground truth) - dist_scipy, _ = batch_scipy_cdt(x_numpy_contiguous, metric=metric, spatial_ndim=spatial_ndim) - dist_ref = torch.as_tensor(dist_scipy, device="cuda", dtype=torch.float32) - - torch.testing.assert_close(dist_cuda, dist_ref, atol=1e-5, rtol=1e-5) +def test_cdt_metric_aliases(alias, canonical): + input = cuda(CASES[2]) + result = tm.chamfer_distance_transform(input, metric=alias) + expected = tm.chamfer_distance_transform(input, metric=canonical) + torch.testing.assert_close(result, expected) @pytest.mark.parametrize( - "alias, canonical", + ("distance_op", "kwargs"), [ - pytest.param("cityblock", "taxicab", id="cityblock"), - pytest.param("manhattan", "taxicab", id="manhattan"), + pytest.param(tm.chamfer_distance_transform, {"metric": "invalid"}, id="cdt"), + pytest.param(tm.brute_force_distance_transform, {"metric": "invalid"}, id="bfdt"), ], ) -def test_cdt_metric_aliases(alias: str, canonical: str) -> None: - """Test that metric aliases produce same results.""" - x_cuda = torch.as_tensor(cdt_case_2d_simple, device="cuda") - - dist_alias = tm.chamfer_distance_transform(x_cuda, metric=alias) - dist_canonical = tm.chamfer_distance_transform(x_cuda, metric=canonical) - - torch.testing.assert_close(dist_alias, dist_canonical, atol=1e-5, rtol=1e-5) - - -def test_cdt_return_flags() -> None: - """Test return_distances and return_indices flags with BCHW format.""" - # (B=1, C=1, H=5, W=6) - x = torch.as_tensor(cdt_case_2d_simple, device="cuda") - - # Only distances (default) - result = tm.chamfer_distance_transform(x, return_distances=True, return_indices=False) - assert isinstance(result, torch.Tensor), "Should return single tensor" - assert result.shape == x.shape - - # Only indices - spatial_ndim=2 for BCHW - result = tm.chamfer_distance_transform(x, return_distances=False, return_indices=True) - assert isinstance(result, torch.Tensor), "Should return single tensor" - assert result.shape == (2, *x.shape) # (spatial_ndim, B, C, H, W) - - # Both - dist, idx = tm.chamfer_distance_transform(x, return_distances=True, return_indices=True) - assert dist.shape == x.shape - assert idx.shape == (2, *x.shape) - - print(">> Return flags test passed.") - - -def test_cdt_preallocated_output() -> None: - """Test pre-allocated output tensors with scipy-style return convention.""" - # (B=1, C=1, H=5, W=6) - x = torch.as_tensor(cdt_case_2d_simple, device="cuda") - - # Pre-allocate distances tensor - dist_out = torch.empty_like(x) - result = tm.chamfer_distance_transform(x, distances=dist_out) - - # Should return None (scipy convention) - assert result is None, "Should return None when distances tensor is provided" - - # But dist_out should be filled - dist_ref, _ = batch_scipy_cdt(cdt_case_2d_simple, metric="chessboard", spatial_ndim=2) - dist_ref_tensor = torch.as_tensor(dist_ref, device="cuda", dtype=torch.float32) - torch.testing.assert_close(dist_out, dist_ref_tensor, atol=1e-5, rtol=1e-5) - - print(">> Pre-allocated output test passed.") - - -def test_cdt_indices_correctness() -> None: - """Test that indices point to correct nearest background pixel with BCHW format.""" - # (B=1, C=1, H=5, W=6) - x = torch.as_tensor(cdt_case_2d_simple, device="cuda") - - dist, idx = tm.chamfer_distance_transform(x, metric="chessboard", return_indices=True) - - # For each foreground pixel, verify the index points to a background pixel - # idx shape: (spatial_ndim=2, B=1, C=1, H=5, W=6) - B, C, H, W = x.shape - x_np = x.cpu().numpy() - idx_np = idx.cpu().numpy() # (2, B, C, H, W) - dist_np = dist.cpu().numpy() - - for b in range(B): - for c in range(C): - for y in range(H): - for x_coord in range(W): - if x_np[b, c, y, x_coord] != 0: # Foreground - idx_y = idx_np[0, b, c, y, x_coord] - idx_x = idx_np[1, b, c, y, x_coord] - # The pointed pixel should be background - assert ( - x_np[b, c, idx_y, idx_x] == 0 - ), f"Index ({idx_y}, {idx_x}) should point to background" - # Chessboard distance should match - expected_dist = max(abs(y - idx_y), abs(x_coord - idx_x)) - assert ( - dist_np[b, c, y, x_coord] == expected_dist - ), f"Distance mismatch at ({b}, {c}, {y}, {x_coord})" - - print(">> Indices correctness test passed.") - - -def test_cdt_indices_correctness_3d() -> None: - """Test that 3D indices point to correct nearest background pixel with BCDHW format.""" - # (B=1, C=1, D=5, H=5, W=5) - x = torch.as_tensor(cdt_case_3d_simple, device="cuda") - - dist, idx = tm.chamfer_distance_transform(x, metric="chessboard", return_indices=True) - - # For each foreground pixel, verify the index points to a background pixel - # idx shape: (spatial_ndim=3, B=1, C=1, D=5, H=5, W=5) - B, C, D, H, W = x.shape - x_np = x.cpu().numpy() - idx_np = idx.cpu().numpy() # (3, B, C, D, H, W) - dist_np = dist.cpu().numpy() - - for b in range(B): - for c in range(C): - for z in range(D): - for y in range(H): - for x_coord in range(W): - if x_np[b, c, z, y, x_coord] != 0: # Foreground - idx_z = idx_np[0, b, c, z, y, x_coord] - idx_y = idx_np[1, b, c, z, y, x_coord] - idx_x = idx_np[2, b, c, z, y, x_coord] - # The pointed pixel should be background - assert ( - x_np[b, c, idx_z, idx_y, idx_x] == 0 - ), f"Index ({idx_z}, {idx_y}, {idx_x}) should point to background" - # Chessboard distance should match - expected_dist = max( - abs(z - idx_z), abs(y - idx_y), abs(x_coord - idx_x) - ) - assert ( - dist_np[b, c, z, y, x_coord] == expected_dist - ), f"Distance mismatch at ({b}, {c}, {z}, {y}, {x_coord})" - - print(">> 3D Indices correctness test passed.") - - -def test_cdt_invalid_metric() -> None: - """Test that invalid metric raises error.""" - x = torch.as_tensor(cdt_case_2d_simple, device="cuda") - +def test_invalid_metric(distance_op, kwargs): with pytest.raises(ValueError, match="metric must be"): - tm.chamfer_distance_transform(x, metric="invalid") - - -def test_cdt_cpu_input_error() -> None: - """Test that CPU input raises error.""" - x = torch.as_tensor(cdt_case_2d_simple, device="cpu") # CPU tensor - - with pytest.raises(ValueError, match="CUDA"): - tm.chamfer_distance_transform(x) - - -@pytest.mark.nd -@pytest.mark.parametrize( - "shape, spatial_ndim, metric", - [ - # 1D-3D cases - pytest.param((2, 1, 32), 1, "chessboard", id="1D_B2C1_32_chessboard"), - pytest.param((2, 1, 32), 1, "taxicab", id="1D_B2C1_32_taxicab"), - pytest.param((2, 1, 32, 32), 2, "chessboard", id="2D_B2C1_32x32_chessboard"), - pytest.param((2, 1, 32, 32), 2, "taxicab", id="2D_B2C1_32x32_taxicab"), - pytest.param((1, 1, 16, 16, 16), 3, "chessboard", id="3D_B1C1_16x16x16_chessboard"), - pytest.param((1, 1, 16, 16, 16), 3, "taxicab", id="3D_B1C1_16x16x16_taxicab"), - # 4D-6D cases - pytest.param((1, 1, 8, 8, 8, 8), 4, "chessboard", id="4D_B1C1_chessboard"), - pytest.param((1, 1, 8, 8, 8, 8), 4, "taxicab", id="4D_B1C1_taxicab"), - pytest.param((1, 1, 5, 5, 5, 5, 5), 5, "chessboard", id="5D_B1C1_chessboard"), - pytest.param((1, 1, 4, 4, 4, 4, 4, 4), 6, "chessboard", id="6D_B1C1_chessboard"), - ], -) -def test_cdt_nd(shape: tuple, spatial_ndim: int, metric: str) -> None: - """Test CDT with ND data (up to 6D).""" - np.random.seed(42) - input_numpy = (np.random.rand(*shape) > 0.5).astype(np.float32) - x_cuda = torch.as_tensor(input_numpy, device="cuda") - - dist_cuda = tm.chamfer_distance_transform(x_cuda, metric=metric) - dist_scipy, _ = batch_scipy_cdt(input_numpy, metric=metric, spatial_ndim=spatial_ndim) - dist_ref = torch.as_tensor(dist_scipy, device="cuda", dtype=torch.float32) - - torch.testing.assert_close(dist_cuda, dist_ref, atol=1e-5, rtol=1e-5) - - -@pytest.mark.parametrize( - "input_numpy, spatial_ndim, metric", - [ - pytest.param(cdt_case_1d, 1, "chessboard", id="1D_indices_chessboard"), - pytest.param(cdt_case_1d, 1, "taxicab", id="1D_indices_taxicab"), - pytest.param(cdt_case_2d_batch, 2, "chessboard", id="2D_batch_indices_chessboard"), - pytest.param(cdt_case_3d_batch, 3, "chessboard", id="3D_batch_indices_chessboard"), - ], -) -def test_cdt_indices_validation( - input_numpy: np.ndarray, - spatial_ndim: int, - metric: str, - request: pytest.FixtureRequest, -) -> None: - """Test that indices correctly point to nearest background pixels in BCHW format.""" - x_numpy_contiguous = np.ascontiguousarray(input_numpy) - x_cuda = torch.as_tensor(x_numpy_contiguous, device="cuda") - - # Run torchmorph CDT with indices - dist_cuda, idx_cuda = tm.chamfer_distance_transform(x_cuda, metric=metric, return_indices=True) - - # Validate indices shape: (spatial_ndim, *input_shape) - expected_idx_shape = (spatial_ndim, *x_cuda.shape) - assert ( - idx_cuda.shape == expected_idx_shape - ), f"Index shape mismatch: {idx_cuda.shape} vs {expected_idx_shape}" - - # Validate that indices point to background pixels and distance matches - spatial_shape = x_cuda.shape[-spatial_ndim:] - batch_shape = x_cuda.shape[:-spatial_ndim] - - # Create coordinate grid for spatial dimensions - coords = [torch.arange(s, device="cuda") for s in spatial_shape] - grid = torch.stack( - torch.meshgrid(*coords, indexing="ij"), dim=0 - ) # (spatial_ndim, *spatial_shape) - - # Expand grid for batch dimensions - for _ in batch_shape: - grid = grid.unsqueeze(1) - grid = grid.expand(spatial_ndim, *batch_shape, *spatial_shape) - - # Calculate distance from indices based on metric - diff = grid.float() - idx_cuda.float() - if metric in ("chessboard",): - # Chessboard: max of absolute differences - dist_calculated = torch.max(torch.abs(diff), dim=0).values - else: - # Taxicab: sum of absolute differences - dist_calculated = torch.sum(torch.abs(diff), dim=0) - - torch.testing.assert_close(dist_calculated, dist_cuda, atol=1e-5, rtol=1e-5) - print(">> Index validation passed.") - - -# ====================================================================== -# BFDT Tests -# ====================================================================== -@pytest.mark.bfdt -@pytest.mark.parametrize( - "input_numpy, spatial_ndim, metric", - [ - pytest.param(cdt_case_1d, 1, "euclidean", id="1D_B2C1_euclidean"), - pytest.param(cdt_case_1d, 1, "taxicab", id="1D_B2C1_taxicab"), - pytest.param(cdt_case_1d, 1, "chessboard", id="1D_B2C1_chessboard"), - pytest.param(cdt_case_2d_batch, 2, "euclidean", id="2D_B2C1_euclidean"), - pytest.param(cdt_case_2d_batch, 2, "taxicab", id="2D_B2C1_taxicab"), - pytest.param(cdt_case_2d_batch, 2, "chessboard", id="2D_B2C1_chessboard"), - pytest.param(cdt_case_3d_batch, 3, "euclidean", id="3D_B2C1_euclidean"), - ], -) -def test_bfdt_basic( - input_numpy: np.ndarray, - spatial_ndim: int, - metric: str, - request: pytest.FixtureRequest, -) -> None: - """Test BFDT distance and indices against scipy.""" - x_numpy_contiguous = np.ascontiguousarray(input_numpy) - x_cuda = torch.as_tensor(x_numpy_contiguous, device="cuda") - - # Run torchmorph BFDT - dist_cuda, idx_cuda = tm.brute_force_distance_transform( - x_cuda, metric=metric, return_indices=True - ) - - # Run scipy BFDT (ground truth for distances) - dist_scipy, _ = batch_scipy_bfdt(x_numpy_contiguous, metric=metric, spatial_ndim=spatial_ndim) - dist_ref = torch.as_tensor(dist_scipy, device="cuda", dtype=torch.float32) - - # Validate distances against scipy - torch.testing.assert_close(dist_cuda, dist_ref, atol=1e-5, rtol=1e-5) - - # Validate indices semantically: distance computed from indices must match dist_cuda - # (direct index comparison can fail due to tie-breaking differences between implementations) - expected_idx_shape = (spatial_ndim, *x_cuda.shape) - assert ( - idx_cuda.shape == expected_idx_shape - ), f"Index shape mismatch: {idx_cuda.shape} vs {expected_idx_shape}" - - spatial_shape = x_cuda.shape[-spatial_ndim:] - batch_shape = x_cuda.shape[:-spatial_ndim] - coords = [torch.arange(s, device="cuda") for s in spatial_shape] - grid = torch.stack(torch.meshgrid(*coords, indexing="ij"), dim=0) - for _ in batch_shape: - grid = grid.unsqueeze(1) - grid = grid.expand(spatial_ndim, *batch_shape, *spatial_shape) - - diff = grid.float() - idx_cuda.float() - if metric == "euclidean": - dist_from_idx = torch.sqrt(torch.sum(diff**2, dim=0)) - elif metric == "taxicab": - dist_from_idx = torch.sum(torch.abs(diff), dim=0) - else: # chessboard - dist_from_idx = torch.max(torch.abs(diff), dim=0).values - - torch.testing.assert_close(dist_from_idx, dist_cuda, atol=1e-5, rtol=1e-5) - - -@pytest.mark.bfdt -def test_bfdt_sampling() -> None: - """Test BFDT with non-unit sampling.""" - x_numpy = cdt_case_2d_simple - x_cuda = torch.as_tensor(x_numpy, device="cuda") - sampling = [0.5, 2.0] - - dist_cuda = tm.brute_force_distance_transform(x_cuda, metric="euclidean", sampling=sampling) - - dist_scipy, _ = batch_scipy_bfdt(x_numpy, metric="euclidean", sampling=sampling, spatial_ndim=2) - dist_ref = torch.as_tensor(dist_scipy, device="cuda", dtype=torch.float32) - - torch.testing.assert_close(dist_cuda, dist_ref, atol=1e-5, rtol=1e-5) - - -@pytest.mark.bfdt -@pytest.mark.nd -@pytest.mark.parametrize( - "shape, spatial_ndim, metric", - [ - pytest.param((2, 1, 16), 1, "euclidean", id="1D_B2C1_16_euclidean"), - pytest.param((1, 1, 16, 16), 2, "euclidean", id="2D_B1C1_16x16_euclidean"), - pytest.param((1, 1, 8, 8, 8), 3, "euclidean", id="3D_B1C1_8x8x8_euclidean"), - pytest.param((1, 1, 6, 6, 6, 6), 4, "euclidean", id="4D_B1C1_6x6x6x6_euclidean"), - pytest.param((1, 1, 4, 4, 4, 4, 4), 5, "euclidean", id="5D_B1C1_4x4x4x4x4_euclidean"), - pytest.param((1, 1, 3, 3, 3, 3, 3, 3), 6, "euclidean", id="6D_B1C1_3x3x3x3x3x3_euclidean"), - ], -) -def test_bfdt_nd(shape: tuple, spatial_ndim: int, metric: str) -> None: - """Test BFDT with ND data (up to 6D).""" - np.random.seed(42) - input_numpy = (np.random.rand(*shape) > 0.5).astype(np.float32) - x_cuda = torch.as_tensor(input_numpy, device="cuda") + distance_op(cuda(CASES[2]), **kwargs) - dist_cuda = tm.brute_force_distance_transform(x_cuda, metric=metric) - dist_scipy, _ = batch_scipy_bfdt(input_numpy, metric=metric, spatial_ndim=spatial_ndim) - dist_ref = torch.as_tensor(dist_scipy, device="cuda", dtype=torch.float32) - torch.testing.assert_close(dist_cuda, dist_ref, atol=1e-5, rtol=1e-5) +@pytest.mark.skipif(torch.cuda.device_count() < 2, reason="requires two CUDA devices") +@pytest.mark.parametrize("distance_op", DISTANCE_OPERATORS) +def test_distance_transform_uses_input_device(distance_op): + input = torch.as_tensor(CASES[2], device="cuda:1") + result = distance_op(input) + assert result.device == input.device diff --git a/test/test_grey.py b/test/test_grey.py index c9236fd..686e577 100644 --- a/test/test_grey.py +++ b/test/test_grey.py @@ -154,6 +154,23 @@ def test_grey_morphology_output(torch_op, scipy_op): torch.testing.assert_close(result.cpu(), torch.as_tensor(expected)) +@pytest.mark.parametrize("torch_op", TORCH_OPERATORS) +def test_grey_morphology_requires_output_shape_and_device(torch_op): + x = torch.as_tensor(CASE_2D, device="cuda") + + with pytest.raises(ValueError, match="output shape"): + torch_op(x, size=3, output=torch.empty((1, 1, 2, 2), device="cuda")) + with pytest.raises(ValueError, match="same device"): + torch_op(x, size=3, output=torch.empty_like(x, device="cpu")) + + +@pytest.mark.parametrize("torch_op", TORCH_OPERATORS) +def test_grey_morphology_rejects_more_than_eight_spatial_dimensions(torch_op): + x = torch.ones((1, 1, *([2] * 9)), device="cuda") + with pytest.raises(ValueError, match="spatial dimensions"): + torch_op(x, size=3) + + @pytest.mark.parametrize("torch_op", TORCH_OPERATORS) def test_grey_morphology_requires_structuring_element(torch_op): x = torch.as_tensor(CASE_2D, device="cuda") diff --git a/torchmorph/distance_transform.py b/torchmorph/distance_transform.py index 03bc1d1..08152d6 100644 --- a/torchmorph/distance_transform.py +++ b/torchmorph/distance_transform.py @@ -1,382 +1,200 @@ +import math from collections.abc import Sequence -import torch +from torch import Tensor -from torchmorph import _C +from . import _C -def euclidean_distance_transform( - input: torch.Tensor, - sampling: float | Sequence[float] | None = None, - return_distances: bool = True, - return_indices: bool = False, - distances: torch.Tensor | None = None, - indices: torch.Tensor | None = None, - algorithm: str = "exact", -) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor] | None: - """Exact Euclidean Distance Transform (EDT) using Felzenszwalb algorithm. - - Args: - input: Binary input tensor (0 = background, non-zero = foreground). - Must be in (B, C, Spatial...) format where Spatial can be 1D, 2D, or 3D. - For single images, use unsqueeze to add batch and channel dims. - sampling: Spacing of elements along each spatial dimension. If a single - number, the spacing is uniform in all spatial dimensions. If a - sequence, it must match the number of spatial dimensions. - Default is None (unit spacing for all spatial dimensions). - Note: When sampling is not unit spacing, only "exact" algorithm is used. - return_distances: Whether to calculate the distance transform. - Default is True. - return_indices: Whether to calculate the feature transform (indices - of closest background element). Default is False. - distances: Optional output tensor for distances. If provided, must have - the same shape as input. If None and return_distances is True, - a new tensor will be created and returned. - indices: Optional output tensor for indices. If provided, must have shape - (spatial_ndim, ...) where ... matches input shape. If None and - return_indices is True, a new tensor will be created and returned. - algorithm: Algorithm to use for distance transform. Options: - - "exact": Use Felzenszwalb's exact algorithm (default). - - "jfa": Use Jump Flooding Algorithm (fast but approximate). - Only available for 2D/3D with unit sampling. - - "auto": Automatically choose based on input (uses JFA when - applicable, otherwise exact). - - Returns: - Depending on return_distances, return_indices, and whether output tensors - are provided: - - Returns distance tensor only when return_distances=True and distances=None - - Returns indices tensor only when return_indices=True and indices=None - - Returns tuple of (distances, indices) when both conditions above are met - - Returns None if output tensors are provided for all requested outputs - - Example: - >>> import torchmorph as tm - >>> # 2D image: (B, C, H, W) - >>> x = torch.zeros(1, 1, 64, 64, device='cuda') - >>> x[0, 0, 10:20, 10:20] = 1 - >>> dist = tm.euclidean_distance_transform(x) - >>> dist, indices = tm.euclidean_distance_transform(x, return_indices=True) - >>> dist = tm.euclidean_distance_transform(x, sampling=[0.5, 1.0]) - >>> # Using JFA algorithm (faster for large images) - >>> dist = tm.euclidean_distance_transform(x, algorithm="jfa") - >>> # Using pre-allocated output tensors - >>> dist_out = torch.empty_like(x) - >>> tm.euclidean_distance_transform(x, distances=dist_out) # Returns None, fills dist_out - >>> # 3D volume: (B, C, D, H, W) - >>> x_3d = torch.zeros(2, 1, 32, 64, 64, device='cuda') - >>> dist_3d = tm.euclidean_distance_transform(x_3d, sampling=[2.0, 1.0, 1.0]) - """ +def _normalize_sampling( + sampling: float | Sequence[float] | None, + spatial_ndim: int, +) -> list[float]: + if sampling is None: + values = [1.0] * spatial_ndim + elif isinstance(sampling, (int, float)): + values = [float(sampling)] * spatial_ndim + else: + values = [float(value) for value in sampling] + if len(values) == 1: + values *= spatial_ndim + elif len(values) != spatial_ndim: + raise ValueError(f"sampling must have length 1 or {spatial_ndim}, got {len(values)}") + + if any(not math.isfinite(value) or value <= 0 for value in values): + raise ValueError("sampling values must be finite and greater than zero") + return values + + +def _prepare_distance_transform( + input: Tensor, + return_distances: bool, + return_indices: bool, + distances: Tensor | None, + indices: Tensor | None, +) -> tuple[int, bool, bool]: if not input.is_cuda: raise ValueError("Input tensor must be on CUDA device.") if input.ndim < 3: raise ValueError( - f"Input must be (B, C, ) format with at least 3 dimensions, got {input.shape}. " - "For single images, use unsqueeze to add batch and channel dims." + f"Input must be (B, C, Spatial...) with at least 3 dimensions, got {input.shape}." ) + spatial_ndim = input.ndim - 2 + if spatial_ndim > 8: + raise ValueError(f"Input spatial dimensions must be in range 1 to 8, got {spatial_ndim}.") if input.numel() == 0: raise ValueError(f"Invalid input: empty tensor with shape {input.shape}.") - # Validate pre-allocated output tensors - if distances is not None: - if distances.shape != input.shape: + expected_indices_shape = (spatial_ndim, *input.shape) + for name, output, expected_shape in ( + ("distances", distances, input.shape), + ("indices", indices, expected_indices_shape), + ): + if output is None: + continue + if tuple(output.shape) != tuple(expected_shape): raise ValueError( - f"distances shape {distances.shape} must match input shape {input.shape}" + f"{name} shape {output.shape} must match expected shape {expected_shape}" ) - if not distances.is_cuda: - raise ValueError("distances tensor must be on CUDA device.") - return_distances = True - - if indices is not None: - if not indices.is_cuda: - raise ValueError("indices tensor must be on CUDA device.") - return_indices = True - - if not return_distances and not return_indices: - raise ValueError( - "At least one of return_distances or return_indices must be True, " - "or output tensors must be provided." - ) - - input = input.float().contiguous() - total_ndim = input.ndim - spatial_ndim = total_ndim - 2 # Exclude B and C dimensions - - # Process sampling parameter for spatial dimensions only - if sampling is None: - # Unit spacing for all spatial dimensions - sampling_list = [1.0] * spatial_ndim - elif isinstance(sampling, (int, float)): - # Single value: same spacing for all spatial dimensions - sampling_list = [float(sampling)] * spatial_ndim - else: - # Sequence: convert to list - sampling_list = [float(s) for s in sampling] - if len(sampling_list) == 1: - # Single element list: broadcast to all spatial dimensions - sampling_list = sampling_list * spatial_ndim - elif len(sampling_list) != spatial_ndim: + if output.device != input.device: raise ValueError( - f"sampling has {len(sampling_list)} but input {spatial_ndim} dimensions " - f"(input shape: {input.shape}, format: (B, C, Spatial...))" + f"{name} must be on the same device as input, " + f"got {output.device} and {input.device}" ) - # Call CUDA kernel - it handles batch dimensions based on sampling size - raw_distances, raw_indices = _C.edt_cuda( - input, sampling_list, return_distances, return_indices, algorithm - ) - - # Copy to pre-allocated tensors if provided + return_distances = return_distances or distances is not None + return_indices = return_indices or indices is not None + if not return_distances and not return_indices: + raise ValueError("At least one distance transform output must be requested.") + return spatial_ndim, return_distances, return_indices + + +def _finish_distance_transform( + raw_distances: Tensor | None, + raw_indices: Tensor | None, + return_distances: bool, + return_indices: bool, + distances: Tensor | None, + indices: Tensor | None, +) -> Tensor | tuple[Tensor, Tensor] | None: if distances is not None and raw_distances is not None: distances.copy_(raw_distances) - if indices is not None and raw_indices is not None: indices.copy_(raw_indices) - # Return based on scipy convention: - # Only return tensors that were NOT provided by the user - return_dist_tensor = return_distances and distances is None - return_idx_tensor = return_indices and indices is None - - if return_dist_tensor and return_idx_tensor: + returned_distances = return_distances and distances is None + returned_indices = return_indices and indices is None + if returned_distances and returned_indices: return raw_distances, raw_indices - elif return_dist_tensor: + if returned_distances: return raw_distances - elif return_idx_tensor: + if returned_indices: return raw_indices - else: - return None + return None -def chamfer_distance_transform( - input: torch.Tensor, - metric: str = "chessboard", +def euclidean_distance_transform( + input: Tensor, + sampling: float | Sequence[float] | None = None, return_distances: bool = True, return_indices: bool = False, - distances: torch.Tensor | None = None, - indices: torch.Tensor | None = None, -) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor] | None: - """Chamfer Distance Transform (CDT). - - Calculates the distance transform of the input using a chamfer metric. - The input is treated as a binary image where non-zero values are foreground - and zero values are background. Distances are computed from each foreground - pixel to the nearest background pixel. - - Args: - input: Binary input tensor (0 = background, non-zero = foreground). - Must be in (B, C, H, W) or (B, C, D, H, W) format for batch processing, - or (H, W) / (D, H, W) for single images. - metric: Distance metric to use: - - "chessboard": L-infinity norm (default). Also known as Chebyshev distance. - - "taxicab": L1 norm. Also known as Manhattan or city-block distance. - - "cityblock": Alias for "taxicab". - - "manhattan": Alias for "taxicab". - return_distances: Whether to calculate the distance transform. Default is True. - return_indices: Whether to calculate the feature transform (indices of closest - background element). Default is False. - distances: Optional output tensor for distances. If provided, must have - the same shape as input. If None and return_distances is True, - a new tensor will be created. - indices: Optional output tensor for indices. If provided, must have shape - (..., ndim) where ... matches input shape. If None and return_indices - is True, a new tensor will be created. - - Returns: - Depending on return_distances, return_indices, and whether output tensors - are provided: - - Returns distance tensor only when return_distances=True and distances=None - - Returns indices tensor only when return_indices=True and indices=None - - Returns tuple of (distances, indices) when both conditions above are met - - Returns None if output tensors are provided for all requested outputs - - Example: - >>> import torchmorph as tm - >>> # 2D image with batch: (B, C, H, W) - >>> x = torch.zeros(1, 1, 64, 64, device='cuda') - >>> x[0, 0, 10:20, 10:20] = 1 - >>> dist = tm.chamfer_distance_transform(x) # chessboard by default - >>> dist = tm.chamfer_distance_transform(x, metric='taxicab') - >>> dist, indices = tm.chamfer_distance_transform(x, return_indices=True) - >>> # Using pre-allocated output tensors - >>> dist_out = torch.empty_like(x) - >>> tm.chamfer_distance_transform(x, distances=dist_out) # Returns None, fills dist_out - """ - if not input.is_cuda: - raise ValueError("Input tensor must be on CUDA device.") - if input.ndim < 2 or input.numel() == 0: - raise ValueError(f"Invalid input dimension: {input.shape}.") + distances: Tensor | None = None, + indices: Tensor | None = None, + algorithm: str = "auto", +) -> Tensor | tuple[Tensor, Tensor] | None: + """Euclidean distance transform for (B, C, Spatial...) CUDA tensors.""" + spatial_ndim, return_distances, return_indices = _prepare_distance_transform( + input, + return_distances, + return_indices, + distances, + indices, + ) + raw_distances, raw_indices = _C.edt_cuda( + input.float().contiguous(), + _normalize_sampling(sampling, spatial_ndim), + return_distances, + return_indices, + algorithm, + ) + return _finish_distance_transform( + raw_distances, + raw_indices, + return_distances, + return_indices, + distances, + indices, + ) - # Normalize metric aliases - if metric in ("cityblock", "manhattan"): - metric = "taxicab" - if metric not in ("chessboard", "taxicab"): +def chamfer_distance_transform( + input: Tensor, + metric: str = "chessboard", + return_distances: bool = True, + return_indices: bool = False, + distances: Tensor | None = None, + indices: Tensor | None = None, +) -> Tensor | tuple[Tensor, Tensor] | None: + """Chamfer distance transform for (B, C, Spatial...) CUDA tensors.""" + _, return_distances, return_indices = _prepare_distance_transform( + input, + return_distances, + return_indices, + distances, + indices, + ) + metric = {"cityblock": "taxicab", "manhattan": "taxicab"}.get(metric, metric) + if metric not in {"chessboard", "taxicab"}: raise ValueError("metric must be 'chessboard', 'taxicab', 'cityblock', or 'manhattan'.") - if not return_distances and not return_indices: - if distances is None and indices is None: - raise ValueError( - "At least one of return_distances or return_indices must be True, " - "or output tensors must be provided." - ) - - input = input.float().contiguous() - - # Validate pre-allocated output tensors - if distances is not None: - if distances.shape != input.shape: - raise ValueError( - f"distances shape {distances.shape} must match input shape {input.shape}" - ) - if not distances.is_cuda: - raise ValueError("distances tensor must be on CUDA device.") - return_distances = True - - if indices is not None: - if not indices.is_cuda: - raise ValueError("indices tensor must be on CUDA device.") - return_indices = True - # Call CUDA kernel - raw_distances, raw_indices = _C.cdt_cuda(input, metric, return_distances, return_indices) - - # Copy to pre-allocated tensors if provided - if distances is not None and raw_distances is not None: - distances.copy_(raw_distances) - - if indices is not None and raw_indices is not None: - indices.copy_(raw_indices) - - # Return based on scipy convention: - # Only return tensors that were NOT provided by the user - return_dist_tensor = return_distances and distances is None - return_idx_tensor = return_indices and indices is None - - if return_dist_tensor and return_idx_tensor: - return raw_distances, raw_indices - elif return_dist_tensor: - return raw_distances - elif return_idx_tensor: - return raw_indices - else: - return None + raw_distances, raw_indices = _C.cdt_cuda( + input.float().contiguous(), + metric, + return_distances, + return_indices, + ) + return _finish_distance_transform( + raw_distances, + raw_indices, + return_distances, + return_indices, + distances, + indices, + ) def brute_force_distance_transform( - input: torch.Tensor, + input: Tensor, metric: str = "euclidean", sampling: float | Sequence[float] | None = None, return_distances: bool = True, return_indices: bool = False, - distances: torch.Tensor | None = None, - indices: torch.Tensor | None = None, -) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor] | None: - """Brute-force distance transform. - - Calculates the distance transform of the input using a brute-force algorithm. - The algorithm computes the distance from each foreground pixel to ALL - background pixels and finds the minimum. This is $O(N*M)$ where $N$ is - number of foreground pixels and $M$ is number of background pixels. - - Args: - input: Binary input tensor (0 = background, non-zero = foreground). - Must be in (B, C, Spatial...) format. - metric: Distance metric to use: - - "euclidean": L2 norm (default). - - "taxicab": L1 norm (Manhattan distance). - - "chessboard": L-infinity norm (Chebyshev distance). - sampling: Spacing of elements along each spatial dimension. If a single - number, the spacing is uniform in all spatial dimensions. If a - sequence, it must match the number of spatial dimensions. - Default is None (unit spacing for all spatial dimensions). - return_distances: Whether to calculate the distance transform. - Default is True. - return_indices: Whether to calculate the feature transform (indices - of closest background element). Default is False. - distances: Optional output tensor for distances. If provided, must have - the same shape as input. - indices: Optional output tensor for indices. If provided, must have shape - (spatial_ndim, ...) where ... matches input shape. - - Returns: - Depending on return_distances, return_indices, and whether output tensors - are provided: - - Returns distance tensor only when return_distances=True and distances=None - - Returns indices tensor only when return_indices=True and indices=None - - Returns tuple of (distances, indices) when both conditions above are met - - Returns None if output tensors are provided for all requested outputs - """ - if not input.is_cuda: - raise ValueError("Input tensor must be on CUDA device.") - if input.ndim < 3: - raise ValueError( - f"Input must be (B, C, ) format with at least 3 dimensions, got {input.shape}." - ) - if input.numel() == 0: - raise ValueError(f"Invalid input: empty tensor with shape {input.shape}.") - - if metric not in ("euclidean", "taxicab", "chessboard"): + distances: Tensor | None = None, + indices: Tensor | None = None, +) -> Tensor | tuple[Tensor, Tensor] | None: + """Brute-force distance transform for (B, C, Spatial...) CUDA tensors.""" + spatial_ndim, return_distances, return_indices = _prepare_distance_transform( + input, + return_distances, + return_indices, + distances, + indices, + ) + if metric not in {"euclidean", "taxicab", "chessboard"}: raise ValueError("metric must be 'euclidean', 'taxicab', or 'chessboard'.") - # Validate pre-allocated output tensors - if distances is not None: - if distances.shape != input.shape: - raise ValueError( - f"distances shape {distances.shape} must match input shape {input.shape}" - ) - if not distances.is_cuda: - raise ValueError("distances tensor must be on CUDA device.") - return_distances = True - - if indices is not None: - if not indices.is_cuda: - raise ValueError("indices tensor must be on CUDA device.") - return_indices = True - - if not return_distances and not return_indices: - raise ValueError( - "At least one of return_distances or return_indices must be True, " - "or output tensors must be provided." - ) - - input = input.float().contiguous() - spatial_ndim = input.ndim - 2 - - # Process sampling parameter - if sampling is None: - sampling_list = [1.0] * spatial_ndim - elif isinstance(sampling, (int, float)): - sampling_list = [float(sampling)] * spatial_ndim - else: - sampling_list = [float(s) for s in sampling] - if len(sampling_list) == 1: - sampling_list = sampling_list * spatial_ndim - elif len(sampling_list) != spatial_ndim: - raise ValueError( - f"sampling has {len(sampling_list)} but input {spatial_ndim} dimensions." - ) - - # Call CUDA kernel raw_distances, raw_indices = _C.bfdt_cuda( - input, metric, sampling_list, return_distances, return_indices + input.float().contiguous(), + metric, + _normalize_sampling(sampling, spatial_ndim), + return_distances, + return_indices, + ) + return _finish_distance_transform( + raw_distances, + raw_indices, + return_distances, + return_indices, + distances, + indices, ) - - # Copy to pre-allocated tensors if provided - if distances is not None and raw_distances is not None: - distances.copy_(raw_distances) - - if indices is not None and raw_indices is not None: - indices.copy_(raw_indices) - - # Return based on scipy convention - return_dist_tensor = return_distances and distances is None - return_idx_tensor = return_indices and indices is None - - if return_dist_tensor and return_idx_tensor: - return raw_distances, raw_indices - elif return_dist_tensor: - return raw_distances - elif return_idx_tensor: - return raw_indices - else: - return None diff --git a/torchmorph/morphology/binary.py b/torchmorph/morphology/binary.py index 98f6133..f383cfa 100644 --- a/torchmorph/morphology/binary.py +++ b/torchmorph/morphology/binary.py @@ -4,8 +4,6 @@ from .. import _C from .structure import _normalize_origin, generate_binary_structure -_BINARY_MORPH_MAX_NDIM = 8 - def _validate_binary_input(input: Tensor) -> int: if not input.is_cuda: @@ -14,20 +12,15 @@ def _validate_binary_input(input: Tensor) -> int: raise ValueError( f"Input must be (B, C, Spatial...) with at least 3 dimensions, got {input.shape}." ) - spatial_ndim = input.ndim - 2 - if spatial_ndim > _BINARY_MORPH_MAX_NDIM: - raise ValueError( - f"Input spatial dimensions must be in range 1 to {_BINARY_MORPH_MAX_NDIM}, " - f"got {spatial_ndim}." - ) + if spatial_ndim > 8: + raise ValueError(f"Input spatial dimensions must be in range 1 to 8, got {spatial_ndim}.") if input.numel() == 0: raise ValueError(f"Invalid input: empty tensor with shape {input.shape}.") - return spatial_ndim -def _validate_output(input: Tensor, output: Tensor | None) -> None: +def _validate_binary_output(input: Tensor, output: Tensor | None) -> None: if output is None: return if output.shape != input.shape: @@ -78,7 +71,7 @@ def _binary_morphology( ) -> Tensor: iterate_until_stable = iterations < 1 spatial_ndim = _validate_binary_input(input) - _validate_output(input, output) + _validate_binary_output(input, output) if structure is None: structure = generate_binary_structure(spatial_ndim, 1) @@ -205,7 +198,7 @@ def binary_fill_holes( ) -> Tensor: """Fill holes in binary objects for `(B, C, Spatial...)` CUDA tensors.""" _validate_binary_input(input) - _validate_output(input, output) + _validate_binary_output(input, output) mask = input == 0 seed = torch.zeros_like(mask, dtype=torch.bool) @@ -235,7 +228,7 @@ def binary_hit_or_miss( ) -> Tensor: """N-dimensional binary hit-or-miss transform for `(B, C, Spatial...)` CUDA tensors.""" spatial_ndim = _validate_binary_input(input) - _validate_output(input, output) + _validate_binary_output(input, output) origin1 = _normalize_origin(origin1, spatial_ndim) origin2 = origin1 if origin2 is None else _normalize_origin(origin2, spatial_ndim) diff --git a/torchmorph/morphology/grey.py b/torchmorph/morphology/grey.py index b3a52e2..fc8b1ae 100644 --- a/torchmorph/morphology/grey.py +++ b/torchmorph/morphology/grey.py @@ -3,6 +3,33 @@ from .. import _C + +def _validate_grey_input(input: Tensor) -> int: + if not input.is_cuda: + raise ValueError("Input tensor must be on CUDA device.") + if input.ndim < 3: + raise ValueError( + f"Input must be (B, C, Spatial...) with at least 3 dimensions, got {input.shape}." + ) + spatial_ndim = input.ndim - 2 + if spatial_ndim > 8: + raise ValueError(f"Input spatial dimensions must be in range 1 to 8, got {spatial_ndim}.") + if input.numel() == 0: + raise ValueError(f"Invalid input: empty tensor with shape {input.shape}.") + return spatial_ndim + + +def _validate_grey_output(input: Tensor, output: Tensor | None) -> None: + if output is None: + return + if output.shape != input.shape: + raise ValueError(f"output shape {output.shape} must match input shape {input.shape}") + if output.device != input.device: + raise ValueError( + f"output must be on the same device as input, got {output.device} and {input.device}" + ) + + _MODE_MAP = { 'constant': 0, 'reflect': 1, @@ -25,16 +52,8 @@ def _grey_morphology( operation: str, ) -> Tensor: """Shared argument normalization for grey erosion and dilation.""" - if not input.is_cuda: - raise ValueError('Input tensor must be on CUDA device.') - if input.ndim < 3: - raise ValueError( - f'Input must be (B, C, Spatial...) with at least 3 dimensions, ' f'got {input.shape}.' - ) - if input.numel() == 0: - raise ValueError(f'Invalid input: empty tensor with shape {input.shape}.') - - spatial_ndim = input.ndim - 2 + spatial_ndim = _validate_grey_input(input) + _validate_grey_output(input, output) footprint_cpu = torch.empty(0, dtype=torch.bool) @@ -243,6 +262,9 @@ def morphological_gradient( origin: int | tuple[int, ...] = 0, ) -> Tensor: """N-dimensional morphological gradient for ``(B, C, Spatial...)`` CUDA tensors.""" + _validate_grey_input(input) + _validate_grey_output(input, output) + dilated = grey_dilation( input, size=size, @@ -279,6 +301,9 @@ def white_tophat( origin: int | tuple[int, ...] = 0, ) -> Tensor: """N-dimensional white top-hat filter for ``(B, C, Spatial...)`` CUDA tensors.""" + _validate_grey_input(input) + _validate_grey_output(input, output) + opened = grey_opening( input, size=size, @@ -306,6 +331,9 @@ def black_tophat( origin: int | tuple[int, ...] = 0, ) -> Tensor: """N-dimensional black top-hat filter for ``(B, C, Spatial...)`` CUDA tensors.""" + _validate_grey_input(input) + _validate_grey_output(input, output) + closed = grey_closing( input, size=size, @@ -333,6 +361,9 @@ def morphological_laplace( origin: int | tuple[int, ...] = 0, ) -> Tensor: """N-dimensional morphological Laplace for ``(B, C, Spatial...)`` CUDA tensors.""" + _validate_grey_input(input) + _validate_grey_output(input, output) + dilated = grey_dilation( input, size=size, From 2d59fe2d3b543ebbafefe343e6007c653a3561d7 Mon Sep 17 00:00:00 2001 From: Yu Han Deng Date: Mon, 13 Jul 2026 15:48:11 +0800 Subject: [PATCH 19/25] delete JFA argoithm --- benchmark/distance_transform.py | 5 +- test/test_distance_transform.py | 182 ++++++- torchmorph/csrc/edt_kernel.cu | 808 ++++--------------------------- torchmorph/csrc/torchmorph.cpp | 6 +- torchmorph/distance_transform.py | 5 +- 5 files changed, 255 insertions(+), 751 deletions(-) diff --git a/benchmark/distance_transform.py b/benchmark/distance_transform.py index a5927ed..3dd09b1 100644 --- a/benchmark/distance_transform.py +++ b/benchmark/distance_transform.py @@ -127,10 +127,7 @@ def benchmark_transform( def bench_edt(args): - operations = { - "exact": partial(tm.euclidean_distance_transform, algorithm="exact"), - "jfa": partial(tm.euclidean_distance_transform, algorithm="jfa"), - } + operations = {"torchmorph": tm.euclidean_distance_transform} benchmark_transform( "EDT 2D", ndi.distance_transform_edt, diff --git a/test/test_distance_transform.py b/test/test_distance_transform.py index 75b461a..905db96 100644 --- a/test/test_distance_transform.py +++ b/test/test_distance_transform.py @@ -76,10 +76,32 @@ def scipy_batch(input_array, scipy_op, **kwargs): return np.stack(results).reshape(input_array.shape) +def scipy_batch_with_indices(input_array, scipy_op, **kwargs): + spatial_shape = input_array.shape[2:] + samples = input_array.reshape(-1, *spatial_shape) + results = [scipy_op(sample, return_indices=True, **kwargs) for sample in samples] + distances = np.stack([result[0] for result in results]).reshape(input_array.shape) + indices = np.stack([result[1] for result in results], axis=1).reshape( + len(spatial_shape), *input_array.shape + ) + return distances, indices + + def cuda(array): return torch.as_tensor(np.ascontiguousarray(array), device="cuda") +def own_indices(input): + spatial_ndim = input.ndim - 2 + coordinates = torch.stack( + torch.meshgrid( + *[torch.arange(size, device=input.device) for size in input.shape[2:]], + indexing="ij", + ) + ).to(torch.int32) + return coordinates.view(spatial_ndim, 1, 1, *input.shape[2:]).expand(spatial_ndim, *input.shape) + + def assert_indices_describe_distances(input, distances, indices, metric, sampling=None): spatial_ndim = input.ndim - 2 assert indices.shape == (spatial_ndim, *input.shape) @@ -259,21 +281,131 @@ def test_sampling_validation(distance_op, _, kwargs, sampling): distance_op(cuda(CASES[2]), sampling=sampling, **kwargs) -@pytest.mark.parametrize("algorithm", ["exact", "jfa", "auto"]) -@pytest.mark.parametrize("spatial_ndim", [2, 3]) -def test_edt_algorithms_match_scipy(algorithm, spatial_ndim): - input_array = CASES[spatial_ndim] - result = tm.euclidean_distance_transform(cuda(input_array), algorithm=algorithm) - expected = scipy_batch(input_array, ndi.distance_transform_edt) - torch.testing.assert_close(result.cpu(), torch.as_tensor(expected).float()) +def test_edt_rejects_removed_algorithm_argument(): + with pytest.raises(TypeError, match="algorithm"): + tm.euclidean_distance_transform(cuda(CASES[2]), algorithm="exact") -def test_edt_jfa_falls_back_for_nonunit_sampling(): - input = cuda(CASES[2]) - kwargs = {"sampling": [0.5, 2.0]} - exact = tm.euclidean_distance_transform(input, algorithm="exact", **kwargs) - jfa = tm.euclidean_distance_transform(input, algorithm="jfa", **kwargs) - torch.testing.assert_close(jfa, exact) +@pytest.mark.parametrize(("torch_op", "scipy_op", "torch_kwargs", "scipy_kwargs"), SCIPY_CASES) +@pytest.mark.parametrize("value", [0.0, 1.0], ids=["all-background", "all-foreground"]) +def test_uniform_inputs_match_scipy(torch_op, scipy_op, torch_kwargs, scipy_kwargs, value): + input_array = np.full((2, 2, 3, 4), value, dtype=np.float32) + distances, indices = torch_op(cuda(input_array), return_indices=True, **torch_kwargs) + expected_distances, expected_indices = scipy_batch_with_indices( + input_array, scipy_op, **scipy_kwargs + ) + + torch.testing.assert_close(distances.cpu(), torch.as_tensor(expected_distances).float()) + torch.testing.assert_close(indices.cpu(), torch.as_tensor(expected_indices).int()) + + +@pytest.mark.parametrize(("torch_op", "scipy_op", "torch_kwargs", "scipy_kwargs"), SCIPY_CASES) +@pytest.mark.parametrize("spatial_ndim", [4, 5, 6, 7, 8]) +def test_high_dimensional_inputs_match_scipy( + torch_op, scipy_op, torch_kwargs, scipy_kwargs, spatial_ndim +): + rng = np.random.default_rng(42) + input_array = rng.integers(0, 2, size=(2, 2, *([2] * spatial_ndim))).astype(np.float32) + distances, indices = torch_op(cuda(input_array), return_indices=True, **torch_kwargs) + expected_distances, expected_indices = scipy_batch_with_indices( + input_array, scipy_op, **scipy_kwargs + ) + torch.testing.assert_close(distances.cpu(), torch.as_tensor(expected_distances).float()) + assert indices.dtype == torch.int32 + assert_indices_describe_distances( + cuda(input_array), + distances, + indices, + torch_kwargs.get("metric", "euclidean"), + ) + + +@pytest.mark.parametrize(("torch_op", "scipy_op", "torch_kwargs", "scipy_kwargs"), SCIPY_CASES) +def test_noncontiguous_batches_and_channels_match_scipy( + torch_op, scipy_op, torch_kwargs, scipy_kwargs +): + input_array = np.ones((2, 2, 5, 6), dtype=np.float32) + input_array[0, 0] = 0 + input_array[0, 1, ::2, ::2] = 0 + input_array[1, 0, 1::2, ::2] = 0 + input_array[1, 1, 2, 3] = 0 + input = cuda(input_array).transpose(-1, -2) + assert not input.is_contiguous() + + transposed_array = input.cpu().numpy() + distances, indices = torch_op(input, return_indices=True, **torch_kwargs) + expected_distances, expected_indices = scipy_batch_with_indices( + transposed_array, scipy_op, **scipy_kwargs + ) + torch.testing.assert_close(distances.cpu(), torch.as_tensor(expected_distances).float()) + assert indices.dtype == torch.int32 + assert_indices_describe_distances( + input, + distances, + indices, + torch_kwargs.get("metric", "euclidean"), + ) + + +@pytest.mark.parametrize("distance_op", DISTANCE_OPERATORS) +def test_distance_transform_is_not_differentiable(distance_op): + input = cuda(CASES[2]).requires_grad_() + result = distance_op(input) + assert not result.requires_grad + + +@pytest.mark.parametrize("distance_op", DISTANCE_OPERATORS) +@pytest.mark.parametrize("spatial_shape", [(16, 16), (2, 2, 2, 2)], ids=["2d", "4d"]) +@pytest.mark.parametrize("output_mode", ["distances", "indices", "preallocated"]) +def test_distance_transform_uses_current_cuda_stream(distance_op, spatial_shape, output_mode): + input = torch.ones((1, 1, *spatial_shape), device="cuda") + stream = torch.cuda.Stream() + distances = None + indices = None + + with torch.cuda.stream(stream): + torch.cuda._sleep(5_000_000) + input.zero_() + if output_mode == "distances": + distances = distance_op(input) + elif output_mode == "indices": + indices = distance_op(input, return_distances=False, return_indices=True) + else: + distances = torch.full_like(input, float("nan")) + indices = torch.full( + (len(spatial_shape), *input.shape), + -2, + dtype=torch.int32, + device=input.device, + ) + result = distance_op( + input, + return_distances=False, + return_indices=False, + distances=distances, + indices=indices, + ) + assert result is None + + stream.synchronize() + if distances is not None: + torch.testing.assert_close(distances, torch.zeros_like(distances)) + if indices is not None: + torch.testing.assert_close(indices, own_indices(input)) + + +@pytest.mark.parametrize("spatial_ndim", [1, 2, 4]) +def test_edt_all_foreground_sampling_matches_scipy(spatial_ndim): + input_array = np.ones((2, 2, *([3] * spatial_ndim)), dtype=np.float32) + sampling = [0.5 + dim for dim in range(spatial_ndim)] + distances, indices = tm.euclidean_distance_transform( + cuda(input_array), sampling=sampling, return_indices=True + ) + expected_distances, expected_indices = scipy_batch_with_indices( + input_array, ndi.distance_transform_edt, sampling=sampling + ) + torch.testing.assert_close(distances.cpu(), torch.as_tensor(expected_distances).float()) + torch.testing.assert_close(indices.cpu(), torch.as_tensor(expected_indices).int()) @pytest.mark.parametrize( @@ -300,7 +432,25 @@ def test_invalid_metric(distance_op, kwargs): @pytest.mark.skipif(torch.cuda.device_count() < 2, reason="requires two CUDA devices") @pytest.mark.parametrize("distance_op", DISTANCE_OPERATORS) -def test_distance_transform_uses_input_device(distance_op): +@pytest.mark.parametrize("output_mode", ["distances", "indices", "preallocated"]) +def test_distance_transform_uses_input_device(distance_op, output_mode): input = torch.as_tensor(CASES[2], device="cuda:1") - result = distance_op(input) - assert result.device == input.device + with torch.cuda.device(0): + if output_mode == "distances": + outputs = (distance_op(input),) + elif output_mode == "indices": + outputs = (distance_op(input, return_distances=False, return_indices=True),) + else: + distances = torch.empty_like(input) + indices = torch.empty((2, *input.shape), dtype=torch.int32, device=input.device) + result = distance_op( + input, + return_distances=False, + return_indices=False, + distances=distances, + indices=indices, + ) + assert result is None + outputs = (distances, indices) + + assert all(output.device == input.device for output in outputs) diff --git a/torchmorph/csrc/edt_kernel.cu b/torchmorph/csrc/edt_kernel.cu index a0a98d7..98317b8 100644 --- a/torchmorph/csrc/edt_kernel.cu +++ b/torchmorph/csrc/edt_kernel.cu @@ -1,4 +1,7 @@ #include +#include +#include +#include #include #include #include @@ -16,463 +19,6 @@ #define MAX_THREADS 256 #define SHARED_MEM_LIMIT 2048 // Max dimension size for shared memory path (48KB limit) -// JFA Configuration -#define BLOCK_SIZE 256 -#define SMEM_LIMIT_ELEMENTS 4096 -#define JFA_BLOCK_DIM 32 -#define JFA_FUSED_STEPS 4 -#define JFA_MAX_OFFSET 8 -#define JFA_SMEM_DIM (JFA_BLOCK_DIM + 2 * JFA_MAX_OFFSET) -#define JFA_3D_BLOCK 8 -#define JFA_3D_HALO 1 - -// ============================================================================== -// JFA Device Helpers -// ============================================================================== -__device__ __forceinline__ float sqr(float x) { return x * x; } - -__device__ __forceinline__ float dist_sq_2d(int y1, int x1, int y2, int x2) { - return sqr((float)(y1 - y2)) + sqr((float)(x1 - x2)); -} - -__device__ __forceinline__ float dist_sq_3d_soa(int z1, int y1, int x1, int z2, int y2, int x2) { - if (z2 == -1) return INF_VAL; - float dz = (float)(z1 - z2); - float dy = (float)(y1 - y2); - float dx = (float)(x1 - x2); - return dz*dz + dy*dy + dx*dx; -} - -__device__ __forceinline__ float compute_cost(int q, int p, float val_p) { - if (p < 0 || val_p >= INF_VAL) return INF_VAL; - return sqr((float)q - (float)p) + val_p; -} - -__device__ __forceinline__ float dist_sq_int2(int y, int x, int2 seed) { - if (seed.x == -1) return INF_VAL; - float dy = (float)(y - seed.x); - float dx = (float)(x - seed.y); - return dy*dy + dx*dx; -} - -// ============================================================================== -// JFA 2D Kernels (Vectorized int2 + Block Shared) -// ============================================================================== -__global__ void init_jfa_2d_opt_kernel( - const float* __restrict__ input, - int2* __restrict__ output, - int64_t total_elements, - int H, int W -) { - int64_t tid = blockIdx.x * blockDim.x + threadIdx.x; - if (tid >= total_elements) return; - - if (input[tid] == 0.0f) { - int64_t spatial_size = (int64_t)H * W; - int64_t rem = tid % spatial_size; - int w = (int)(rem % W); - int h = (int)(rem / W); - output[tid] = make_int2(h, w); - } else { - output[tid] = make_int2(-1, -1); - } -} - -__global__ void jfa_block_fused_2d_kernel( - const int2* __restrict__ in_idx, - int2* __restrict__ out_idx, - int H, int W, - int64_t num_images -) { - __shared__ int2 smem[JFA_SMEM_DIM][JFA_SMEM_DIM]; - - int tx = threadIdx.x; - int ty = threadIdx.y; - - int bx = blockIdx.x * blockDim.x; - int by = blockIdx.y * blockDim.y; - int img_idx = blockIdx.z; - int64_t batch_offset = (int64_t)img_idx * (H * W); - - int gx = bx + tx; - int gy = by + ty; - - // Phase 1: load data to Shared Memory - int smem_linear_size = JFA_SMEM_DIM * JFA_SMEM_DIM; - int total_threads = blockDim.x * blockDim.y; - int thread_linear_idx = ty * blockDim.x + tx; - - int base_x = bx - JFA_MAX_OFFSET; - int base_y = by - JFA_MAX_OFFSET; - - for (int i = thread_linear_idx; i < smem_linear_size; i += total_threads) { - int s_y = i / JFA_SMEM_DIM; - int s_x = i % JFA_SMEM_DIM; - int global_y = base_y + s_y; - int global_x = base_x + s_x; - int2 val = make_int2(-1, -1); - if (global_y >= 0 && global_y < H && global_x >= 0 && global_x < W) { - val = in_idx[batch_offset + global_y * W + global_x]; - } - smem[s_y][s_x] = val; - } - __syncthreads(); - - // Phase 2: Iterate in Shared Memory - if (gx < W && gy < H) { - int center_sy = ty + JFA_MAX_OFFSET; - int center_sx = tx + JFA_MAX_OFFSET; - - int2 best_seed = smem[center_sy][center_sx]; - float best_dist = dist_sq_int2(gy, gx, best_seed); - - int step = 1; - #pragma unroll - for (int k = 0; k < JFA_FUSED_STEPS; ++k) { - #pragma unroll - for (int dy = -1; dy <= 1; ++dy) { - #pragma unroll - for (int dx = -1; dx <= 1; ++dx) { - if (dy == 0 && dx == 0) continue; - int2 neighbor_seed = smem[center_sy + dy * step][center_sx + dx * step]; - if (neighbor_seed.x != -1) { - float d = dist_sq_int2(gy, gx, neighbor_seed); - if (d < best_dist) { - best_dist = d; - best_seed = neighbor_seed; - } - } - } - } - __syncthreads(); - smem[center_sy][center_sx] = best_seed; - __syncthreads(); - step *= 2; - } - out_idx[batch_offset + gy * W + gx] = best_seed; - } -} - -__global__ void jfa_step_global_2d_opt_kernel( - const int2* __restrict__ in_idx, - int2* __restrict__ out_idx, - int step, - int H, int W, - int64_t total_pixels -) { - int64_t tid = blockIdx.x * blockDim.x + threadIdx.x; - if (tid >= total_pixels) return; - - int64_t spatial_size = (int64_t)H * W; - int64_t rem = tid % spatial_size; - int64_t batch_offset = tid - rem; - int w = (int)(rem % W); - int h = (int)(rem / W); - - int2 best_seed = in_idx[tid]; - float best_dist = dist_sq_int2(h, w, best_seed); - - #pragma unroll - for (int dy = -1; dy <= 1; ++dy) { - #pragma unroll - for (int dx = -1; dx <= 1; ++dx) { - if (dx == 0 && dy == 0) continue; - - int ny = h + dy * step; - int nx = w + dx * step; - - if (ny >= 0 && ny < H && nx >= 0 && nx < W) { - int2 neighbor_seed = in_idx[batch_offset + ny * W + nx]; - if (neighbor_seed.x != -1) { - float d = dist_sq_int2(h, w, neighbor_seed); - if (d < best_dist) { - best_dist = d; - best_seed = neighbor_seed; - } - } - } - } - } - out_idx[tid] = best_seed; -} - -__global__ void calc_dist_2d_opt_kernel( - const int2* __restrict__ indices, - float* __restrict__ dist_out, - int64_t total_elements, - int H, int W -) { - int64_t tid = blockIdx.x * blockDim.x + threadIdx.x; - if (tid >= total_elements) return; - - int2 s = indices[tid]; - if (s.x == -1) { - dist_out[tid] = INF_VAL; - } else { - int64_t spatial_size = (int64_t)H * W; - int64_t rem = tid % spatial_size; - int cur_w = (int)(rem % W); - int cur_h = (int)(rem / W); - dist_out[tid] = sqrtf(dist_sq_int2(cur_h, cur_w, s)); - } -} - -// ============================================================================== -// JFA 3D Kernels (Optimized SoA Layout) -// ============================================================================== -template -__global__ void init_jfa_3d_soa_kernel( - const float* __restrict__ input, - IndexType* __restrict__ indices_z, - IndexType* __restrict__ indices_y, - IndexType* __restrict__ indices_x, - int64_t total_elements, - int D, int H, int W -) { - int64_t tid = blockIdx.x * blockDim.x + threadIdx.x; - if (tid >= total_elements) return; - - if (input[tid] == 0.0f) { - int64_t spatial_size = (int64_t)D * H * W; - int64_t rem = tid % spatial_size; - int w = (int)(rem % W); - int h = (int)((rem / W) % H); - int d = (int)(rem / (W * H)); - - indices_z[tid] = (IndexType)d; - indices_y[tid] = (IndexType)h; - indices_x[tid] = (IndexType)w; - } else { - indices_z[tid] = (IndexType)-1; - indices_y[tid] = (IndexType)-1; - indices_x[tid] = (IndexType)-1; - } -} - -template -__global__ void jfa_block_fused_3d_soa_kernel( - const IndexType* __restrict__ in_z, - const IndexType* __restrict__ in_y, - const IndexType* __restrict__ in_x, - IndexType* __restrict__ out_z, - IndexType* __restrict__ out_y, - IndexType* __restrict__ out_x, - int D, int H, int W, - int blocks_per_d -) { - const int BLOCK_DIM = 8; - const int HALO = 3; - const int SMEM_DIM = BLOCK_DIM + 2 * HALO; // 14 - const int SMEM_SIZE = SMEM_DIM * SMEM_DIM * SMEM_DIM; - - extern __shared__ char smem_raw[]; - IndexType* smem_z = (IndexType*)smem_raw; - IndexType* smem_y = smem_z + SMEM_SIZE; - IndexType* smem_x = smem_y + SMEM_SIZE; - - int tx = threadIdx.x; int ty = threadIdx.y; int tz = threadIdx.z; - - int b_z_total = blockIdx.z; - int batch_id = b_z_total / blocks_per_d; - int b_z_local = b_z_total % blocks_per_d; - - int bx = blockIdx.x * BLOCK_DIM; - int by = blockIdx.y * BLOCK_DIM; - int bz = b_z_local * BLOCK_DIM; - - int64_t spatial_offset = (int64_t)batch_id * (D * H * W); - - // Phase 1: Load to SoA Shared Memory - int tid = tz * 64 + ty * 8 + tx; - int base_x = bx - HALO; - int base_y = by - HALO; - int base_z = bz - HALO; - - for (int i = tid; i < SMEM_SIZE; i += 512) { - int temp = i; - int sx = temp % SMEM_DIM; temp /= SMEM_DIM; - int sy = temp % SMEM_DIM; - int sz = temp / SMEM_DIM; - - int gx = base_x + sx; - int gy = base_y + sy; - int gz = base_z + sz; - - IndexType val_z = -1, val_y = -1, val_x = -1; - if (gz >= 0 && gz < D && gy >= 0 && gy < H && gx >= 0 && gx < W) { - int64_t idx = spatial_offset + (int64_t)gz * (H * W) + gy * W + gx; - val_z = in_z[idx]; - val_y = in_y[idx]; - val_x = in_x[idx]; - } - smem_z[i] = val_z; - smem_y[i] = val_y; - smem_x[i] = val_x; - } - __syncthreads(); - - // Phase 2: Compute - int center_sz = tz + HALO; - int center_sy = ty + HALO; - int center_sx = tx + HALO; - int my_s_idx = (center_sz * SMEM_DIM + center_sy) * SMEM_DIM + center_sx; - - int best_z = (int)smem_z[my_s_idx]; - int best_y = (int)smem_y[my_s_idx]; - int best_x = (int)smem_x[my_s_idx]; - - int g_cz = bz + tz; - int g_cy = by + ty; - int g_cx = bx + tx; - - float best_dist = dist_sq_3d_soa(g_cz, g_cy, g_cx, best_z, best_y, best_x); - - int step = 1; - #pragma unroll - for (int k = 0; k < 2; ++k) { - #pragma unroll - for (int dz = -1; dz <= 1; ++dz) { - #pragma unroll - for (int dy = -1; dy <= 1; ++dy) { - #pragma unroll - for (int dx = -1; dx <= 1; ++dx) { - if (dz == 0 && dy == 0 && dx == 0) continue; - - int nz = center_sz + dz * step; - int ny = center_sy + dy * step; - int nx = center_sx + dx * step; - int n_idx = (nz * SMEM_DIM + ny) * SMEM_DIM + nx; - - int sz_in = (int)smem_z[n_idx]; - if (sz_in != -1) { - int sy_in = (int)smem_y[n_idx]; - int sx_in = (int)smem_x[n_idx]; - float d = dist_sq_3d_soa(g_cz, g_cy, g_cx, sz_in, sy_in, sx_in); - if (d < best_dist) { - best_dist = d; - best_z = sz_in; - best_y = sy_in; - best_x = sx_in; - } - } - } - } - } - __syncthreads(); - smem_z[my_s_idx] = (IndexType)best_z; - smem_y[my_s_idx] = (IndexType)best_y; - smem_x[my_s_idx] = (IndexType)best_x; - __syncthreads(); - step *= 2; - } - - if (g_cz < D && g_cy < H && g_cx < W) { - int64_t out_idx_g = spatial_offset + (int64_t)g_cz * (H * W) + g_cy * W + g_cx; - out_z[out_idx_g] = (IndexType)best_z; - out_y[out_idx_g] = (IndexType)best_y; - out_x[out_idx_g] = (IndexType)best_x; - } -} - -template -__global__ void jfa_step_3d_soa_kernel( - const IndexType* __restrict__ in_z, - const IndexType* __restrict__ in_y, - const IndexType* __restrict__ in_x, - IndexType* __restrict__ out_z, - IndexType* __restrict__ out_y, - IndexType* __restrict__ out_x, - int step, - int D, int H, int W, - int64_t total_pixels -) { - int64_t tid = blockIdx.x * blockDim.x + threadIdx.x; - if (tid >= total_pixels) return; - - int64_t spatial_size = (int64_t)D * H * W; - int64_t rem = tid % spatial_size; - int64_t batch_offset = tid - rem; - int cur_w = (int)(rem % W); - int cur_h = (int)((rem / W) % H); - int cur_d = (int)(rem / (W * H)); - - int best_z = (int)in_z[tid]; - int best_y = (int)in_y[tid]; - int best_x = (int)in_x[tid]; - - float best_dist = dist_sq_3d_soa(cur_d, cur_h, cur_w, best_z, best_y, best_x); - - #pragma unroll - for (int dz = -1; dz <= 1; ++dz) { - #pragma unroll - for (int dy = -1; dy <= 1; ++dy) { - #pragma unroll - for (int dx = -1; dx <= 1; ++dx) { - if (dz == 0 && dy == 0 && dx == 0) continue; - - int nz = cur_d + dz * step; - int ny = cur_h + dy * step; - int nx = cur_w + dx * step; - - if (nz >= 0 && nz < D && ny >= 0 && ny < H && nx >= 0 && nx < W) { - int64_t n_idx = batch_offset + (int64_t)nz * (H * W) + ny * W + nx; - - int seed_z = (int)in_z[n_idx]; - if (seed_z != -1) { - float dz_val = (float)(cur_d - seed_z); - float dz_sq = dz_val * dz_val; - - if (dz_sq < best_dist) { - int seed_y = (int)in_y[n_idx]; - int seed_x = (int)in_x[n_idx]; - float dist = dz_sq + sqr((float)(cur_h - seed_y)) + sqr((float)(cur_w - seed_x)); - - if (dist < best_dist) { - best_dist = dist; - best_z = seed_z; - best_y = seed_y; - best_x = seed_x; - } - } - } - } - } - } - } - out_z[tid] = (IndexType)best_z; - out_y[tid] = (IndexType)best_y; - out_x[tid] = (IndexType)best_x; -} - -template -__global__ void calc_dist_3d_soa_kernel( - const IndexType* __restrict__ in_z, - const IndexType* __restrict__ in_y, - const IndexType* __restrict__ in_x, - float* __restrict__ dist_out, - int64_t total_elements, - int D, int H, int W -) { - int64_t tid = blockIdx.x * blockDim.x + threadIdx.x; - if (tid >= total_elements) return; - - int seed_d = (int)in_z[tid]; - if (seed_d == -1) { - dist_out[tid] = INF_VAL; - } else { - int seed_h = (int)in_y[tid]; - int seed_w = (int)in_x[tid]; - - int64_t spatial_size = (int64_t)D * H * W; - int64_t rem = tid % spatial_size; - int cur_w = (int)(rem % W); - int cur_h = (int)((rem / W) % H); - int cur_d = (int)(rem / (W * H)); - - dist_out[tid] = sqrtf(dist_sq_3d_soa(cur_d, cur_h, cur_w, seed_d, seed_h, seed_w)); - } -} - // ============================================================================== // 2D Optimized: Initialization kernel // ============================================================================== @@ -635,6 +181,7 @@ __global__ void edt_2d_cols_kernel( int width, int64_t batch_stride, float spacing, + float spacing_x, bool is_final, bool compute_indices ) { @@ -712,10 +259,12 @@ __global__ void edt_2d_cols_kernel( int64_t out_idx = col_base + q * stride; if (k < 0) { - output[out_idx] = INF_VAL; + float delta_y = (float)(q + 1) * spacing; + float delta_x = (float)col_idx * spacing_x; + output[out_idx] = sqrtf(delta_y * delta_y + delta_x * delta_x); if (compute_indices) { - output_idx_y[out_idx] = 0; - output_idx_x[out_idx] = col_idx; + output_idx_y[out_idx] = -1; + output_idx_x[out_idx] = 0; } } else { float q_pos = (float)q * spacing; @@ -756,6 +305,7 @@ std::tuple run_edt_2d_optimized( TORCH_CHECK(input.scalar_type() == torch::kFloat32, "Input must be float32"); input = input.contiguous(); + const auto stream = c10::cuda::getCurrentCUDAStream(input.get_device()); int total_ndim = input.dim(); TORCH_CHECK(total_ndim >= 2, "Input must have at least 2 dimensions"); @@ -787,7 +337,7 @@ std::tuple run_edt_2d_optimized( dim3 block(16, 16); dim3 grid((width + 15) / 16, (height + 15) / 16, batch_size); - init_distance_2d_kernel<<>>( + init_distance_2d_kernel<<>>( input.data_ptr(), distance.data_ptr(), return_indices ? indices_y.data_ptr() : nullptr, @@ -795,6 +345,7 @@ std::tuple run_edt_2d_optimized( height, width, batch_stride, return_indices ); + C10_CUDA_KERNEL_LAUNCH_CHECK(); } // Step 2: Row-wise EDT (X direction) - shared memory @@ -805,7 +356,7 @@ std::tuple run_edt_2d_optimized( width * sizeof(int) + // v_idx (width + 1) * sizeof(float); // z - edt_2d_rows_kernel<<>>( + edt_2d_rows_kernel<<>>( distance.data_ptr(), temp.data_ptr(), return_indices ? indices_y.data_ptr() : nullptr, @@ -816,6 +367,7 @@ std::tuple run_edt_2d_optimized( spacing_x, return_indices ); + C10_CUDA_KERNEL_LAUNCH_CHECK(); } // Step 3: Column-wise EDT (Y direction) - shared memory @@ -829,7 +381,7 @@ std::tuple run_edt_2d_optimized( shared_mem_size += height * sizeof(int); // src_x } - edt_2d_cols_kernel<<>>( + edt_2d_cols_kernel<<>>( temp.data_ptr(), distance.data_ptr(), return_indices ? temp_idx_y.data_ptr() : nullptr, @@ -838,9 +390,11 @@ std::tuple run_edt_2d_optimized( return_indices ? indices_x.data_ptr() : nullptr, height, width, batch_stride, spacing_y, + spacing_x, true, // is_final return_indices ); + C10_CUDA_KERNEL_LAUNCH_CHECK(); } // Combine indices into single tensor with shape [2, ...] @@ -858,6 +412,30 @@ std::tuple run_edt_2d_optimized( return std::make_tuple(distance, indices); } +__device__ __forceinline__ float virtual_background_distance( + int64_t slice_idx, + int64_t slices_per_sample, + int q, + int spatial_ndim, + const int64_t* __restrict__ spatial_shape, + const float* __restrict__ sampling +) { + int64_t rem = slice_idx % slices_per_sample; + float squared_distance = 0.0f; + + for (int d = spatial_ndim - 2; d >= 0; d--) { + int64_t coordinate = rem % spatial_shape[d]; + rem /= spatial_shape[d]; + float shifted_coordinate = (d == 0) ? (float)(coordinate + 1) : (float)coordinate; + float delta = shifted_coordinate * sampling[d]; + squared_distance += delta * delta; + } + + float last_coordinate = (spatial_ndim == 1) ? (float)(q + 1) : (float)q; + float last_delta = last_coordinate * sampling[spatial_ndim - 1]; + return sqrtf(squared_distance + last_delta * last_delta); +} + // ============================================================================== // 1D EDT kernel using GLOBAL memory (for large dimensions) // ============================================================================== @@ -876,6 +454,9 @@ __global__ void edt_1d_global_kernel( int spatial_ndim, int current_dim, float spacing, + const int64_t* __restrict__ spatial_shape, + const float* __restrict__ sampling, + int64_t slices_per_sample, bool is_final, bool compute_indices ) { @@ -944,10 +525,18 @@ __global__ void edt_1d_global_kernel( int64_t out_idx = base_offset + q; if (k < 0) { - output[out_idx] = INF_VAL; + output[out_idx] = is_final + ? virtual_background_distance( + slice_idx, + slices_per_sample, + q, + spatial_ndim, + spatial_shape, + sampling) + : INF_VAL; if (compute_indices) { for (int d = 0; d < spatial_ndim; d++) { - output_idx[d * num_pixels + out_idx] = 0; + output_idx[d * num_pixels + out_idx] = (is_final && d == 0) ? -1 : 0; } } } else { @@ -996,6 +585,9 @@ __global__ void edt_1d_kernel( int spatial_ndim, int current_dim, float spacing, + const int64_t* __restrict__ spatial_shape, + const float* __restrict__ sampling, + int64_t slices_per_sample, bool is_final, bool compute_indices ) { @@ -1066,10 +658,18 @@ __global__ void edt_1d_kernel( int64_t out_idx = base_offset + q; if (k < 0) { - output[out_idx] = INF_VAL; + output[out_idx] = is_final + ? virtual_background_distance( + slice_idx, + slices_per_sample, + q, + spatial_ndim, + spatial_shape, + sampling) + : INF_VAL; if (compute_indices) { for (int d = 0; d < spatial_ndim; d++) { - output_idx[d * num_pixels + out_idx] = 0; + output_idx[d * num_pixels + out_idx] = (is_final && d == 0) ? -1 : 0; } } } else { @@ -1156,6 +756,7 @@ std::tuple run_edt_separable( TORCH_CHECK(input.scalar_type() == torch::kFloat32, "Input must be float32"); input = input.contiguous(); + const auto stream = c10::cuda::getCurrentCUDAStream(input.get_device()); int total_ndim = input.dim(); int spatial_ndim = sampling.size(); @@ -1176,12 +777,14 @@ std::tuple run_edt_separable( // Copy shape to device auto shape_tensor = torch::tensor(std::vector(shape.begin(), shape.end()), torch::TensorOptions().dtype(torch::kInt64).device(input.device())); + auto sampling_tensor = torch::tensor(sampling, input.options()); + int64_t batch_size = shape[0] * shape[1]; // Initialize distances and indices int threads = 256; int blocks = (total_pixels + threads - 1) / threads; - init_distance_kernel<<>>( + init_distance_kernel<<>>( input.data_ptr(), distance.data_ptr(), return_indices ? indices.data_ptr() : nullptr, @@ -1191,6 +794,7 @@ std::tuple run_edt_separable( shape_tensor.data_ptr(), return_indices ); + C10_CUDA_KERNEL_LAUNCH_CHECK(); // Global memory buffers (allocated lazily for large dimensions) torch::Tensor g_v_val, g_v_idx, g_z, g_k; @@ -1227,7 +831,7 @@ std::tuple run_edt_separable( slice_len * sizeof(int) + // v_idx (slice_len + 1) * sizeof(float); // z - edt_1d_kernel<<>>( + edt_1d_kernel<<>>( dist_transposed.data_ptr(), dist_out.data_ptr(), return_indices ? idx_transposed.data_ptr() : nullptr, @@ -1238,9 +842,13 @@ std::tuple run_edt_separable( spatial_ndim, dim_idx, spacing, + shape_tensor.data_ptr() + start_dim, + sampling_tensor.data_ptr(), + num_slices / batch_size, is_final, return_indices ); + C10_CUDA_KERNEL_LAUNCH_CHECK(); } else { // Allocate global memory buffers if needed int64_t total_elements = dist_transposed.numel(); @@ -1255,7 +863,7 @@ std::tuple run_edt_separable( g_k = torch::empty({num_slices}, dist_transposed.options().dtype(torch::kInt32)); } - edt_1d_global_kernel<<>>( + edt_1d_global_kernel<<>>( dist_transposed.data_ptr(), dist_out.data_ptr(), return_indices ? idx_transposed.data_ptr() : nullptr, @@ -1270,9 +878,13 @@ std::tuple run_edt_separable( spatial_ndim, dim_idx, spacing, + shape_tensor.data_ptr() + start_dim, + sampling_tensor.data_ptr(), + num_slices / batch_size, is_final, return_indices ); + C10_CUDA_KERNEL_LAUNCH_CHECK(); } // Transpose back @@ -1285,211 +897,6 @@ std::tuple run_edt_separable( return std::make_tuple(distance.contiguous(), return_indices ? indices.contiguous() : torch::Tensor()); } -// ============================================================================== -// JFA Dispatch Helpers -// ============================================================================== -std::tuple run_jfa_2d( - torch::Tensor input, int64_t H, int64_t W, int grid, int block, int64_t numel -) { - auto index_opts = input.options().dtype(torch::kInt32); - auto idx_shape = input.sizes().vec(); - idx_shape.push_back(2); - auto curr_idx = torch::empty(idx_shape, index_opts); - auto next_idx = torch::empty(idx_shape, index_opts); - - int2* d_curr = (int2*)curr_idx.data_ptr(); - int2* d_next = (int2*)next_idx.data_ptr(); - - init_jfa_2d_opt_kernel<<>>( - input.data_ptr(), d_curr, numel, H, W - ); - - { - dim3 dimBlock(JFA_BLOCK_DIM, JFA_BLOCK_DIM); - int64_t batch_size = numel / (H * W); - dim3 dimGrid((W + JFA_BLOCK_DIM - 1) / JFA_BLOCK_DIM, - (H + JFA_BLOCK_DIM - 1) / JFA_BLOCK_DIM, - batch_size); - - jfa_block_fused_2d_kernel<<>>(d_curr, d_next, H, W, batch_size); - std::swap(d_curr, d_next); - std::swap(curr_idx, next_idx); - } - - int max_dim = std::max((int)H, (int)W); - int step = 16; - - while (step < max_dim) { - jfa_step_global_2d_opt_kernel<<>>(d_curr, d_next, step, H, W, numel); - std::swap(d_curr, d_next); - std::swap(curr_idx, next_idx); - step *= 2; - } - - auto final_dist = torch::empty_like(input); - calc_dist_2d_opt_kernel<<>>(d_curr, final_dist.data_ptr(), numel, H, W); - - return std::make_tuple(final_dist, curr_idx); -} - -std::tuple run_jfa_3d( - torch::Tensor input, int64_t D, int64_t H, int64_t W, int grid, int block, int64_t numel -) { - bool use_int16 = (D < 32767 && H < 32767 && W < 32767); - auto index_opts = input.options().dtype(use_int16 ? torch::kInt16 : torch::kInt32); - - int64_t batch = numel / (D * H * W); - - // (3, Batch, D, H, W) - auto curr_idx_soa = torch::empty({3, batch, D, H, W}, index_opts); - auto next_idx_soa = torch::empty({3, batch, D, H, W}, index_opts); - - void* d_curr = curr_idx_soa.data_ptr(); - void* d_next = next_idx_soa.data_ptr(); - int64_t plane_stride = numel; // B*D*H*W - - // 1. Init - if (use_int16) { - int16_t* ptr = (int16_t*)d_curr; - init_jfa_3d_soa_kernel<<>>( - input.data_ptr(), ptr, ptr + plane_stride, ptr + 2 * plane_stride, numel, D, H, W - ); - } else { - int32_t* ptr = (int32_t*)d_curr; - init_jfa_3d_soa_kernel<<>>( - input.data_ptr(), ptr, ptr + plane_stride, ptr + 2 * plane_stride, numel, D, H, W - ); - } - - // 2. Fused Steps - int block_dim = 8; - int blocks_per_d = (D + block_dim - 1) / block_dim; - dim3 fused_block(block_dim, block_dim, block_dim); - dim3 fused_grid((W + block_dim - 1) / block_dim, (H + block_dim - 1) / block_dim, blocks_per_d * batch); - size_t smem_bytes = (14*14*14) * 3 * (use_int16 ? 2 : 4); - - if (use_int16) { - int16_t* c = (int16_t*)d_curr; - int16_t* n = (int16_t*)d_next; - jfa_block_fused_3d_soa_kernel<<>>( - c, c + plane_stride, c + 2 * plane_stride, - n, n + plane_stride, n + 2 * plane_stride, - D, H, W, blocks_per_d - ); - } else { - int32_t* c = (int32_t*)d_curr; - int32_t* n = (int32_t*)d_next; - jfa_block_fused_3d_soa_kernel<<>>( - c, c + plane_stride, c + 2 * plane_stride, - n, n + plane_stride, n + 2 * plane_stride, - D, H, W, blocks_per_d - ); - } - std::swap(d_curr, d_next); - - // 3. Global Steps - int max_dim = std::max({(int)D, (int)H, (int)W}); - int step = 4; - while (step < max_dim) { - if (use_int16) { - int16_t* c = (int16_t*)d_curr; - int16_t* n = (int16_t*)d_next; - jfa_step_3d_soa_kernel<<>>( - c, c + plane_stride, c + 2 * plane_stride, - n, n + plane_stride, n + 2 * plane_stride, - step, D, H, W, numel - ); - } else { - int32_t* c = (int32_t*)d_curr; - int32_t* n = (int32_t*)d_next; - jfa_step_3d_soa_kernel<<>>( - c, c + plane_stride, c + 2 * plane_stride, - n, n + plane_stride, n + 2 * plane_stride, - step, D, H, W, numel - ); - } - std::swap(d_curr, d_next); - step *= 2; - } - - // 4. Final Dist - auto final_dist = torch::empty_like(input); - if (use_int16) { - int16_t* c = (int16_t*)d_curr; - calc_dist_3d_soa_kernel<<>>( - c, c + plane_stride, c + 2 * plane_stride, - final_dist.data_ptr(), numel, D, H, W - ); - } else { - int32_t* c = (int32_t*)d_curr; - calc_dist_3d_soa_kernel<<>>( - c, c + plane_stride, c + 2 * plane_stride, - final_dist.data_ptr(), numel, D, H, W - ); - } - - // Permute result indices back to (Batch, D, H, W, 3) - torch::Tensor result_indices; - if (d_curr == curr_idx_soa.data_ptr()) result_indices = curr_idx_soa; - else result_indices = next_idx_soa; - - result_indices = result_indices.permute({1, 2, 3, 4, 0}).contiguous(); - - return std::make_tuple(final_dist, result_indices); -} - -// ============================================================================== -// JFA Main Entry Point -// ============================================================================== -std::tuple distance_transform_cuda(torch::Tensor input) { - TORCH_CHECK(input.is_cuda(), "Input must be CUDA tensor"); - input = input.contiguous(); - - int64_t dims = input.dim(); - int64_t numel = input.numel(); - int block = BLOCK_SIZE; - int grid = (numel + block - 1) / block; - - if (dims >= 5) { - // For 4D+ spatial, fall back to separable algorithm - int spatial_ndim = dims - 1; - std::vector sampling(spatial_ndim, 1.0f); - return run_edt_separable(input, sampling, true); - } - else if (dims == 4) { - int64_t dim1 = input.size(1); - if (dim1 == 1) { - int64_t H = input.size(-2); - int64_t W = input.size(-1); - return run_jfa_2d(input, H, W, grid, block, numel); - } - else { - int64_t D = dim1; - int64_t H = input.size(-2); - int64_t W = input.size(-1); - return run_jfa_3d(input, D, H, W, grid, block, numel); - } - } - else if (dims == 3) { - int64_t H = input.size(-2); - int64_t W = input.size(-1); - return run_jfa_2d(input, H, W, grid, block, numel); - } - else if (dims == 2) { - int64_t H = 1; - int64_t W = input.size(-1); - auto result = run_jfa_2d(input, H, W, grid, block, numel); - torch::Tensor dist = std::get<0>(result); - torch::Tensor idx_2d = std::get<1>(result); - auto idx_1d = idx_2d.slice(/*dim=*/-1, /*start=*/1, /*end=*/2).contiguous(); - return std::make_tuple(dist, idx_1d); - } - else { - TORCH_CHECK(false, "Unsupported dimensions."); - return std::make_tuple(torch::Tensor(), torch::Tensor()); - } -} - // ============================================================================== // Python binding entry point // ============================================================================== @@ -1498,10 +905,10 @@ std::tuple edt_cuda( torch::Tensor input, std::vector sampling, bool return_distances, - bool return_indices, - const std::string& algorithm + bool return_indices ) { TORCH_CHECK(input.is_cuda(), "Input must be a CUDA tensor"); + const c10::cuda::CUDAGuard device_guard(input.device()); int total_ndim = input.dim(); @@ -1515,53 +922,6 @@ std::tuple edt_cuda( int spatial_ndim = sampling.size(); - // Check if we can use JFA algorithm - bool can_use_jfa = true; - - // JFA doesn't support non-unit sampling - for (float s : sampling) { - if (std::abs(s - 1.0f) > 1e-6f) { - can_use_jfa = false; - break; - } - } - - // JFA only supports 2D and 3D (spatial dimensions) - if (spatial_ndim > 3) { - can_use_jfa = false; - } - - // Determine which algorithm to use - bool use_jfa = false; - if (algorithm == "jfa") { - if (can_use_jfa) { - use_jfa = true; - } else { - // Fall back to exact with warning (or we can throw) - // For now, silently fall back to exact - use_jfa = false; - } - } else if (algorithm == "exact") { - use_jfa = false; - } else if (algorithm == "auto") { - // Auto mode: use JFA only for 2D with unit sampling - // For 3D, exact algorithm performs better in practice - use_jfa = can_use_jfa && (spatial_ndim == 2); - } else { - TORCH_CHECK(false, "algorithm must be 'exact', 'jfa', or 'auto', got: ", algorithm); - } - - if (use_jfa) { - // Use JFA algorithm - auto [distances, indices_result] = distance_transform_cuda(input); - - if (!return_indices) { - indices_result = torch::Tensor(); - } - - return std::make_tuple(distances, indices_result); - } - // Use exact (Felzenszwalb) algorithm // Use 2D optimized path only when both dimensions fit in shared memory // For larger dimensions, the N-D general version with transpose is faster diff --git a/torchmorph/csrc/torchmorph.cpp b/torchmorph/csrc/torchmorph.cpp index 6321765..4d3e0c8 100644 --- a/torchmorph/csrc/torchmorph.cpp +++ b/torchmorph/csrc/torchmorph.cpp @@ -36,8 +36,7 @@ std::tuple edt_cuda( torch::Tensor input, std::vector sampling, bool return_distances, - bool return_indices, - const std::string& algorithm + bool return_indices ); std::tuple cdt_cuda( @@ -94,8 +93,7 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { py::arg("input"), py::arg("sampling"), py::arg("return_distances") = true, - py::arg("return_indices") = false, - py::arg("algorithm") = "exact"); + py::arg("return_indices") = false); m.def("cdt_cuda", &cdt_cuda, "Chessboard/Manhattan distance transform", diff --git a/torchmorph/distance_transform.py b/torchmorph/distance_transform.py index 08152d6..bc7e6d9 100644 --- a/torchmorph/distance_transform.py +++ b/torchmorph/distance_transform.py @@ -100,7 +100,6 @@ def euclidean_distance_transform( return_indices: bool = False, distances: Tensor | None = None, indices: Tensor | None = None, - algorithm: str = "auto", ) -> Tensor | tuple[Tensor, Tensor] | None: """Euclidean distance transform for (B, C, Spatial...) CUDA tensors.""" spatial_ndim, return_distances, return_indices = _prepare_distance_transform( @@ -110,12 +109,12 @@ def euclidean_distance_transform( distances, indices, ) + normalized_sampling = _normalize_sampling(sampling, spatial_ndim) raw_distances, raw_indices = _C.edt_cuda( input.float().contiguous(), - _normalize_sampling(sampling, spatial_ndim), + normalized_sampling, return_distances, return_indices, - algorithm, ) return _finish_distance_transform( raw_distances, From 9b4fba592fc93d6e6cdd6241ad7d44f46c3c1ec8 Mon Sep 17 00:00:00 2001 From: Yu Han Deng Date: Mon, 13 Jul 2026 15:48:53 +0800 Subject: [PATCH 20/25] add CUDAGuard --- torchmorph/csrc/bfdt_kernel.cu | 93 ++++++++++++++++++++++------------ torchmorph/csrc/cdt_kernel.cu | 61 ++++++++++++++-------- 2 files changed, 99 insertions(+), 55 deletions(-) diff --git a/torchmorph/csrc/bfdt_kernel.cu b/torchmorph/csrc/bfdt_kernel.cu index 408c6c3..0ec30ec 100644 --- a/torchmorph/csrc/bfdt_kernel.cu +++ b/torchmorph/csrc/bfdt_kernel.cu @@ -1,10 +1,15 @@ #include +#include +#include +#include #include #include #include #include #include #include +#include +#include #define BFDT_INF_VAL 1e20f #define BFDT_MAX_NDIM 8 @@ -130,14 +135,14 @@ __global__ void bfdt_kernel( #define DISPATCH_NDIM_AND_METRIC(METRIC_TYPE, NDIM_VAL, ...) \ switch (NDIM_VAL) { \ - case 1: bfdt_kernel<<>>(__VA_ARGS__); break; \ - case 2: bfdt_kernel<<>>(__VA_ARGS__); break; \ - case 3: bfdt_kernel<<>>(__VA_ARGS__); break; \ - case 4: bfdt_kernel<<>>(__VA_ARGS__); break; \ - case 5: bfdt_kernel<<>>(__VA_ARGS__); break; \ - case 6: bfdt_kernel<<>>(__VA_ARGS__); break; \ - case 7: bfdt_kernel<<>>(__VA_ARGS__); break; \ - case 8: bfdt_kernel<<>>(__VA_ARGS__); break; \ + case 1: bfdt_kernel<<>>(__VA_ARGS__); break; \ + case 2: bfdt_kernel<<>>(__VA_ARGS__); break; \ + case 3: bfdt_kernel<<>>(__VA_ARGS__); break; \ + case 4: bfdt_kernel<<>>(__VA_ARGS__); break; \ + case 5: bfdt_kernel<<>>(__VA_ARGS__); break; \ + case 6: bfdt_kernel<<>>(__VA_ARGS__); break; \ + case 7: bfdt_kernel<<>>(__VA_ARGS__); break; \ + case 8: bfdt_kernel<<>>(__VA_ARGS__); break; \ default: TORCH_CHECK(false, "Unsupported number of dimensions: ", NDIM_VAL); \ } @@ -154,8 +159,10 @@ std::tuple bfdt_cuda( bool return_indices ) { TORCH_CHECK(input.is_cuda(), "Input must be a CUDA tensor"); + const c10::cuda::CUDAGuard device_guard(input.device()); input = input.contiguous(); + const auto stream = c10::cuda::getCurrentCUDAStream(input.get_device()); auto shape = input.sizes(); int ndim = input.dim() - 2; // Extract spatial dimensions (excluding Batch and Channel) int64_t batch_size = shape[0] * shape[1]; @@ -192,9 +199,26 @@ std::tuple bfdt_cuda( int num_fg = fg_indices_flat.size(0); int num_bg = bg_indices_flat.size(0); - - if (num_fg == 0) continue; - if (num_bg == 0) continue; // No background: distances remain initialized to INF + + if (num_fg == 0 && !return_indices) { + dist_out.view({batch_size, spatial_size})[b].zero_(); + continue; + } + + if (num_bg == 0) { + if (return_distances) { + float value = metric_type == 0 + ? std::numeric_limits::infinity() + : static_cast(std::numeric_limits::max()); + dist_out.view({batch_size, spatial_size})[b].fill_(value); + } + if (return_indices) { + for (int d = 0; d < ndim; d++) { + indices_out.select(0, d).view({batch_size, spatial_size})[b].zero_(); + } + } + continue; + } auto spatial_shape = input.sizes().slice(2); auto fg_coords = torch::empty({num_fg, ndim}, input.options().dtype(torch::kInt32)); @@ -223,29 +247,32 @@ std::tuple bfdt_cuda( int blocks = (num_fg + threads - 1) / threads; size_t shared_mem_bytes = BFDT_BLOCK_SIZE * ndim * sizeof(float); - // Dispatch to different template instantiations based on metric type + // Dispatch to different template instantiations based on metric type // to guarantee zero-overhead branching inside the inner loops. - if (metric_type == 0) { - DISPATCH_NDIM_AND_METRIC(0, ndim, - fg_coords.data_ptr(), bg_coords.data_ptr(), - batch_dist.data_ptr(), batch_indices.data_ptr(), - num_fg, num_bg, sampling_tensor.data_ptr(), - return_distances, return_indices - ); - } else if (metric_type == 1) { - DISPATCH_NDIM_AND_METRIC(1, ndim, - fg_coords.data_ptr(), bg_coords.data_ptr(), - batch_dist.data_ptr(), batch_indices.data_ptr(), - num_fg, num_bg, sampling_tensor.data_ptr(), - return_distances, return_indices - ); - } else if (metric_type == 2) { - DISPATCH_NDIM_AND_METRIC(2, ndim, - fg_coords.data_ptr(), bg_coords.data_ptr(), - batch_dist.data_ptr(), batch_indices.data_ptr(), - num_fg, num_bg, sampling_tensor.data_ptr(), - return_distances, return_indices - ); + if (num_fg > 0) { + if (metric_type == 0) { + DISPATCH_NDIM_AND_METRIC(0, ndim, + fg_coords.data_ptr(), bg_coords.data_ptr(), + batch_dist.data_ptr(), batch_indices.data_ptr(), + num_fg, num_bg, sampling_tensor.data_ptr(), + return_distances, return_indices + ); + } else if (metric_type == 1) { + DISPATCH_NDIM_AND_METRIC(1, ndim, + fg_coords.data_ptr(), bg_coords.data_ptr(), + batch_dist.data_ptr(), batch_indices.data_ptr(), + num_fg, num_bg, sampling_tensor.data_ptr(), + return_distances, return_indices + ); + } else if (metric_type == 2) { + DISPATCH_NDIM_AND_METRIC(2, ndim, + fg_coords.data_ptr(), bg_coords.data_ptr(), + batch_dist.data_ptr(), batch_indices.data_ptr(), + num_fg, num_bg, sampling_tensor.data_ptr(), + return_distances, return_indices + ); + } + C10_CUDA_KERNEL_LAUNCH_CHECK(); } // Scatter the computed results back to their spatial positions diff --git a/torchmorph/csrc/cdt_kernel.cu b/torchmorph/csrc/cdt_kernel.cu index 355c26a..511d65f 100644 --- a/torchmorph/csrc/cdt_kernel.cu +++ b/torchmorph/csrc/cdt_kernel.cu @@ -1,4 +1,7 @@ #include +#include +#include +#include #include #include #include @@ -28,27 +31,29 @@ __global__ void cdt_init_kernel( int64_t tid = blockIdx.x * blockDim.x + threadIdx.x; if (tid >= total_elements) return; - if (input[tid] == 0.0f) { - dist[tid] = 0; - if (compute_indices) { - int64_t spatial_idx = tid % spatial_elements; - int64_t rem = spatial_idx; - for (int d = 0; d < spatial_ndim; d++) { - int64_t coord = rem / spatial_strides[d]; - rem = rem % spatial_strides[d]; - indices[d * total_elements + tid] = (int32_t)coord; - } - } - } else { - dist[tid] = CDT_INF_VAL; - if (compute_indices) { - for (int d = 0; d < spatial_ndim; d++) { - indices[d * total_elements + tid] = -1; - } + dist[tid] = (input[tid] == 0.0f) ? 0 : CDT_INF_VAL; + + if (compute_indices) { + int64_t spatial_idx = tid % spatial_elements; + int64_t rem = spatial_idx; + for (int d = 0; d < spatial_ndim; d++) { + int64_t coord = rem / spatial_strides[d]; + rem %= spatial_strides[d]; + indices[d * total_elements + tid] = (int32_t)coord; } } } +__global__ void cdt_finalize_distance_kernel( + const int32_t* __restrict__ dist, + float* __restrict__ output, + int64_t total_elements +) { + int64_t tid = blockIdx.x * blockDim.x + threadIdx.x; + if (tid >= total_elements) return; + output[tid] = (dist[tid] == CDT_INF_VAL) ? -1.0f : (float)dist[tid]; +} + // ============================================================================ // Dimension-wise sweep kernels for chessboard metric // Each thread handles one "line" along the scan dimension @@ -268,6 +273,7 @@ std::tuple cdt_cuda( bool return_indices ) { TORCH_CHECK(input.is_cuda(), "Input must be CUDA tensor"); + const c10::cuda::CUDAGuard device_guard(input.device()); TORCH_CHECK(input.scalar_type() == torch::kFloat32, "Input must be float32"); TORCH_CHECK(metric == "chessboard" || metric == "taxicab", "metric must be 'chessboard' or 'taxicab'"); @@ -275,6 +281,7 @@ std::tuple cdt_cuda( "At least one of return_distances or return_indices must be True"); input = input.contiguous(); + const auto stream = c10::cuda::getCurrentCUDAStream(input.get_device()); bool is_taxicab = (metric == "taxicab"); int total_ndim = input.dim(); @@ -319,7 +326,7 @@ std::tuple cdt_cuda( int block = CDT_BLOCK_SIZE; int grid = (total_elements + block - 1) / block; - cdt_init_kernel<<>>( + cdt_init_kernel<<>>( input.data_ptr(), dist.data_ptr(), return_indices ? indices.data_ptr() : nullptr, @@ -329,6 +336,7 @@ std::tuple cdt_cuda( spatial_strides_tensor.data_ptr(), return_indices ); + C10_CUDA_KERNEL_LAUNCH_CHECK(); // For each dimension, do forward and backward sweeps for (int d = 0; d < spatial_ndim; d++) { @@ -340,7 +348,7 @@ std::tuple cdt_cuda( int sweep_grid = (num_lines + sweep_block - 1) / sweep_block; // Forward sweep - cdt_sweep_forward_chessboard_kernel<<>>( + cdt_sweep_forward_chessboard_kernel<<>>( dist.data_ptr(), return_indices ? indices.data_ptr() : nullptr, total_elements, @@ -355,9 +363,10 @@ std::tuple cdt_cuda( spatial_shape_tensor.data_ptr(), return_indices ); + C10_CUDA_KERNEL_LAUNCH_CHECK(); // Backward sweep - cdt_sweep_backward_chessboard_kernel<<>>( + cdt_sweep_backward_chessboard_kernel<<>>( dist.data_ptr(), return_indices ? indices.data_ptr() : nullptr, total_elements, @@ -372,6 +381,7 @@ std::tuple cdt_cuda( spatial_shape_tensor.data_ptr(), return_indices ); + C10_CUDA_KERNEL_LAUNCH_CHECK(); } // For chessboard metric, we need additional diagonal passes @@ -408,7 +418,7 @@ std::tuple cdt_cuda( // Need more passes for higher dimensions to ensure full propagation int num_passes = spatial_ndim * 2; // Scale with dimensions for (int pass = 0; pass < num_passes; pass++) { - cdt_diagonal_pass_kernel<<>>( + cdt_diagonal_pass_kernel<<>>( dist.data_ptr(), return_indices ? indices.data_ptr() : nullptr, total_elements, @@ -422,6 +432,7 @@ std::tuple cdt_cuda( return_indices, pass % 2 == 0 ); + C10_CUDA_KERNEL_LAUNCH_CHECK(); } } } @@ -431,7 +442,13 @@ std::tuple cdt_cuda( torch::Tensor result_indices; if (return_distances) { - result_dist = dist.to(torch::kFloat32).view(input.sizes()); + result_dist = torch::empty_like(input); + cdt_finalize_distance_kernel<<>>( + dist.data_ptr(), + result_dist.data_ptr(), + total_elements + ); + C10_CUDA_KERNEL_LAUNCH_CHECK(); } if (return_indices) { From 4bc15233a4a8bbdd44902edfc3d47aada9ed5d4c Mon Sep 17 00:00:00 2001 From: Yu Han Deng Date: Mon, 13 Jul 2026 15:49:15 +0800 Subject: [PATCH 21/25] add validation test --- test/test_binary.py | 7 +++++++ test/test_grey.py | 15 +++++++++++++++ torchmorph/morphology/grey.py | 8 +++++--- 3 files changed, 27 insertions(+), 3 deletions(-) diff --git a/test/test_binary.py b/test/test_binary.py index e6102a4..f9e681c 100644 --- a/test/test_binary.py +++ b/test/test_binary.py @@ -439,6 +439,13 @@ def test_binary_morphology_uses_current_cuda_stream(torch_op): assert result.device == x.device +@pytest.mark.parametrize("torch_op", ALL_TORCH_OPERATORS) +def test_binary_morphology_is_not_differentiable(torch_op): + x = torch.as_tensor(CASE_2D, dtype=torch.float32, device="cuda").requires_grad_() + result = torch_op(x) + assert not result.requires_grad + + @pytest.mark.skipif(torch.cuda.device_count() < 2, reason="requires two CUDA devices") @pytest.mark.parametrize("torch_op", ALL_TORCH_OPERATORS) def test_binary_morphology_uses_input_cuda_device(torch_op): diff --git a/test/test_grey.py b/test/test_grey.py index 686e577..2862613 100644 --- a/test/test_grey.py +++ b/test/test_grey.py @@ -207,6 +207,21 @@ def test_grey_morphology_rejects_invalid_origin(torch_op, origin): torch_op(x, size=3, origin=origin) +@pytest.mark.parametrize("torch_op", TORCH_OPERATORS) +@pytest.mark.parametrize("size", [0, -1, (3, 0), (3, -1)]) +def test_grey_morphology_rejects_nonpositive_size(torch_op, size): + x = torch.as_tensor(CASE_2D, device="cuda") + with pytest.raises(ValueError, match="greater than zero"): + torch_op(x, size=size) + + +@pytest.mark.parametrize("torch_op", TORCH_OPERATORS) +def test_grey_morphology_is_not_differentiable(torch_op): + x = torch.as_tensor(CASE_2D, device="cuda").requires_grad_() + result = torch_op(x, size=3) + assert not result.requires_grad + + @pytest.mark.parametrize(("torch_op", "scipy_op"), GREY_OPERATORS) def test_grey_morphology_uses_current_cuda_stream(torch_op, scipy_op): stream = torch.cuda.Stream() diff --git a/torchmorph/morphology/grey.py b/torchmorph/morphology/grey.py index fc8b1ae..958671e 100644 --- a/torchmorph/morphology/grey.py +++ b/torchmorph/morphology/grey.py @@ -79,6 +79,8 @@ def _grey_morphology( elif size is not None: if isinstance(size, int): size = (size,) * spatial_ndim + if any(value <= 0 for value in size): + raise ValueError("size values must be greater than zero") struct = torch.zeros(size, dtype=torch.float32) else: raise ValueError('At least one of size, footprint, or structure must be specified.') @@ -313,7 +315,7 @@ def white_tophat( cval=cval, origin=origin, ) - result = input.float() - opened + result = input.detach().float() - opened if output is not None: output.copy_(result) return output @@ -343,7 +345,7 @@ def black_tophat( cval=cval, origin=origin, ) - result = closed - input.float() + result = closed - input.detach().float() if output is not None: output.copy_(result) return output @@ -382,7 +384,7 @@ def morphological_laplace( cval=cval, origin=origin, ) - result = dilated + eroded - 2 * input.float() + result = dilated + eroded - 2 * input.detach().float() if output is not None: output.copy_(result) return output From 627303d2312428201c8de9acb9be6dcf5247ea7d Mon Sep 17 00:00:00 2001 From: Kai Zhao <694691@qq.com> Date: Mon, 13 Jul 2026 16:40:27 +0800 Subject: [PATCH 22/25] optimal transport (#6) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * optimal transport * debug ot function and add test * add basic benchmark * decouple the iteration from the derivation * Implement Sinkhorn iterations in CUDA * Add Sinkhorn CUDA benchmark and correctness tests * Add CUDA support for Sinkhorn log-domain iterations and related tests. * Refactor Sinkhorn functions and test gradient accuracy * more test * try test * test on cuda kernal * tighten numeric gradient test and support more dimensions * cancel unnecessary change in test.yml Removed TORCH_CUDA_ARCH_LIST environment variable from the workflow. * Rework Sinkhorn OT: standard epsilon convention, log-space end-to-end, batched (n, d) API - Use the standard convention K = exp(-C / epsilon); POT baselines need no parameter inversion anymore. - Keep log potentials in log space end-to-end: the log-domain CUDA solver returns (log_u, log_v), gradients are computed as f = eps * log_u, and the transport plan is reconstructed as exp(log_u - C/eps + log_v) for all backends so it cannot underflow at small epsilon. - Every backend takes strictly (n, d) batched histograms sharing one (d, d) cost matrix; the CUDA kernels are batched over a 2-D launch grid, removing the previous B=C=1 restriction. - Merge the four kernels into two (scaling_update, log_scaling_update), fixing the uncoalesced K^T reads; collapse entry points into solve(); delete the dead sinkhorn__nobatch and loss_test. - CUDA hygiene: current-stream launches, CUDAGuard, TORCH_CHECK input validation, kernel-launch checks, and a NaN guard for all -inf LSE rows. - Export SinkhornSolver / build_cost_matrix from the package; tests import it directly, run the torch backend on CPU, and add a small-epsilon stability test against POT's sinkhorn_log. Co-Authored-By: Claude Fable 5 * Apply black formatting to optimal transport files Co-Authored-By: Claude Fable 5 * Make SinkhornSolver a differentiable nn.Module with automatic device dispatch - SinkhornSolver now inherits torch.nn.Module: forward(source, target, cost_matrix) returns the (n,) transport distances, and a custom autograd.Function implements backward via the centered dual potentials (envelope theorem), chained through the differentiable normalization. - Replace the confused torch/cuda/cuda_log backend strings with one orthogonal option: log_space=True/False selects the algorithm, while the implementation is picked from the input automatically (fused CUDA kernels for float32 CUDA tensors, pure torch ops elsewhere - including a logsumexp-based log-space path that also runs on CPU and float64). - Convenience methods plan() and potentials() expose the transport plan and dual potentials; extra_repr() reports the module configuration. - Tests cover both log_space modes across forward shapes, marginals, POT baselines, fused-vs-CPU consistency, small-epsilon stability, and autograd-vs-numeric gradients; benchmark updated to the module API. Co-Authored-By: Claude Fable 5 * Expand Sinkhorn tests to a full device x log_space x dtype matrix Previously most tests picked a single device (CUDA when available), so the CPU paths went untested on GPU machines, POT baselines only covered the plain iteration, and backward never ran through the fused kernels. Now: - Shared DEVICES parametrization runs every applicable test on both CPU and CUDA (skipping cleanly when CUDA is absent). - POT distance/plan baselines and the small-epsilon stability test cover both log_space modes and both devices (torch logsumexp path on CPU, fused log kernel on CUDA). - New tests: dtype/device consistency (float64 CUDA exercises the torch fallback on GPU), backward through the fused float32 kernels vs CPU, fused-vs-CPU potentials, and threshold early stopping reaching the same solution on both algorithms. Co-Authored-By: Claude Fable 5 * Optimize Sinkhorn CUDA path: batch tiling, online logsumexp, CUDA graphs Kernel changes (both memory-bandwidth-bound, so all wins target matrix traffic and launch latency): - Batch tiling: one block now streams a matrix row once and applies it to up to 8 batch items held in registers, cutting d^2 matrix reads by up to 8x for batched inputs. A TILE=1 template instantiation keeps the n=1 fast path free of predication overhead. - Online logsumexp: the log-domain update now computes the LSE in a single pass over the row (running max with rescaled running sum, NaN-free -inf merging), halving matrix traffic vs the two-pass max-then-sum version. - Warp-shuffle reductions replace the shared-memory tree (one barrier instead of eight per reduction). Python-side: - Long fused runs (max_iter >= 100) capture a 25-iteration chunk into a torch.cuda.CUDAGraph and replay it, hiding the 2-launches-per-iteration overhead that dominates small problems; falls back to plain launches when capture is unavailable (extra fixed-point iterations are harmless). - Skip the (d, d) transpose copy when the cost/kernel matrix is symmetric, which build_cost_matrix outputs always are. Tests: fused-vs-CPU now covers single-item, partial-tile, and multi-tile batches (n = 1, 3, 20) plus a short-run case below the CUDA-graph threshold to exercise the plain-launch fused path. Co-Authored-By: Claude Fable 5 * Test gradients w.r.t. both source and target The backward pass already returned gradients for both marginals (the centered dual potentials f and g), but the tests only verified the source side. The numeric-derivative test now perturbs source and target independently, and the fused-kernel backward test compares both gradients against CPU. Co-Authored-By: Claude Fable 5 --------- Co-authored-by: maokovski <1498743918@qq.com> Co-authored-by: maokovski <“1498743918@qq.com”> Co-authored-by: Claude Fable 5 --- benchmark/optimal_transport.py | 134 +++++++++++++ test/test_optimal_transport.py | 327 ++++++++++++++++++++++++++++++++ torchmorph/__init__.py | 3 + torchmorph/csrc/ot_kernel.cu | 283 +++++++++++++++++++++++++++ torchmorph/csrc/torchmorph.cpp | 44 +++++ torchmorph/optimal_transport.py | 249 ++++++++++++++++++++++++ 6 files changed, 1040 insertions(+) create mode 100644 benchmark/optimal_transport.py create mode 100644 test/test_optimal_transport.py create mode 100644 torchmorph/csrc/ot_kernel.cu create mode 100644 torchmorph/optimal_transport.py diff --git a/benchmark/optimal_transport.py b/benchmark/optimal_transport.py new file mode 100644 index 0000000..b800933 --- /dev/null +++ b/benchmark/optimal_transport.py @@ -0,0 +1,134 @@ +import ot +import torch +from torch.utils import benchmark + +from torchmorph import SinkhornSolver, build_cost_matrix + + +def _grid_problem(n, H, W, epsilon, max_iter, log_space=False, seed=42): + """Random (n, H*W) marginals with a 2-D grid cost matrix.""" + device = "cuda" if torch.cuda.is_available() else "cpu" + torch.manual_seed(seed) + source = torch.rand(n, H * W, device=device) + target = torch.rand(n, H * W, device=device) + cost_matrix = build_cost_matrix((H, W), device=device) + solver = SinkhornSolver(epsilon=epsilon, max_iter=max_iter, log_space=log_space) + return solver, source, target, cost_matrix + + +def run_forward_benchmark(n=1, H=32, W=32, epsilon=1.0, max_iter=100, log_space=False): + solver, source, target, cost_matrix = _grid_problem(n, H, W, epsilon, max_iter, log_space) + if torch.cuda.is_available(): + torch.cuda.reset_peak_memory_stats() + timer = benchmark.Timer( + stmt="solver(source, target, cost_matrix)", + globals={ + "solver": solver, + "source": source, + "target": target, + "cost_matrix": cost_matrix, + }, + ) + result = timer.blocked_autorange(min_run_time=3) + if torch.cuda.is_available(): + print(f"peak allocated: {torch.cuda.max_memory_allocated() / 1024**2:.2f} MB") + print(f"peak reserved: {torch.cuda.max_memory_reserved() / 1024**2:.2f} MB") + return result + + +def run_pot_sinkhorn_benchmark(H=32, W=32, epsilon=1.0, max_iter=100): + solver, source, target, cost_matrix = _grid_problem(1, H, W, epsilon, max_iter) + source, target, cost_matrix = solver.data_preprocess(source, target, cost_matrix) + timer = benchmark.Timer( + stmt="ot.sinkhorn(source, target, M, reg=epsilon, numItermax=max_iter, stopThr=1e-5)", + globals={ + "ot": ot, + "source": source[0], + "target": target[0], + "M": cost_matrix, + "epsilon": epsilon, + "max_iter": max_iter, + }, + ) + return timer.blocked_autorange(min_run_time=2) + + +@torch.no_grad() +def run_sinkhorn_relative_error(H=32, W=32, epsilon=1.0, max_iter=100): + solver, source, target, cost_matrix = _grid_problem(1, H, W, epsilon, max_iter) + plan = solver.plan(source, target, cost_matrix) + a, b, cost_matrix = solver.data_preprocess(source, target, cost_matrix) + ot_plan = ot.sinkhorn(a[0], b[0], cost_matrix, reg=epsilon, numItermax=max_iter, stopThr=1e-5) + return torch.linalg.norm(plan[0] - ot_plan) / torch.clamp(torch.linalg.norm(ot_plan), min=1e-12) + + +def reference_error_check(size=32, epsilon=10.0, max_iter=200, verbose=False): + """Compare log-space fused-kernel potentials against POT's sinkhorn_log.""" + if not torch.cuda.is_available(): + raise RuntimeError("CUDA device not available.") + solver, source, target, cost_matrix = _grid_problem( + 1, size, size, epsilon, max_iter, log_space=True, seed=2 + ) + + grad_f, grad_g = solver.potentials(source, target, cost_matrix) + grad_f, grad_g = grad_f[0], grad_g[0] + + a, b, cost_matrix = solver.data_preprocess(source, target, cost_matrix) + _, log = ot.bregman.sinkhorn_log( + a[0], + b[0], + cost_matrix, + reg=epsilon, + numItermax=max_iter, + stopThr=0, + log=True, + ) + f = epsilon * log["log_u"] + g = epsilon * log["log_v"] + f = f - f.mean() + g = g - g.mean() + + if verbose: + print(f, g, grad_f, grad_g, sep="\n") + + f_rel_l2 = torch.linalg.norm(f - grad_f) / torch.clamp(torch.linalg.norm(f), min=1e-12) + g_rel_l2 = torch.linalg.norm(g - grad_g) / torch.clamp(torch.linalg.norm(g), min=1e-12) + return { + "size": size, + "points": size * size, + "epsilon": epsilon, + "max_iter": max_iter, + "f_grad_max_abs": (f - grad_f).abs().max().item(), + "f_grad_rel_l2": f_rel_l2.item(), + "g_grad_max_abs": (g - grad_g).abs().max().item(), + "g_grad_rel_l2": g_rel_l2.item(), + "f_allclose": torch.allclose(f, grad_f, rtol=1e-4, atol=1e-4), + "g_allclose": torch.allclose(g, grad_g, rtol=1e-4, atol=1e-4), + } + + +def print_reference_error_table( + sizes=(2, 4, 8, 16, 32, 64), max_iters=(100, 200, 500), epsilon=10.0 +): + """Print a size and iteration sweep for POT-vs-CUDA gradient error.""" + headers = ("size", "N", "eps", "itr", "f_max", "f_rel_l2", "g_max", "g_rel_l2", "close") + print("| " + " | ".join(headers) + " |") + print("| " + " | ".join(["---"] * len(headers)) + " |") + for size in sizes: + for max_iter in max_iters: + row = reference_error_check(size=size, epsilon=epsilon, max_iter=max_iter) + print( + f"| {row['size']}x{row['size']} " + f"| {row['points']} " + f"| {row['epsilon']:.3g} " + f"| {row['max_iter']} " + f"| {row['f_grad_max_abs']:.3e} " + f"| {row['f_grad_rel_l2']:.3e} " + f"| {row['g_grad_max_abs']:.3e} " + f"| {row['g_grad_rel_l2']:.3e} " + f"| {row['f_allclose']}/{row['g_allclose']} |" + ) + + +if __name__ == "__main__": + print_reference_error_table() diff --git a/test/test_optimal_transport.py b/test/test_optimal_transport.py new file mode 100644 index 0000000..785e069 --- /dev/null +++ b/test/test_optimal_transport.py @@ -0,0 +1,327 @@ +import numpy as np +import pytest +import torch + +from torchmorph import SinkhornSolver, build_cost_matrix + +try: + import ot +except ImportError: # pragma: no cover - exercised only when POT is absent + ot = None + +# Every test parametrized over DEVICES runs on CPU and, when available, CUDA, +# so both the pure-torch and the fused-kernel implementations are covered. +DEVICES = [ + "cpu", + pytest.param( + "cuda", + marks=pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available"), + ), +] +LOG_SPACE = [False, True] + + +def _require_cuda(): + if not torch.cuda.is_available(): + pytest.skip("CUDA is required for the fused Sinkhorn kernels") + + +def _require_pot(): + if ot is None: + pytest.skip("POT is required for baseline Sinkhorn tests") + + +def _make_positive(shape, device, seed): + generator = torch.Generator(device=device).manual_seed(seed) + return torch.rand(shape, generator=generator, device=device) + 0.1 + + +def _pot_plan(source, target, cost_matrix, solver): + return ot.sinkhorn( + source.detach().cpu().numpy().astype(np.float64), + target.detach().cpu().numpy().astype(np.float64), + cost_matrix.detach().cpu().numpy().astype(np.float64), + reg=solver.epsilon, + numItermax=solver.max_iter, + stopThr=1e-9, + ) + + +@pytest.mark.parametrize("device", DEVICES) +def test_build_cost_matrix(device): + # 1-D grid: zero diagonal, symmetry, unit spacing + cost_1d = build_cost_matrix((5,), device=device) + assert cost_1d.shape == (5, 5) + assert cost_1d.device.type == device + assert torch.allclose(cost_1d.diag(), torch.zeros(5, device=device)) + assert torch.allclose(cost_1d, cost_1d.T) + assert cost_1d[0, 1].item() == pytest.approx(1.0) + + # 2-D grid: adjacent columns are distance 1 apart + cost_2d = build_cost_matrix((3, 4), device=device) + assert cost_2d.shape == (12, 12) + assert torch.allclose(cost_2d, cost_2d.T) + assert cost_2d[0, 1].item() == pytest.approx(1.0) + + # p-norms differ on diagonal moves: grid points (0, 0) and (1, 1) + assert build_cost_matrix((2, 2), p=1, device=device)[0, 3].item() == pytest.approx(2.0) + assert build_cost_matrix((2, 2), p=2, device=device)[0, 3].item() == pytest.approx(2.0**0.5) + + +@pytest.mark.parametrize("device", DEVICES) +@pytest.mark.parametrize("log_space", LOG_SPACE) +@pytest.mark.parametrize("n,d", [(1, 4), (3, 6), (2, 16)]) +def test_forward_shapes_and_marginals(n, d, log_space, device): + solver = SinkhornSolver(epsilon=1.5, max_iter=300, log_space=log_space) + + source = _make_positive((n, d), device=device, seed=11) + target = _make_positive((n, d), device=device, seed=23) + + distance = solver(source, target) + plan = solver.plan(source, target) + f, g = solver.potentials(source, target) + assert distance.shape == (n,) + assert distance.device.type == device + assert plan.shape == (n, d, d) + assert f.shape == (n, d) and g.shape == (n, d) + + a, b, cost_matrix = solver.data_preprocess(source, target) + assert torch.allclose(distance, (plan * cost_matrix).sum(dim=(-2, -1)), rtol=1e-5, atol=1e-6) + assert torch.allclose(plan.sum(dim=-1), a, atol=1e-4) + assert torch.allclose(plan.sum(dim=-2), b, atol=1e-4) + + +@pytest.mark.parametrize("shape", [(4,), (2, 2, 2)]) +def test_forward_rejects_non_2d_inputs(shape): + solver = SinkhornSolver(epsilon=1.5, max_iter=10) + source = _make_positive(shape, device="cpu", seed=911) + target = _make_positive(shape, device="cpu", seed=919) + + with pytest.raises(ValueError, match=r"\(n, d\)"): + solver(source, target) + + +def test_forward_rejects_mismatched_cost_matrix(): + solver = SinkhornSolver(epsilon=1.5, max_iter=10) + source = _make_positive((2, 4), device="cpu", seed=5) + target = _make_positive((2, 4), device="cpu", seed=7) + + with pytest.raises(ValueError, match="cost_matrix"): + solver(source, target, cost_matrix=torch.zeros(3, 3)) + + +@pytest.mark.parametrize("device", DEVICES) +@pytest.mark.parametrize("log_space", LOG_SPACE) +def test_distance_matches_pot_sinkhorn(log_space, device): + _require_pot() + solver = SinkhornSolver(epsilon=1.5, max_iter=1000, log_space=log_space) + + source = _make_positive((3, 8), device=device, seed=101) + target = _make_positive((3, 8), device=device, seed=211) + distance = solver(source, target) + + a, b, cost_matrix = solver.data_preprocess(source, target) + cost_np = cost_matrix.cpu().numpy() + expected = torch.tensor( + [ + float(np.sum(_pot_plan(a[i], b[i], cost_matrix, solver) * cost_np)) + for i in range(a.shape[0]) + ], + device=device, + dtype=distance.dtype, + ) + assert torch.allclose(distance, expected, rtol=5e-3, atol=1e-5) + + +@pytest.mark.parametrize("device", DEVICES) +@pytest.mark.parametrize("log_space", LOG_SPACE) +def test_plan_matches_pot_sinkhorn_grid_cost(log_space, device): + """A flattened 2-D image with an explicit grid cost matrix matches POT.""" + _require_pot() + solver = SinkhornSolver(epsilon=1.5, max_iter=1000, log_space=log_space) + cost_matrix = build_cost_matrix((2, 3), device=device) + + source = _make_positive((1, 6), device=device, seed=307) + target = _make_positive((1, 6), device=device, seed=401) + plan = solver.plan(source, target, cost_matrix) + + a, b, _ = solver.data_preprocess(source, target, cost_matrix) + expected = torch.tensor( + _pot_plan(a[0], b[0], cost_matrix, solver), device=device, dtype=plan.dtype + ) + assert torch.allclose(plan[0], expected, rtol=5e-3, atol=1e-5) + + +@pytest.mark.parametrize("log_space", LOG_SPACE) +@pytest.mark.parametrize("n", [1, 3, 20]) # single-item, partial-tile, and multi-tile batches +def test_fused_cuda_kernels_match_cpu(log_space, n): + """CUDA float32 inputs dispatch to the fused kernels; results must match CPU.""" + _require_cuda() + solver = SinkhornSolver(epsilon=1.5, max_iter=1000, log_space=log_space) + + source = _make_positive((n, 6), device="cpu", seed=503) + target = _make_positive((n, 6), device="cpu", seed=607) + + cpu_distance = solver(source, target) + cpu_plan = solver.plan(source, target) + cpu_f, cpu_g = solver.potentials(source, target) + cuda_distance = solver(source.cuda(), target.cuda()) + cuda_plan = solver.plan(source.cuda(), target.cuda()) + cuda_f, cuda_g = solver.potentials(source.cuda(), target.cuda()) + + assert torch.allclose(cuda_distance.cpu(), cpu_distance, rtol=1e-4, atol=1e-5) + assert torch.allclose(cuda_plan.cpu(), cpu_plan, rtol=1e-4, atol=1e-5) + assert torch.allclose(cuda_f.cpu(), cpu_f, rtol=1e-4, atol=1e-4) + assert torch.allclose(cuda_g.cpu(), cpu_g, rtol=1e-4, atol=1e-4) + + +@pytest.mark.parametrize("log_space", LOG_SPACE) +def test_fused_short_run_without_cuda_graph_matches_cpu(log_space): + """max_iter below the CUDA-graph threshold takes the plain-launch fused path.""" + _require_cuda() + solver = SinkhornSolver(epsilon=1.5, max_iter=60, log_space=log_space) + + source = _make_positive((2, 6), device="cpu", seed=1409) + target = _make_positive((2, 6), device="cpu", seed=1423) + + cpu_distance = solver(source, target) + cuda_distance = solver(source.cuda(), target.cuda()) + assert torch.allclose(cuda_distance.cpu(), cpu_distance, rtol=1e-4, atol=1e-5) + + +@pytest.mark.parametrize("log_space", LOG_SPACE) +def test_dtype_device_combinations_agree(log_space): + """float32/float64 on CPU/CUDA must all reach the same solution. + + float64 CUDA inputs exercise the torch fallback on the GPU, float32 CUDA + inputs the fused kernels, and CPU inputs the pure-torch paths. + """ + solver = SinkhornSolver(epsilon=1.5, max_iter=2000, log_space=log_space) + source = _make_positive((2, 8), device="cpu", seed=41) + target = _make_positive((2, 8), device="cpu", seed=43) + reference = solver(source.double(), target.double()) + + combos = [(source, target)] # cpu float32 + if torch.cuda.is_available(): + combos += [ + (source.cuda(), target.cuda()), # cuda float32: fused kernels + (source.cuda().double(), target.cuda().double()), # cuda float64: torch ops + ] + for src, tgt in combos: + distance = solver(src, tgt) + tol = 1e-4 if distance.dtype == torch.float32 else 1e-8 + assert torch.allclose(distance.cpu().double(), reference, rtol=tol, atol=tol) + + +@pytest.mark.parametrize("device", DEVICES) +def test_log_space_survives_small_epsilon(device): + """At this epsilon exp(-C/eps) underflows float32, so only log_space=True is usable. + + On CPU this exercises the torch logsumexp path, on CUDA the fused log kernel. + """ + _require_pot() + solver = SinkhornSolver(epsilon=0.1, max_iter=5000, log_space=True) + cost_matrix = build_cost_matrix((8, 8), device=device) + + source = _make_positive((1, 64), device=device, seed=31) + target = _make_positive((1, 64), device=device, seed=37) + + f, g = solver.potentials(source, target, cost_matrix) + assert torch.isfinite(f).all() + assert torch.isfinite(g).all() + + plan = solver.plan(source, target, cost_matrix)[0] + a, b, _ = solver.data_preprocess(source, target, cost_matrix) + assert torch.isfinite(plan).all() + assert (plan >= 0).all() + assert torch.allclose(plan.sum(dim=0), b[0], atol=1e-5) + assert torch.allclose(plan.sum(dim=1), a[0], atol=5e-3) + + cost_np = cost_matrix.cpu().numpy().astype(np.float64) + pot_plan = ot.bregman.sinkhorn_log( + a[0].cpu().numpy().astype(np.float64), + b[0].cpu().numpy().astype(np.float64), + cost_np, + reg=solver.epsilon, + numItermax=solver.max_iter, + stopThr=1e-12, + ) + pot_distance = float(np.sum(pot_plan * cost_np)) + distance = solver(source, target, cost_matrix) + assert distance[0].item() == pytest.approx(pot_distance, rel=2e-2) + + +@pytest.mark.parametrize("device", DEVICES) +@pytest.mark.parametrize("log_space", LOG_SPACE) +def test_backward_matches_numeric_directional_derivative(log_space, device): + """Autograd gradients w.r.t. BOTH source and target must match numeric + derivatives of the entropic cost. + + Runs in float64, i.e. through the torch paths on both CPU and CUDA. + """ + n, d = 3, 8 + solver = SinkhornSolver(epsilon=1.5, max_iter=3000, log_space=log_space) + cost_matrix = build_cost_matrix((d,), device=device).double() + + source = _make_positive((n, d), device=device, seed=1009).double().requires_grad_(True) + target = _make_positive((n, d), device=device, seed=2003).double().requires_grad_(True) + direction = _make_positive((n, d), device=device, seed=3001).double() - 0.6 + + solver(source, target, cost_matrix).sum().backward() + assert source.grad is not None and target.grad is not None + + def entropic_cost(src, tgt): + plan = solver.plan(src, tgt, cost_matrix) + entropy = (plan * plan.clamp(min=1e-300).log()).sum() + return (plan * cost_matrix).sum() + solver.epsilon * entropy + + eps_fd = 1e-4 + src, tgt = source.detach(), target.detach() + for grad, plus, minus in [ + (source.grad, (src + eps_fd * direction, tgt), (src - eps_fd * direction, tgt)), + (target.grad, (src, tgt + eps_fd * direction), (src, tgt - eps_fd * direction)), + ]: + analytic = (grad * direction).sum() + numeric = (entropic_cost(*plus) - entropic_cost(*minus)) / (2 * eps_fd) + assert torch.allclose(analytic, numeric, rtol=1e-5, atol=1e-7) + + +@pytest.mark.parametrize("log_space", LOG_SPACE) +def test_backward_through_fused_kernels_matches_cpu(log_space): + """float32 CUDA backward runs through the fused kernels; grads must match CPU.""" + _require_cuda() + solver = SinkhornSolver(epsilon=1.5, max_iter=1000, log_space=log_space) + + source = _make_positive((3, 6), device="cpu", seed=1213) + target = _make_positive((3, 6), device="cpu", seed=1217) + + cpu_source = source.clone().requires_grad_(True) + cpu_target = target.clone().requires_grad_(True) + solver(cpu_source, cpu_target).sum().backward() + cuda_source = source.cuda().requires_grad_(True) + cuda_target = target.cuda().requires_grad_(True) + solver(cuda_source, cuda_target).sum().backward() + + for cuda_grad, cpu_grad in [ + (cuda_source.grad, cpu_source.grad), + (cuda_target.grad, cpu_target.grad), + ]: + assert torch.isfinite(cuda_grad).all() + assert torch.allclose(cuda_grad.cpu(), cpu_grad, rtol=1e-4, atol=1e-5) + + +@pytest.mark.parametrize("device", DEVICES) +@pytest.mark.parametrize("log_space", LOG_SPACE) +def test_threshold_early_stopping_reaches_same_solution(log_space, device): + """threshold > 0 stops the torch paths early without changing the solution.""" + source = _make_positive((2, 8), device=device, seed=1301) + target = _make_positive((2, 8), device=device, seed=1303) + source, target = source.double(), target.double() # force the torch paths + + full = SinkhornSolver(epsilon=1.5, max_iter=5000, log_space=log_space) + early = SinkhornSolver(epsilon=1.5, max_iter=5000, threshold=1e-12, log_space=log_space) + + assert torch.allclose(early(source, target), full(source, target), rtol=1e-8, atol=1e-10) + assert torch.allclose( + early.plan(source, target), full.plan(source, target), rtol=1e-6, atol=1e-10 + ) diff --git a/torchmorph/__init__.py b/torchmorph/__init__.py index 096562b..de23f1c 100644 --- a/torchmorph/__init__.py +++ b/torchmorph/__init__.py @@ -19,6 +19,7 @@ morphological_laplace, white_tophat, ) +from .optimal_transport import SinkhornSolver, build_cost_matrix __all__ = [ "euclidean_distance_transform", @@ -38,4 +39,6 @@ "morphological_laplace", "white_tophat", "black_tophat", + "SinkhornSolver", + "build_cost_matrix", ] diff --git a/torchmorph/csrc/ot_kernel.cu b/torchmorph/csrc/ot_kernel.cu new file mode 100644 index 0000000..193d1ac --- /dev/null +++ b/torchmorph/csrc/ot_kernel.cu @@ -0,0 +1,283 @@ +#include +#include +#include +#include + +namespace { + +constexpr int kThreads = 256; // multiple of the warp size, assumed below +constexpr int kBatchTile = 8; // batch items processed per block +constexpr int kWarps = kThreads / 32; + +__device__ inline float warp_reduce_sum(float v){ + for (int offset = 16; offset > 0; offset >>= 1){ + v += __shfl_down_sync(0xffffffff, v, offset); + } + return v; +} + +// Block-level sum reduction; the result is valid in thread 0 only. +__device__ inline float block_reduce_sum(float v){ + __shared__ float partial[kWarps]; + const int lane = threadIdx.x & 31; + const int warp = threadIdx.x >> 5; + v = warp_reduce_sum(v); + if (lane == 0) partial[warp] = v; + __syncthreads(); + if (warp == 0){ + v = lane < kWarps ? partial[lane] : 0.0f; + v = warp_reduce_sum(v); + } + __syncthreads(); // let the next call reuse `partial` + return v; +} + +// Merge two online-logsumexp states; (m, s) represents m + log(s). +// s is 0 exactly when m is -inf (the empty state), which the guards below +// keep NaN-free. +__device__ inline void lse_merge(float& m, float& s, float m_other, float s_other){ + const float hi = fmaxf(m, m_other); + if (hi == -INFINITY){ + m = -INFINITY; + s = 0.0f; + return; + } + s = (m == -INFINITY ? 0.0f : s * expf(m - hi)) + + (m_other == -INFINITY ? 0.0f : s_other * expf(m_other - hi)); + m = hi; +} + +__device__ inline void warp_reduce_lse(float& m, float& s){ + for (int offset = 16; offset > 0; offset >>= 1){ + const float m_other = __shfl_down_sync(0xffffffff, m, offset); + const float s_other = __shfl_down_sync(0xffffffff, s, offset); + lse_merge(m, s, m_other, s_other); + } +} + +// Block-level logsumexp reduction; the result is valid in thread 0 only. +__device__ inline void block_reduce_lse(float& m, float& s){ + __shared__ float partial_m[kWarps]; + __shared__ float partial_s[kWarps]; + const int lane = threadIdx.x & 31; + const int warp = threadIdx.x >> 5; + warp_reduce_lse(m, s); + if (lane == 0){ + partial_m[warp] = m; + partial_s[warp] = s; + } + __syncthreads(); + if (warp == 0){ + m = lane < kWarps ? partial_m[lane] : -INFINITY; + s = lane < kWarps ? partial_s[lane] : 0.0f; + warp_reduce_lse(m, s); + } + __syncthreads(); // let the next call reuse the partial arrays +} + +// One Sinkhorn scaling update. One block per (row, batch tile): the matrix +// row is streamed once and applied to TILE batch items, so the d^2 matrix is +// read ceil(n / TILE) times per update instead of n times. +// out[b, i] = num[b, i] / (sum_j mat[i, j] * scale[b, j] + 1e-12) +// The u-update passes (a, K, v); the v-update passes (b, K^T, u). +template +__global__ void scaling_update( + const float* __restrict__ num, + const float* __restrict__ mat, + const float* __restrict__ scale, + float* __restrict__ out, + int d, + int n_batch +){ + const int row = blockIdx.x; + const int batch0 = blockIdx.y * TILE; + const int tile = min(TILE, n_batch - batch0); + const float* mat_row = mat + (size_t)row * d; + + float acc[TILE]; +#pragma unroll + for (int t = 0; t < TILE; t++) acc[t] = 0.0f; + + for (int j = threadIdx.x; j < d; j += blockDim.x){ + const float k = mat_row[j]; +#pragma unroll + for (int t = 0; t < TILE; t++){ + if (t < tile) acc[t] += k * scale[(size_t)(batch0 + t) * d + j]; + } + } + + for (int t = 0; t < tile; t++){ + const float total = block_reduce_sum(acc[t]); + if (threadIdx.x == 0){ + const size_t idx = (size_t)(batch0 + t) * d + row; + out[idx] = num[idx] / (total + 1e-12f); + } + } +} + +// One log-domain Sinkhorn update with the same batch tiling: +// log_out[b, i] = log_num[b, i] - logsumexp_j(-cost[i, j] / eps + log_scale[b, j]) +// The logsumexp is evaluated online in a single pass over the row (running +// max with a rescaled running sum), halving the matrix traffic of the +// classic two-pass max-then-sum evaluation. +// The u-update passes (log_a, M, log_v); the v-update passes (log_b, M^T, log_u). +template +__global__ void log_scaling_update( + const float* __restrict__ log_num, + const float* __restrict__ cost, + const float* __restrict__ log_scale, + float* __restrict__ log_out, + int d, + int n_batch, + float inv_eps +){ + const int row = blockIdx.x; + const int batch0 = blockIdx.y * TILE; + const int tile = min(TILE, n_batch - batch0); + const float* cost_row = cost + (size_t)row * d; + + float m[TILE], s[TILE]; +#pragma unroll + for (int t = 0; t < TILE; t++){ + m[t] = -INFINITY; + s[t] = 0.0f; + } + + for (int j = threadIdx.x; j < d; j += blockDim.x){ + const float c = -cost_row[j] * inv_eps; +#pragma unroll + for (int t = 0; t < TILE; t++){ + if (t >= tile) continue; + const float z = c + log_scale[(size_t)(batch0 + t) * d + j]; + if (z > m[t]){ + s[t] = s[t] * expf(m[t] - z) + 1.0f; // expf(-inf) == 0 covers the first element + m[t] = z; + } else if (m[t] != -INFINITY){ // skip z == m == -inf, whose contribution is 0 + s[t] += expf(z - m[t]); + } + } + } + + for (int t = 0; t < tile; t++){ + block_reduce_lse(m[t], s[t]); + if (threadIdx.x == 0){ + const size_t idx = (size_t)(batch0 + t) * d + row; + // m == -inf means the opposite marginal is all zero; pin the + // potential instead of producing NaN. + log_out[idx] = m[t] == -INFINITY ? -INFINITY : log_num[idx] - m[t] - logf(s[t]); + } + } +} + +void check_input(const torch::Tensor& t, int64_t numel, const char* name){ + TORCH_CHECK(t.is_cuda(), name, " must be a CUDA tensor"); + TORCH_CHECK(t.scalar_type() == torch::kFloat32, name, " must be float32"); + TORCH_CHECK(t.is_contiguous(), name, " must be contiguous"); + TORCH_CHECK(t.numel() == numel, name, " must have ", numel, " elements, got ", t.numel()); +} + +// Validate batched (n, d) vectors against a shared (d, d) matrix. +void check_shapes(const torch::Tensor& u, const torch::Tensor& mat){ + TORCH_CHECK(mat.dim() == 2 && mat.size(0) == mat.size(1), "cost/kernel matrix must be square"); + TORCH_CHECK(u.dim() == 2 && u.size(1) == mat.size(0), + "scaling vectors must be (n, d) with d matching the matrix"); +} + +} // namespace + +std::tuple sinkhorn_fastiter( + const torch::Tensor& a, + const torch::Tensor& b, + const torch::Tensor& K, + const torch::Tensor& K_T, + torch::Tensor u, + torch::Tensor v, + int64_t n_iter +){ + check_shapes(u, K); + const int64_t n = u.size(0); + const int64_t d = u.size(1); + check_input(a, n * d, "a"); + check_input(b, n * d, "b"); + check_input(K, d * d, "K"); + check_input(K_T, d * d, "K_T"); + check_input(u, n * d, "u"); + check_input(v, n * d, "v"); + + const at::cuda::CUDAGuard guard(u.device()); + const auto stream = at::cuda::getCurrentCUDAStream(); + const bool single = n == 1; + const dim3 grid(static_cast(d), + single ? 1u : static_cast((n + kBatchTile - 1) / kBatchTile)); + + for (int64_t i = 0; i < n_iter; i++){ + if (single){ + scaling_update<1><<>>( + a.data_ptr(), K.data_ptr(), v.data_ptr(), + u.data_ptr(), (int)d, (int)n); + scaling_update<1><<>>( + b.data_ptr(), K_T.data_ptr(), u.data_ptr(), + v.data_ptr(), (int)d, (int)n); + } else { + scaling_update<<>>( + a.data_ptr(), K.data_ptr(), v.data_ptr(), + u.data_ptr(), (int)d, (int)n); + scaling_update<<>>( + b.data_ptr(), K_T.data_ptr(), u.data_ptr(), + v.data_ptr(), (int)d, (int)n); + } + } + C10_CUDA_KERNEL_LAUNCH_CHECK(); + + return std::make_tuple(u, v); +} + +std::tuple sinkhorn_logiter( + const torch::Tensor& log_a, + const torch::Tensor& log_b, + const torch::Tensor& M, + const torch::Tensor& M_T, + torch::Tensor log_u, + torch::Tensor log_v, + int64_t n_iter, + double epsilon +){ + TORCH_CHECK(epsilon > 0, "epsilon must be positive, got ", epsilon); + check_shapes(log_u, M); + const int64_t n = log_u.size(0); + const int64_t d = log_u.size(1); + check_input(log_a, n * d, "log_a"); + check_input(log_b, n * d, "log_b"); + check_input(M, d * d, "M"); + check_input(M_T, d * d, "M_T"); + check_input(log_u, n * d, "log_u"); + check_input(log_v, n * d, "log_v"); + + const at::cuda::CUDAGuard guard(log_u.device()); + const auto stream = at::cuda::getCurrentCUDAStream(); + const float inv_eps = static_cast(1.0 / epsilon); + const bool single = n == 1; + const dim3 grid(static_cast(d), + single ? 1u : static_cast((n + kBatchTile - 1) / kBatchTile)); + + for (int64_t i = 0; i < n_iter; i++){ + if (single){ + log_scaling_update<1><<>>( + log_a.data_ptr(), M.data_ptr(), log_v.data_ptr(), + log_u.data_ptr(), (int)d, (int)n, inv_eps); + log_scaling_update<1><<>>( + log_b.data_ptr(), M_T.data_ptr(), log_u.data_ptr(), + log_v.data_ptr(), (int)d, (int)n, inv_eps); + } else { + log_scaling_update<<>>( + log_a.data_ptr(), M.data_ptr(), log_v.data_ptr(), + log_u.data_ptr(), (int)d, (int)n, inv_eps); + log_scaling_update<<>>( + log_b.data_ptr(), M_T.data_ptr(), log_u.data_ptr(), + log_v.data_ptr(), (int)d, (int)n, inv_eps); + } + } + C10_CUDA_KERNEL_LAUNCH_CHECK(); + + return std::make_tuple(log_u, log_v); +} diff --git a/torchmorph/csrc/torchmorph.cpp b/torchmorph/csrc/torchmorph.cpp index e553463..669c56d 100644 --- a/torchmorph/csrc/torchmorph.cpp +++ b/torchmorph/csrc/torchmorph.cpp @@ -41,6 +41,27 @@ std::tuple bfdt_cuda( bool return_indices ); +// Optimal Transport functions +std::tuple sinkhorn_fastiter( + const torch::Tensor& a, + const torch::Tensor& b, + const torch::Tensor& K, + const torch::Tensor& K_T, + torch::Tensor u, + torch::Tensor v, + int64_t n_iter +); + +std::tuple sinkhorn_logiter( + const torch::Tensor& log_a, + const torch::Tensor& log_b, + const torch::Tensor& M, + const torch::Tensor& M_T, + torch::Tensor log_u, + torch::Tensor log_v, + int64_t n_iter, + double epsilon +); PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { m.def("grey_erosion_cuda", &grey_erosion_cuda, @@ -82,4 +103,27 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { py::arg("sampling") = std::vector(), py::arg("return_distances") = true, py::arg("return_indices") = false); + + // Optimal Transport + m.def("sinkhorn_fastiter", &sinkhorn_fastiter, + "Sinkhorn scaling-form iterations (CUDA)", + py::arg("a"), + py::arg("b"), + py::arg("K"), + py::arg("K_T"), + py::arg("u"), + py::arg("v"), + py::arg("n_iter")); + + m.def("sinkhorn_logiter", &sinkhorn_logiter, + "Sinkhorn log-domain iterations (CUDA)", + py::arg("log_a"), + py::arg("log_b"), + py::arg("M"), + py::arg("M_T"), + py::arg("log_u"), + py::arg("log_v"), + py::arg("n_iter"), + py::arg("epsilon")); } + diff --git a/torchmorph/optimal_transport.py b/torchmorph/optimal_transport.py new file mode 100644 index 0000000..3652845 --- /dev/null +++ b/torchmorph/optimal_transport.py @@ -0,0 +1,249 @@ +from typing import Optional + +import torch +from torch import Tensor, nn + + +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`. + """ + coords = torch.stack( + torch.meshgrid([torch.arange(s, device=device) for s in shape], indexing="ij"), dim=-1 + ) + coords = coords.reshape(-1, len(shape)).float() + return torch.cdist(coords, coords, p=p) + + +def _use_fused_kernels(t: Tensor) -> bool: + return t.is_cuda and t.dtype == torch.float32 + + +def _transpose_or_self(matrix: Tensor) -> Tensor: + """Return the contiguous transpose, reusing the input when it is symmetric. + + Cost matrices from build_cost_matrix are always symmetric, so this + usually saves a (d, d) copy. + """ + return matrix if torch.equal(matrix, matrix.mT) else matrix.mT.contiguous() + + +_GRAPH_MIN_ITER = 100 # below this, plain launches beat the graph-capture setup cost +_GRAPH_CHUNK = 25 + + +def _run_fused(launch, max_iter, device): + """Run `launch(k)`, which enqueues k in-place Sinkhorn iterations, max_iter times in total. + + Long runs are launch-latency-bound (two tiny kernels per iteration), so a + chunk of iterations is captured into a CUDA graph once and replayed. The + iterations are fixed-point steps on ping-pong buffers, so a fallback that + runs extra iterations is always safe. + """ + if max_iter < _GRAPH_MIN_ITER or torch.cuda.is_current_stream_capturing(): + launch(max_iter) + return + try: + with torch.cuda.device(device): + # Warm up on a side stream (required before capture), doing the + # first chunk of real iterations in the process. + side_stream = torch.cuda.Stream() + side_stream.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(side_stream): + launch(_GRAPH_CHUNK) + torch.cuda.current_stream().wait_stream(side_stream) + + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + launch(_GRAPH_CHUNK) + for _ in range(max_iter // _GRAPH_CHUNK - 1): + graph.replay() + remainder = max_iter % _GRAPH_CHUNK + if remainder: + launch(remainder) + except RuntimeError: + launch(max_iter) # graphs unavailable; extra iterations are harmless + + +def _scaling_iterations(a, b, cost_matrix, epsilon, max_iter, threshold): + """Plain scaling-form Sinkhorn on (n, d) marginals; returns (log_u, log_v).""" + K = torch.exp(-cost_matrix / epsilon) + if _use_fused_kernels(a): + from torchmorph import _C + + a, b = a.contiguous(), b.contiguous() + K_T = _transpose_or_self(K) + u = torch.ones_like(a) + v = torch.ones_like(b) + _run_fused(lambda k: _C.sinkhorn_fastiter(a, b, K, K_T, u, v, k), max_iter, a.device) + return u.log(), v.log() + + u = torch.ones_like(a) + v = torch.ones_like(b) + for i in range(max_iter): + u_prev = u + u = a / (v @ K.mT + 1e-12) + v = b / (u @ K + 1e-12) + if threshold > 0 and i % 10 == 0: + if (u - u_prev).abs().sum(dim=-1).max() <= threshold: + break + return u.log(), v.log() + + +def _log_iterations(a, b, cost_matrix, epsilon, max_iter, threshold): + """Log-domain Sinkhorn on (n, d) marginals; returns (log_u, log_v).""" + log_a, log_b = a.log(), b.log() + if _use_fused_kernels(a): + from torchmorph import _C + + log_a, log_b = log_a.contiguous(), log_b.contiguous() + cost_T = _transpose_or_self(cost_matrix) + log_u = torch.zeros_like(a) + log_v = torch.zeros_like(b) + _run_fused( + lambda k: _C.sinkhorn_logiter( + log_a, log_b, cost_matrix, cost_T, log_u, log_v, k, epsilon + ), + max_iter, + a.device, + ) + return log_u, log_v + + log_K = -cost_matrix / epsilon + log_u = torch.zeros_like(a) + log_v = torch.zeros_like(b) + for i in range(max_iter): + log_u_prev = log_u + log_u = log_a - torch.logsumexp(log_K + log_v.unsqueeze(-2), dim=-1) + log_v = log_b - torch.logsumexp(log_K.mT + log_u.unsqueeze(-2), dim=-1) + if threshold > 0 and i % 10 == 0: + if (log_u - log_u_prev).abs().sum(dim=-1).max() <= threshold: + break + return log_u, log_v + + +class _SinkhornDistance(torch.autograd.Function): + """Transport distance with envelope-theorem gradients. + + The backward pass returns the centered dual potentials, i.e. the exact + gradients of the entropic OT cost w.r.t. the (normalized) marginals. + """ + + @staticmethod + def forward(ctx, a, b, cost_matrix, epsilon, max_iter, threshold, log_space): + iterate = _log_iterations if log_space else _scaling_iterations + log_u, log_v = iterate(a, b, cost_matrix, epsilon, max_iter, threshold) + plan = torch.exp(log_u.unsqueeze(-1) - cost_matrix / epsilon + log_v.unsqueeze(-2)) + distance = (plan * cost_matrix).sum(dim=(-2, -1)) + f = epsilon * log_u + g = epsilon * log_v + ctx.save_for_backward(f - f.mean(dim=-1, keepdim=True), g - g.mean(dim=-1, keepdim=True)) + return distance + + @staticmethod + def backward(ctx, grad_output): + f, g = ctx.saved_tensors + grad = grad_output.unsqueeze(-1) + return grad * f, grad * g, None, None, None, None, None + + +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). + """ + + def __init__(self, epsilon=1.0, max_iter=100, threshold=0.0, p=2, log_space=False): + super().__init__() + if epsilon <= 0: + raise ValueError("epsilon must be positive.") + if max_iter <= 0: + raise ValueError("max_iter must be positive.") + if threshold < 0: + raise ValueError("threshold must be non-negative.") + self.epsilon = epsilon + self.max_iter = max_iter + self.threshold = threshold + self.p = p + self.log_space = log_space + + def extra_repr(self): + return ( + f"epsilon={self.epsilon}, max_iter={self.max_iter}, " + 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. + """ + if source.ndim != 2 or source.shape != target.shape: + raise ValueError("source and target must be (n, d) tensors of the same shape.") + if not source.is_floating_point(): + raise ValueError("source and target must be floating-point tensors.") + d = source.shape[1] + if cost_matrix is None: + cost_matrix = build_cost_matrix((d,), self.p, source.device) + elif cost_matrix.shape != (d, d): + raise ValueError( + f"cost_matrix must have shape ({d}, {d}), got {tuple(cost_matrix.shape)}." + ) + cost_matrix = cost_matrix.to(device=source.device, dtype=source.dtype).contiguous() + + source = source.clamp(min=0) + target = target.clamp(min=0) + source = source / source.sum(dim=-1, keepdim=True).clamp(min=1e-12) + 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.""" + source, target, cost_matrix = self.data_preprocess(source, target, cost_matrix) + return _SinkhornDistance.apply( + source, + target, + cost_matrix, + self.epsilon, + self.max_iter, + self.threshold, + self.log_space, + ) + + @torch.no_grad() + def plan(self, source: Tensor, target: Tensor, cost_matrix: Optional[Tensor] = None): + """(n, d, d) transport plans, reconstructed in log space.""" + 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. + """ + source, target, cost_matrix = self.data_preprocess(source, target, cost_matrix) + log_u, log_v = self._iterate(source, target, cost_matrix) + f = self.epsilon * log_u + g = self.epsilon * log_v + return f - f.mean(dim=-1, keepdim=True), g - g.mean(dim=-1, keepdim=True) + + def _iterate(self, a, b, cost_matrix): + iterate = _log_iterations if self.log_space else _scaling_iterations + return iterate(a, b, cost_matrix, self.epsilon, self.max_iter, self.threshold) From 51a300867450ccdcae3e75a2b99204429b62a1d4 Mon Sep 17 00:00:00 2001 From: Kai ZHAO <694691@qq.com> Date: Wed, 15 Jul 2026 12:27:30 +0800 Subject: [PATCH 23/25] refactor: share validation, cut duplication, drop redundant work - Add torchmorph/_validation.py with shared validate_bcs_input / validate_output; remove three near-identical copies from binary, grey, and distance_transform (error messages preserved). - binary: drop per-iteration tensor .clone() in the stabilize loop (step/where already return fresh tensors) and skip the redundant bool conversion when the value is already bool + contiguous. - Move _validate_origin to structure.py; grey now reuses _normalize_origin/_validate_origin instead of hand-rolling it. - grey: pack the shared structuring-element args via _element_kwargs across the composite operators. Net ~-108 lines; behavior unchanged. Co-Authored-By: Claude Opus 4.8 (1M context) --- torchmorph/_validation.py | 42 +++++++ torchmorph/distance_transform.py | 32 +---- torchmorph/morphology/binary.py | 53 ++------ torchmorph/morphology/grey.py | 186 +++++++---------------------- torchmorph/morphology/structure.py | 8 ++ 5 files changed, 106 insertions(+), 215 deletions(-) create mode 100644 torchmorph/_validation.py diff --git a/torchmorph/_validation.py b/torchmorph/_validation.py new file mode 100644 index 0000000..4a04bbe --- /dev/null +++ b/torchmorph/_validation.py @@ -0,0 +1,42 @@ +from collections.abc import Sequence + +from torch import Tensor + +MAX_SPATIAL_NDIM = 8 + + +def validate_bcs_input(input: Tensor) -> int: + """Validate a ``(B, C, Spatial...)`` CUDA input and return its spatial ndim.""" + if not input.is_cuda: + raise ValueError("Input tensor must be on CUDA device.") + if input.ndim < 3: + raise ValueError( + f"Input must be (B, C, Spatial...) with at least 3 dimensions, got {input.shape}." + ) + spatial_ndim = input.ndim - 2 + if spatial_ndim > MAX_SPATIAL_NDIM: + raise ValueError(f"Input spatial dimensions must be in range 1 to 8, got {spatial_ndim}.") + if input.numel() == 0: + raise ValueError(f"Invalid input: empty tensor with shape {input.shape}.") + return spatial_ndim + + +def validate_output( + input: Tensor, + output: Tensor | None, + expected_shape: Sequence[int] | None = None, + name: str = "output", +) -> None: + """Validate that an optional pre-allocated output matches shape and device.""" + if output is None: + return + expected = input.shape if expected_shape is None else expected_shape + if tuple(output.shape) != tuple(expected): + raise ValueError( + f"{name} shape {tuple(output.shape)} must match expected shape {tuple(expected)}" + ) + if output.device != input.device: + raise ValueError( + f"{name} must be on the same device as input, " + f"got {output.device} and {input.device}" + ) diff --git a/torchmorph/distance_transform.py b/torchmorph/distance_transform.py index bc7e6d9..f5ebf3e 100644 --- a/torchmorph/distance_transform.py +++ b/torchmorph/distance_transform.py @@ -4,6 +4,7 @@ from torch import Tensor from . import _C +from ._validation import validate_bcs_input, validate_output def _normalize_sampling( @@ -33,34 +34,9 @@ def _prepare_distance_transform( distances: Tensor | None, indices: Tensor | None, ) -> tuple[int, bool, bool]: - if not input.is_cuda: - raise ValueError("Input tensor must be on CUDA device.") - if input.ndim < 3: - raise ValueError( - f"Input must be (B, C, Spatial...) with at least 3 dimensions, got {input.shape}." - ) - spatial_ndim = input.ndim - 2 - if spatial_ndim > 8: - raise ValueError(f"Input spatial dimensions must be in range 1 to 8, got {spatial_ndim}.") - if input.numel() == 0: - raise ValueError(f"Invalid input: empty tensor with shape {input.shape}.") - - expected_indices_shape = (spatial_ndim, *input.shape) - for name, output, expected_shape in ( - ("distances", distances, input.shape), - ("indices", indices, expected_indices_shape), - ): - if output is None: - continue - if tuple(output.shape) != tuple(expected_shape): - raise ValueError( - f"{name} shape {output.shape} must match expected shape {expected_shape}" - ) - if output.device != input.device: - raise ValueError( - f"{name} must be on the same device as input, " - f"got {output.device} and {input.device}" - ) + spatial_ndim = validate_bcs_input(input) + validate_output(input, distances, name="distances") + validate_output(input, indices, (spatial_ndim, *input.shape), name="indices") return_distances = return_distances or distances is not None return_indices = return_indices or indices is not None diff --git a/torchmorph/morphology/binary.py b/torchmorph/morphology/binary.py index f383cfa..d47a0c1 100644 --- a/torchmorph/morphology/binary.py +++ b/torchmorph/morphology/binary.py @@ -2,33 +2,8 @@ from torch import Tensor from .. import _C -from .structure import _normalize_origin, generate_binary_structure - - -def _validate_binary_input(input: Tensor) -> int: - if not input.is_cuda: - raise ValueError("Input tensor must be on CUDA device.") - if input.ndim < 3: - raise ValueError( - f"Input must be (B, C, Spatial...) with at least 3 dimensions, got {input.shape}." - ) - spatial_ndim = input.ndim - 2 - if spatial_ndim > 8: - raise ValueError(f"Input spatial dimensions must be in range 1 to 8, got {spatial_ndim}.") - if input.numel() == 0: - raise ValueError(f"Invalid input: empty tensor with shape {input.shape}.") - return spatial_ndim - - -def _validate_binary_output(input: Tensor, output: Tensor | None) -> None: - if output is None: - return - if output.shape != input.shape: - raise ValueError(f"output shape {output.shape} must match input shape {input.shape}") - if output.device != input.device: - raise ValueError( - f"output must be on the same device as input, got {output.device} and {input.device}" - ) +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: @@ -37,14 +12,6 @@ def _normalize_structure(structure: Tensor, spatial_ndim: int, name: str = "stru return (structure != 0).detach().to(device="cpu", dtype=torch.bool).contiguous() -def _validate_origin(origin: tuple[int, ...], structure: Tensor, name: str = "origin") -> None: - for origin_value, structure_size in zip(origin, structure.shape): - min_origin = -(structure_size // 2) - max_origin = (structure_size - 1) // 2 - if not min_origin <= origin_value <= max_origin: - raise ValueError(f"invalid {name}") - - def _binary_morphology_cuda_step( input: Tensor, structure: Tensor, @@ -53,7 +20,7 @@ def _binary_morphology_cuda_step( *, mode: str, ) -> Tensor: - x = (input != 0).contiguous() + 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)) @@ -70,8 +37,8 @@ def _binary_morphology( mode: str, ) -> Tensor: iterate_until_stable = iterations < 1 - spatial_ndim = _validate_binary_input(input) - _validate_binary_output(input, output) + spatial_ndim = validate_bcs_input(input) + validate_output(input, output) if structure is None: structure = generate_binary_structure(spatial_ndim, 1) @@ -107,7 +74,7 @@ def step(value: Tensor) -> Tensor: if old is not None and torch.equal(x, old): break - old = x.clone() + old = x else: for _ in range(iterations): x = step(x) @@ -197,8 +164,8 @@ def binary_fill_holes( origin: int | tuple[int, ...] = 0, ) -> Tensor: """Fill holes in binary objects for `(B, C, Spatial...)` CUDA tensors.""" - _validate_binary_input(input) - _validate_binary_output(input, output) + validate_bcs_input(input) + validate_output(input, output) mask = input == 0 seed = torch.zeros_like(mask, dtype=torch.bool) @@ -227,8 +194,8 @@ def binary_hit_or_miss( origin2: int | tuple[int, ...] | None = None, ) -> Tensor: """N-dimensional binary hit-or-miss transform for `(B, C, Spatial...)` CUDA tensors.""" - spatial_ndim = _validate_binary_input(input) - _validate_binary_output(input, output) + 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) diff --git a/torchmorph/morphology/grey.py b/torchmorph/morphology/grey.py index 958671e..0f53b37 100644 --- a/torchmorph/morphology/grey.py +++ b/torchmorph/morphology/grey.py @@ -2,33 +2,8 @@ from torch import Tensor from .. import _C - - -def _validate_grey_input(input: Tensor) -> int: - if not input.is_cuda: - raise ValueError("Input tensor must be on CUDA device.") - if input.ndim < 3: - raise ValueError( - f"Input must be (B, C, Spatial...) with at least 3 dimensions, got {input.shape}." - ) - spatial_ndim = input.ndim - 2 - if spatial_ndim > 8: - raise ValueError(f"Input spatial dimensions must be in range 1 to 8, got {spatial_ndim}.") - if input.numel() == 0: - raise ValueError(f"Invalid input: empty tensor with shape {input.shape}.") - return spatial_ndim - - -def _validate_grey_output(input: Tensor, output: Tensor | None) -> None: - if output is None: - return - if output.shape != input.shape: - raise ValueError(f"output shape {output.shape} must match input shape {input.shape}") - if output.device != input.device: - raise ValueError( - f"output must be on the same device as input, got {output.device} and {input.device}" - ) - +from .._validation import validate_bcs_input, validate_output +from .structure import _normalize_origin, _validate_origin _MODE_MAP = { 'constant': 0, @@ -39,6 +14,18 @@ def _validate_grey_output(input: Tensor, output: Tensor | None) -> None: } +def _element_kwargs(size, footprint, structure, mode, cval, origin) -> dict: + """Pack the structuring-element arguments shared by every grey operator.""" + return dict( + size=size, + footprint=footprint, + structure=structure, + mode=mode, + cval=cval, + origin=origin, + ) + + def _grey_morphology( input: Tensor, size: int | tuple[int, ...] | None = None, @@ -52,8 +39,8 @@ def _grey_morphology( operation: str, ) -> Tensor: """Shared argument normalization for grey erosion and dilation.""" - spatial_ndim = _validate_grey_input(input) - _validate_grey_output(input, output) + spatial_ndim = validate_bcs_input(input) + validate_output(input, output) footprint_cpu = torch.empty(0, dtype=torch.bool) @@ -91,21 +78,10 @@ def _grey_morphology( ) struct = struct.detach().to(device='cpu', dtype=torch.float32).contiguous() - # Normalize origin - if isinstance(origin, int): - origin_list = [origin] * spatial_ndim - else: - origin_list = list(origin) - if len(origin_list) != spatial_ndim: - raise ValueError( - f'Origin length {len(origin_list)} must match spatial dimension ' f'{spatial_ndim}.' - ) - - for origin_value, structure_size in zip(origin_list, struct.shape): - min_origin = -(structure_size // 2) - max_origin = (structure_size - 1) // 2 - if not min_origin <= origin_value <= max_origin: - raise ValueError("invalid origin") + # Normalize and validate origin + origin_tuple = _normalize_origin(origin, spatial_ndim) + _validate_origin(origin_tuple, struct) + origin_list = list(origin_tuple) # Map mode string to int if mode not in _MODE_MAP: @@ -197,25 +173,9 @@ def grey_opening( Computation uses float32. The result is float32 unless ``output`` is given. """ - eroded = grey_erosion( - input, - size=size, - footprint=footprint, - structure=structure, - mode=mode, - cval=cval, - origin=origin, - ) - return grey_dilation( - eroded, - size=size, - footprint=footprint, - structure=structure, - output=output, - mode=mode, - cval=cval, - origin=origin, - ) + kwargs = _element_kwargs(size, footprint, structure, mode, cval, origin) + eroded = grey_erosion(input, **kwargs) + return grey_dilation(eroded, output=output, **kwargs) def grey_closing( @@ -232,25 +192,9 @@ def grey_closing( Computation uses float32. The result is float32 unless ``output`` is given. """ - dilated = grey_dilation( - input, - size=size, - footprint=footprint, - structure=structure, - mode=mode, - cval=cval, - origin=origin, - ) - return grey_erosion( - dilated, - size=size, - footprint=footprint, - structure=structure, - output=output, - mode=mode, - cval=cval, - origin=origin, - ) + kwargs = _element_kwargs(size, footprint, structure, mode, cval, origin) + dilated = grey_dilation(input, **kwargs) + return grey_erosion(dilated, output=output, **kwargs) def morphological_gradient( @@ -264,27 +208,12 @@ def morphological_gradient( origin: int | tuple[int, ...] = 0, ) -> Tensor: """N-dimensional morphological gradient for ``(B, C, Spatial...)`` CUDA tensors.""" - _validate_grey_input(input) - _validate_grey_output(input, output) + validate_bcs_input(input) + validate_output(input, output) - dilated = grey_dilation( - input, - size=size, - footprint=footprint, - structure=structure, - mode=mode, - cval=cval, - origin=origin, - ) - eroded = grey_erosion( - input, - size=size, - footprint=footprint, - structure=structure, - mode=mode, - cval=cval, - origin=origin, - ) + kwargs = _element_kwargs(size, footprint, structure, mode, cval, origin) + dilated = grey_dilation(input, **kwargs) + eroded = grey_erosion(input, **kwargs) result = dilated - eroded if output is not None: output.copy_(result) @@ -303,18 +232,10 @@ def white_tophat( origin: int | tuple[int, ...] = 0, ) -> Tensor: """N-dimensional white top-hat filter for ``(B, C, Spatial...)`` CUDA tensors.""" - _validate_grey_input(input) - _validate_grey_output(input, output) + validate_bcs_input(input) + validate_output(input, output) - opened = grey_opening( - input, - size=size, - footprint=footprint, - structure=structure, - mode=mode, - cval=cval, - origin=origin, - ) + opened = grey_opening(input, **_element_kwargs(size, footprint, structure, mode, cval, origin)) result = input.detach().float() - opened if output is not None: output.copy_(result) @@ -333,18 +254,10 @@ def black_tophat( origin: int | tuple[int, ...] = 0, ) -> Tensor: """N-dimensional black top-hat filter for ``(B, C, Spatial...)`` CUDA tensors.""" - _validate_grey_input(input) - _validate_grey_output(input, output) + validate_bcs_input(input) + validate_output(input, output) - closed = grey_closing( - input, - size=size, - footprint=footprint, - structure=structure, - mode=mode, - cval=cval, - origin=origin, - ) + closed = grey_closing(input, **_element_kwargs(size, footprint, structure, mode, cval, origin)) result = closed - input.detach().float() if output is not None: output.copy_(result) @@ -363,27 +276,12 @@ def morphological_laplace( origin: int | tuple[int, ...] = 0, ) -> Tensor: """N-dimensional morphological Laplace for ``(B, C, Spatial...)`` CUDA tensors.""" - _validate_grey_input(input) - _validate_grey_output(input, output) + validate_bcs_input(input) + validate_output(input, output) - dilated = grey_dilation( - input, - size=size, - footprint=footprint, - structure=structure, - mode=mode, - cval=cval, - origin=origin, - ) - eroded = grey_erosion( - input, - size=size, - footprint=footprint, - structure=structure, - mode=mode, - cval=cval, - origin=origin, - ) + kwargs = _element_kwargs(size, footprint, structure, mode, cval, origin) + dilated = grey_dilation(input, **kwargs) + eroded = grey_erosion(input, **kwargs) result = dilated + eroded - 2 * input.detach().float() if output is not None: output.copy_(result) diff --git a/torchmorph/morphology/structure.py b/torchmorph/morphology/structure.py index ac42a77..71427fc 100644 --- a/torchmorph/morphology/structure.py +++ b/torchmorph/morphology/structure.py @@ -49,6 +49,14 @@ def _normalize_origin(origin: int | tuple[int, ...], ndim: int) -> tuple[int, .. return origin_tuple +def _validate_origin(origin: tuple[int, ...], structure: Tensor, name: str = "origin") -> None: + for origin_value, structure_size in zip(origin, structure.shape): + min_origin = -(structure_size // 2) + max_origin = (structure_size - 1) // 2 + if not min_origin <= origin_value <= max_origin: + raise ValueError(f"invalid {name}") + + def iterate_structure( structure: Tensor, iterations: int, From c8c8c51a77461636450af816d331bf762a6d082e Mon Sep 17 00:00:00 2001 From: Yuhandeng Date: Wed, 15 Jul 2026 19:42:22 +0800 Subject: [PATCH 24/25] update requirements --- requirements-test.txt | 1 + 1 file changed, 1 insertion(+) diff --git a/requirements-test.txt b/requirements-test.txt index 216f071..b571a83 100644 --- a/requirements-test.txt +++ b/requirements-test.txt @@ -3,4 +3,5 @@ # Testing and validation pytest>=7.0 pytest-cov>=4.0 +POT From 7d22b2d89764c13e51a99f29372de4a26c673c94 Mon Sep 17 00:00:00 2001 From: Yuhandeng Date: Wed, 15 Jul 2026 19:43:35 +0800 Subject: [PATCH 25/25] smaller head file --- torchmorph/csrc/ot_kernel.cu | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/torchmorph/csrc/ot_kernel.cu b/torchmorph/csrc/ot_kernel.cu index 193d1ac..4efbb7c 100644 --- a/torchmorph/csrc/ot_kernel.cu +++ b/torchmorph/csrc/ot_kernel.cu @@ -1,5 +1,5 @@ #include -#include +#include #include #include