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 d849f48..0000000 --- a/CONTEXT.md +++ /dev/null @@ -1,53 +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 - -**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 4d4f69d..ede098a 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: diff --git a/benchmark/binary.py b/benchmark/binary.py index 319df3c..f1bdbd7 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,35 @@ 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), + "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), +} -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 +48,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/benchmark/distance_transform.py b/benchmark/distance_transform.py index 465f011..3dd09b1 100644 --- a/benchmark/distance_transform.py +++ b/benchmark/distance_transform.py @@ -1,398 +1,216 @@ 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 = {"torchmorph": tm.euclidean_distance_transform} + 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/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/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 diff --git a/test/test_binary.py b/test/test_binary.py index c961b45..f9e681c 100644 --- a/test/test_binary.py +++ b/test/test_binary.py @@ -1,13 +1,16 @@ -import numpy as np # noqa: F401 +import importlib + +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_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 import torchmorph as tm @@ -15,312 +18,511 @@ not torch.cuda.is_available(), reason="CUDA is required for torchmorph tests" ) +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"}), +] + + +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 + + +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]]]], + 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 + +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) +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) + sample_output = None 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) + result = scipy_op(sample, **sample_kwargs) + if result is None: + result = sample_output 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( + return np.stack(results).reshape(*batch_shape, *spatial_shape) + + +def optional_cuda_tensor(value, dtype): + if value is None: + return None + return torch.as_tensor(value, dtype=dtype, device="cuda") + + +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", "structure", "iterations", "origin", "border_value"), + [ + 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_morphology_matches_scipy( + torch_op, scipy_op, np_input, structure, iterations, origin, border_value +): + 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 = apply_scipy_to_batch( np_input, - scipy_func, + scipy_op, structure=structure, iterations=iterations, - mask=mask, - output=output, - border_value=border_value, origin=origin, + border_value=border_value, ) + torch.testing.assert_close(result.cpu(), torch.as_tensor(expected)) -# case -case_2d = np.array( +@pytest.mark.parametrize(("torch_op", "scipy_op"), BINARY_OPERATORS) +@pytest.mark.parametrize( + ("np_input", "mask"), [ - [ - [ - [0, 1, 0], - [1, 1, 1], - [0, 1, 0], - ] - ] + 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_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 = 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"), 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"), ], - 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 +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)) + -# structure -structure_2d = np.array([[0, 1, 0], [1, 1, 1], [0, 1, 0]]) +@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)) -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) +@pytest.mark.parametrize(("torch_op", "scipy_op"), BINARY_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_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) -# 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) + result = torch_op(x, output=output) -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 + 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)) -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]) +@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) -output_3d_np = np.empty([5, 5, 5]) -output_3d_tm = torch.empty([2, 1, 5, 5, 5]) + result = torch_op(x, output=output) -output_4d_np = np.empty([4, 4, 4, 4]) -output_4d_tm = torch.empty([2, 1, 4, 4, 4, 4]) + 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"), FILL_HOLES_OPERATORS) @pytest.mark.parametrize( - ("np_input, scipy_func, structure, iterations, origin, border_value"), + ("np_input", "structure", "origin"), [ - 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_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_basic( - np_input, - scipy_func, - 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, - iterations=iterations, +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, - 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"), 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( - ("input_np, scipy_func, mask_np, mask_tm"), + ("np_input", "structure1", "structure2", "origin1", "origin2"), [ - 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, 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_mask( - input_np, - scipy_func, - mask_np, - mask_tm, +def test_binary_hit_or_miss_matches_scipy( + torch_op, scipy_op, np_input, structure1, structure2, origin1, origin2 ): - x = torch.as_tensor(input_np, dtype=torch.float32) - tm_func = getattr(tm, scipy_func.__name__) - actual = tm_func( - x, - mask=mask_tm, + 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_np = batch_scipy( - input_np, - scipy_func, - mask=mask_np, + expected = apply_scipy_to_batch( + np_input, + scipy_op, + structure1=structure1, + structure2=structure2, + origin1=origin1, + origin2=origin2, ) - 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"), 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( - "input_np,scipy_func, output_np, output_tm", + ("torch_op", "expected_value"), [ - 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(tm.binary_erosion, True, id="erosion"), + pytest.param(tm.binary_dilation, False, id="dilation"), ], ) -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_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") - actual = tm_func( - x, - output=output_tm, - ) + result = torch_op(x, structure=structure) - 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) + expected = torch.full_like(x, expected_value, dtype=torch.bool) + torch.testing.assert_close(result, 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", 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"): + torch_op(x) + + +@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)) + + +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): + x = torch.as_tensor(CASE_2D, dtype=torch.float32, device="cuda") + result = torch_op(x) + stream.synchronize() + + 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): + 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 + + +@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/test/test_distance_transform.py b/test/test_distance_transform.py index bf232b3..905db96 100644 --- a/test/test_distance_transform.py +++ b/test/test_distance_transform.py @@ -1,865 +1,456 @@ -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, - ) - 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 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 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 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], - ] - ] - ], - 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 - -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) +def cuda(array): + return torch.as_tensor(np.ascontiguousarray(array), device="cuda") -# 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], +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) + 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", + ) + deltas = [ + (axis - sample_indices[dim]).abs() * sampling[dim] for dim, axis in enumerate(axes) ] - ] - ], - dtype=np.float32, -) + 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()) -# 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) + assert result is None + torch.testing.assert_close(distances, expected_distances.double()) + torch.testing.assert_close(indices, expected_indices.int()) - torch.testing.assert_close(dist_cuda, dist_ref, atol=1e-5, rtol=1e-5) - -@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) - - torch.testing.assert_close(dist_cuda, dist_ref, atol=1e-5, rtol=1e-5) - - -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) + assert isinstance(result, torch.Tensor) + assert result.shape == (2, *input.shape) - # 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") +@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), + ) - 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" - ) - - # 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) +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()) - 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() - dist_exact = tm.euclidean_distance_transform(x, algorithm="exact") - dist_jfa = tm.euclidean_distance_transform(x, algorithm="jfa") +def test_edt_rejects_removed_algorithm_argument(): + with pytest.raises(TypeError, match="algorithm"): + tm.euclidean_distance_transform(cuda(CASES[2]), algorithm="exact") - # 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() - # 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"), - ], -) -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) +@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( - "alias, canonical", - [ - pytest.param("cityblock", "taxicab", id="cityblock"), - pytest.param("manhattan", "taxicab", id="manhattan"), - ], -) -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") - with pytest.raises(ValueError, match="metric must be"): - tm.chamfer_distance_transform(x, metric="invalid") +@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"), + ) -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 +@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"), + ) - with pytest.raises(ValueError, match="CUDA"): - tm.chamfer_distance_transform(x) +@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 -@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") + 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)) - 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("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( - "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"), - ], + ("alias", "canonical"), [("cityblock", "taxicab"), ("manhattan", "taxicab")] ) -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) - +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.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", + ("distance_op", "kwargs"), [ - 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"), + pytest.param(tm.chamfer_distance_transform, {"metric": "invalid"}, id="cdt"), + pytest.param(tm.brute_force_distance_transform, {"metric": "invalid"}, id="bfdt"), ], ) -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") - - 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) +def test_invalid_metric(distance_op, kwargs): + with pytest.raises(ValueError, match="metric must be"): + distance_op(cuda(CASES[2]), **kwargs) + + +@pytest.mark.skipif(torch.cuda.device_count() < 2, reason="requires two CUDA devices") +@pytest.mark.parametrize("distance_op", DISTANCE_OPERATORS) +@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") + 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) - torch.testing.assert_close(dist_cuda, dist_ref, atol=1e-5, rtol=1e-5) + assert all(output.device == input.device for output in outputs) diff --git a/test/test_grey.py b/test/test_grey.py index c9236fd..2862613 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") @@ -190,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/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..ac15e25 100644 --- a/torchmorph/__init__.py +++ b/torchmorph/__init__.py @@ -7,7 +7,10 @@ binary_closing, binary_dilation, binary_erosion, + binary_fill_holes, + binary_hit_or_miss, binary_opening, + binary_propagation, black_tophat, generate_binary_structure, grey_closing, @@ -19,6 +22,7 @@ morphological_laplace, white_tophat, ) +from .optimal_transport import SinkhornSolver, build_cost_matrix __all__ = [ "euclidean_distance_transform", @@ -28,8 +32,11 @@ "iterate_structure", "binary_erosion", "binary_dilation", + "binary_fill_holes", + "binary_hit_or_miss", "binary_opening", "binary_closing", + "binary_propagation", "grey_erosion", "grey_dilation", "grey_opening", @@ -38,4 +45,6 @@ "morphological_laplace", "white_tophat", "black_tophat", + "SinkhornSolver", + "build_cost_matrix", ] 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/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/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/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) { 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/ot_kernel.cu b/torchmorph/csrc/ot_kernel.cu new file mode 100644 index 0000000..4efbb7c --- /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..5a87c4c 100644 --- a/torchmorph/csrc/torchmorph.cpp +++ b/torchmorph/csrc/torchmorph.cpp @@ -18,12 +18,25 @@ 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, bool return_distances, - bool return_indices, - const std::string& algorithm + bool return_indices ); std::tuple cdt_cuda( @@ -41,6 +54,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, @@ -61,13 +95,27 @@ 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"), 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", py::arg("input"), @@ -82,4 +130,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/distance_transform.py b/torchmorph/distance_transform.py index 03bc1d1..f5ebf3e 100644 --- a/torchmorph/distance_transform.py +++ b/torchmorph/distance_transform.py @@ -1,382 +1,175 @@ +import math from collections.abc import Sequence -import torch +from torch import Tensor -from torchmorph import _C +from . import _C +from ._validation import validate_bcs_input, validate_output -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]) - """ - 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." - ) - 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: - 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() - total_ndim = input.ndim - spatial_ndim = total_ndim - 2 # Exclude B and C dimensions - - # Process sampling parameter for spatial dimensions only +def _normalize_sampling( + sampling: float | Sequence[float] | None, + spatial_ndim: int, +) -> list[float]: if sampling is None: - # Unit spacing for all spatial dimensions - sampling_list = [1.0] * spatial_ndim + values = [1.0] * spatial_ndim elif isinstance(sampling, (int, float)): - # Single value: same spacing for all spatial dimensions - sampling_list = [float(sampling)] * spatial_ndim + values = [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: - raise ValueError( - f"sampling has {len(sampling_list)} but input {spatial_ndim} dimensions " - f"(input shape: {input.shape}, format: (B, C, Spatial...))" - ) - - # 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 + 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]: + 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 + 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, +) -> 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, + ) + normalized_sampling = _normalize_sampling(sampling, spatial_ndim) + raw_distances, raw_indices = _C.edt_cuda( + input.float().contiguous(), + normalized_sampling, + return_distances, + return_indices, + ) + 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/__init__.py b/torchmorph/morphology/__init__.py index bf6dea0..4388484 100644 --- a/torchmorph/morphology/__init__.py +++ b/torchmorph/morphology/__init__.py @@ -1,4 +1,12 @@ -from .binary import binary_closing, binary_dilation, binary_erosion, binary_opening +from .binary import ( + binary_closing, + binary_dilation, + binary_erosion, + binary_fill_holes, + binary_hit_or_miss, + binary_opening, + binary_propagation, +) from .grey import ( black_tophat, grey_closing, @@ -16,8 +24,11 @@ "iterate_structure", "binary_erosion", "binary_dilation", + "binary_fill_holes", + "binary_hit_or_miss", "binary_opening", "binary_closing", + "binary_propagation", "grey_erosion", "grey_dilation", "grey_opening", 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 49ec8b9..d47a0c1 100644 --- a/torchmorph/morphology/binary.py +++ b/torchmorph/morphology/binary.py @@ -1,222 +1,290 @@ -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 +from torch import Tensor + +from .. import _C +from .._validation import validate_bcs_input, validate_output +from .structure import _normalize_origin, _validate_origin, generate_binary_structure + + +def _normalize_structure(structure: Tensor, spatial_ndim: int, name: str = "structure") -> Tensor: + if structure.ndim != spatial_ndim: + raise ValueError(f"{name} dimension is not {spatial_ndim}, got {structure.ndim}") + return (structure != 0).detach().to(device="cpu", dtype=torch.bool).contiguous() + + +def _binary_morphology_cuda_step( + input: Tensor, + structure: Tensor, + border_value: bool, + origin: tuple[int, ...], + *, + mode: str, +) -> Tensor: + x = input if input.dtype == torch.bool and input.is_contiguous() else (input != 0).contiguous() + kernel = _C.binary_erosion_cuda if mode == "erosion" else _C.binary_dilation_cuda + return kernel(x, structure, list(origin), bool(border_value)) + + +def _binary_morphology( + input: Tensor, + structure: Tensor | None, + iterations: int, + mask: Tensor | None, + output: Tensor | None, + border_value: bool, + origin: int | tuple[int, ...], + *, + mode: str, +) -> Tensor: + iterate_until_stable = iterations < 1 + spatial_ndim = validate_bcs_input(input) + validate_output(input, output) + + if structure is None: + structure = generate_binary_structure(spatial_ndim, 1) + structure = _normalize_structure(structure, spatial_ndim) + + origin = _normalize_origin(origin, spatial_ndim) + _validate_origin(origin, structure) + x = input != 0 + input_bool = x + if mask is not None and mask.shape != input.shape: + raise ValueError(f"mask shape {mask.shape} must match input shape {input.shape}") + mask_bool = mask.to(device=input.device, dtype=torch.bool) if mask is not None else None + structure_is_empty = not structure.any().item() + + def step(value: Tensor) -> Tensor: + if structure_is_empty: + return torch.full_like(value, mode == "erosion", dtype=torch.bool) + return _binary_morphology_cuda_step( + value, + structure, + border_value, + origin, + mode=mode, + ) + + if iterate_until_stable: + old = None + while True: + x = step(x) + if mask_bool is not None: + x = torch.where(mask_bool, x, input_bool) + + if old is not None and torch.equal(x, old): + break + + old = x + else: + for _ in range(iterations): + x = step(x) + if mask_bool is not None: + x = torch.where(mask_bool, x, input_bool) + + result = x.to(dtype=torch.bool) + if output is not None: + output.copy_(result) + return output + return result + + +def binary_erosion( + input: Tensor, + structure: Tensor | None = None, + iterations: int = 1, + mask: Tensor | None = None, + output: Tensor | None = None, + border_value: bool = False, + origin: int | tuple[int, ...] = 0, +) -> Tensor: + """ + N-dimensional binary erosion for `(B, C, Spatial...)` CUDA tensors. + + For a single image or volume, add batch and channel dimensions first. + """ + return _binary_morphology( + input=input, + structure=structure, + iterations=iterations, + mask=mask, + output=output, + border_value=border_value, + origin=origin, + mode="erosion", + ) + + +def binary_dilation( + input: Tensor, + structure: Tensor | None = None, + iterations: int = 1, + mask: Tensor | None = None, + output: Tensor | None = None, + border_value: bool = False, + origin: int | tuple[int, ...] = 0, +) -> Tensor: + """N-dimensional binary dilation for `(B, C, Spatial...)` CUDA tensors.""" + return _binary_morphology( + input, + structure, + iterations, + mask, + output, + border_value, + origin, + mode="dilation", + ) + + +def binary_propagation( + input: Tensor, + structure: Tensor | None = None, + mask: Tensor | None = None, + output: Tensor | None = None, + border_value: bool = False, + origin: int | tuple[int, ...] = 0, +) -> Tensor: + """N-dimensional binary propagation for `(B, C, Spatial...)` CUDA tensors.""" + return _binary_morphology( + input, + structure, + -1, + mask, + output, + border_value, + origin, + mode="dilation", + ) + + +def binary_fill_holes( + input: Tensor, + structure: Tensor | None = None, + output: Tensor | None = None, + origin: int | tuple[int, ...] = 0, +) -> Tensor: + """Fill holes in binary objects for `(B, C, Spatial...)` CUDA tensors.""" + validate_bcs_input(input) + validate_output(input, output) + + mask = input == 0 + seed = torch.zeros_like(mask, dtype=torch.bool) + background = binary_propagation( + seed, + structure=structure, + mask=mask, + output=None, + border_value=True, + origin=origin, + ) + result = torch.logical_not(background) + + if output is not None: + output.copy_(result) + return output + return result + + +def binary_hit_or_miss( + input: Tensor, + structure1: Tensor | None = None, + structure2: Tensor | None = None, + output: Tensor | None = None, + origin1: int | tuple[int, ...] = 0, + origin2: int | tuple[int, ...] | None = None, +) -> Tensor: + """N-dimensional binary hit-or-miss transform for `(B, C, Spatial...)` CUDA tensors.""" + spatial_ndim = validate_bcs_input(input) + validate_output(input, output) + origin1 = _normalize_origin(origin1, spatial_ndim) + origin2 = origin1 if origin2 is None else _normalize_origin(origin2, spatial_ndim) + + if structure1 is None: + structure1 = generate_binary_structure(spatial_ndim, 1) + structure1 = _normalize_structure(structure1, spatial_ndim, "structure1") + _validate_origin(origin1, structure1, "origin1") + + if structure2 is None: + structure2 = torch.logical_not(structure1).contiguous() + else: + structure2 = _normalize_structure(structure2, spatial_ndim, "structure2") + _validate_origin(origin2, structure2, "origin2") + + input_bool = input != 0 + if structure1.any().item(): + hit = binary_erosion(input_bool, structure=structure1, origin=origin1) + else: + hit = torch.ones_like(input_bool, dtype=torch.bool) + + if structure2.any().item(): + miss = binary_erosion(input_bool == 0, structure=structure2, origin=origin2) + result = torch.logical_and(hit, miss) + else: + result = hit + + if output is not None: + output.copy_(result) + return output + return result + + +def binary_opening( + input: Tensor, + structure: Tensor | None = None, + iterations: int = 1, + mask: Tensor | None = None, + output: Tensor | None = None, + border_value: bool = False, + origin: int | tuple[int, ...] = 0, +) -> Tensor: + """N-dimensional binary opening for `(B, C, Spatial...)` CUDA tensors.""" + x = binary_erosion( + input, + structure, + iterations, + mask, + None, + border_value, + origin, + ) + x = binary_dilation( + x, + structure, + iterations, + mask, + output, + border_value, + origin, + ) + return x + + +def binary_closing( + input: Tensor, + structure: Tensor | None = None, + iterations: int = 1, + mask: Tensor | None = None, + output: Tensor | None = None, + border_value: bool = False, + origin: int | tuple[int, ...] = 0, +) -> Tensor: + """N-dimensional binary closing for `(B, C, Spatial...)` CUDA tensors.""" + x = binary_dilation( + input, + structure, + iterations, + mask, + None, + border_value, + origin, + ) + x = binary_erosion( + x, + structure, + iterations, + mask, + output, + border_value, + origin, + ) + return x diff --git a/torchmorph/morphology/grey.py b/torchmorph/morphology/grey.py index b3a52e2..0f53b37 100644 --- a/torchmorph/morphology/grey.py +++ b/torchmorph/morphology/grey.py @@ -2,6 +2,8 @@ from torch import Tensor from .. import _C +from .._validation import validate_bcs_input, validate_output +from .structure import _normalize_origin, _validate_origin _MODE_MAP = { 'constant': 0, @@ -12,6 +14,18 @@ } +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, @@ -25,16 +39,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_bcs_input(input) + validate_output(input, output) footprint_cpu = torch.empty(0, dtype=torch.bool) @@ -60,6 +66,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.') @@ -70,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: @@ -176,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( @@ -211,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( @@ -243,24 +208,12 @@ def morphological_gradient( 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, - ) + validate_bcs_input(input) + validate_output(input, output) + + 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) @@ -279,16 +232,11 @@ def white_tophat( 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 + validate_bcs_input(input) + validate_output(input, output) + + 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) return output @@ -306,16 +254,11 @@ def black_tophat( 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() + validate_bcs_input(input) + validate_output(input, output) + + 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) return output @@ -333,25 +276,13 @@ def morphological_laplace( 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() + validate_bcs_input(input) + validate_output(input, output) + + 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) return output diff --git a/torchmorph/morphology/structure.py b/torchmorph/morphology/structure.py index 573f29b..71427fc 100644 --- a/torchmorph/morphology/structure.py +++ b/torchmorph/morphology/structure.py @@ -39,14 +39,22 @@ 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 _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( @@ -85,5 +93,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] 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)