From 9ac76774005fd1fc3cd4ed9072bcecba704001a2 Mon Sep 17 00:00:00 2001 From: Juan Miguel Giraldo Date: Sat, 27 Jun 2026 14:59:31 -0400 Subject: [PATCH 1/2] Vectorize decohesion causal smear --- .../batchgeneratorsv2.egg-info/SOURCES.txt | 3 ++ .../transforms/spatial/decohesion.py | 42 ++++++++++++++----- 2 files changed, 34 insertions(+), 11 deletions(-) diff --git a/segmentation/models/batchgeneratorsv2/batchgeneratorsv2.egg-info/SOURCES.txt b/segmentation/models/batchgeneratorsv2/batchgeneratorsv2.egg-info/SOURCES.txt index e2d11c3f49..82bc90294b 100644 --- a/segmentation/models/batchgeneratorsv2/batchgeneratorsv2.egg-info/SOURCES.txt +++ b/segmentation/models/batchgeneratorsv2/batchgeneratorsv2.egg-info/SOURCES.txt @@ -32,10 +32,13 @@ batchgeneratorsv2/transforms/noise/__init__.py batchgeneratorsv2/transforms/noise/extranoisetransforms.py batchgeneratorsv2/transforms/noise/gaussian_blur.py batchgeneratorsv2/transforms/spatial/__init__.py +batchgeneratorsv2/transforms/spatial/decohesion.py batchgeneratorsv2/transforms/spatial/low_resolution.py batchgeneratorsv2/transforms/spatial/mirroring.py batchgeneratorsv2/transforms/spatial/spatial.py +batchgeneratorsv2/transforms/spatial/squeeze.py batchgeneratorsv2/transforms/spatial/transpose.py +batchgeneratorsv2/transforms/spatial/warp.py batchgeneratorsv2/transforms/utils/__init__.py batchgeneratorsv2/transforms/utils/compose.py batchgeneratorsv2/transforms/utils/cropping.py diff --git a/segmentation/models/batchgeneratorsv2/batchgeneratorsv2/transforms/spatial/decohesion.py b/segmentation/models/batchgeneratorsv2/batchgeneratorsv2/transforms/spatial/decohesion.py index 843de7c5ce..98aed6173f 100644 --- a/segmentation/models/batchgeneratorsv2/batchgeneratorsv2/transforms/spatial/decohesion.py +++ b/segmentation/models/batchgeneratorsv2/batchgeneratorsv2/transforms/spatial/decohesion.py @@ -9,7 +9,7 @@ IMAGE-ONLY: decohesion is an imaging artifact -- it changes appearance, NOT geometry, so it must not touch the segmentation/labels. """ -from typing import Optional, Sequence, Tuple +from typing import Optional, Tuple import torch import torch.nn.functional as F @@ -49,15 +49,35 @@ def get_parameters(self, **data_dict) -> dict: @staticmethod def _causal_smear(img: torch.Tensor, taxis: int, k: torch.Tensor) -> torch.Tensor: - out = k[0] * img - n = img.shape[taxis] - for i in range(1, k.shape[0]): - dst = [slice(None)] * img.ndim - src = [slice(None)] * img.ndim - dst[taxis] = slice(i, None) - src[taxis] = slice(0, n - i) - out[tuple(dst)] = out[tuple(dst)] + k[i] * img[tuple(src)] - return out + C = img.shape[0] + dim = img.ndim - 1 + K = int(k.shape[0]) + + if K == 1: + return k[0] * img + + if dim not in (2, 3): + raise ValueError(f"_causal_smear supports 2D or 3D inputs, got dim={dim}") + + spatial_axis = taxis - 1 + if not (0 <= spatial_axis < dim): + raise ValueError(f"taxis={taxis} out of range for img of shape {tuple(img.shape)}") + + pad = [0, 0] * dim + pad_idx = (dim - 1 - spatial_axis) * 2 + pad[pad_idx] = K - 1 + x = F.pad(img.unsqueeze(0), pad, mode='constant', value=0.0) + + ks = [1] * dim + ks[spatial_axis] = K + w = torch.flip(k, dims=[0]).to(dtype=x.dtype, device=x.device) + w = w.view([1, 1] + ks).expand(C, 1, *ks).contiguous() + + if dim == 3: + out = F.conv3d(x, w, groups=C) + else: + out = F.conv2d(x, w, groups=C) + return out.squeeze(0) def _density_weight(self, img: torch.Tensor, strength: float) -> torch.Tensor: if not self.density_modulated: @@ -134,4 +154,4 @@ def _apply_to_image(self, img: torch.Tensor, **params) -> torch.Tensor: if dev == 'cuda': torch.cuda.synchronize() print(f"[ok] {shape}: {(time.time() - st) / n * 1000:.2f} ms/sample on {dev}") - print("\nAll checks passed.") + print("\nAll checks passed.") \ No newline at end of file From d414942395aa7fd0b43a204382eeda1232894e9b Mon Sep 17 00:00:00 2001 From: Juan Miguel Giraldo Date: Sat, 27 Jun 2026 17:26:10 -0400 Subject: [PATCH 2/2] Update for GPU --- .../transforms/spatial/decohesion.py | 36 +++++++++++-------- 1 file changed, 22 insertions(+), 14 deletions(-) diff --git a/segmentation/models/batchgeneratorsv2/batchgeneratorsv2/transforms/spatial/decohesion.py b/segmentation/models/batchgeneratorsv2/batchgeneratorsv2/transforms/spatial/decohesion.py index 98aed6173f..c6b798145b 100644 --- a/segmentation/models/batchgeneratorsv2/batchgeneratorsv2/transforms/spatial/decohesion.py +++ b/segmentation/models/batchgeneratorsv2/batchgeneratorsv2/transforms/spatial/decohesion.py @@ -49,7 +49,6 @@ def get_parameters(self, **data_dict) -> dict: @staticmethod def _causal_smear(img: torch.Tensor, taxis: int, k: torch.Tensor) -> torch.Tensor: - C = img.shape[0] dim = img.ndim - 1 K = int(k.shape[0]) @@ -63,21 +62,30 @@ def _causal_smear(img: torch.Tensor, taxis: int, k: torch.Tensor) -> torch.Tenso if not (0 <= spatial_axis < dim): raise ValueError(f"taxis={taxis} out of range for img of shape {tuple(img.shape)}") - pad = [0, 0] * dim - pad_idx = (dim - 1 - spatial_axis) * 2 - pad[pad_idx] = K - 1 - x = F.pad(img.unsqueeze(0), pad, mode='constant', value=0.0) + last = img.ndim - 1 + if taxis == last: + moved = img.contiguous() + permuted = False + perm = None + else: + perm = list(range(img.ndim)) + perm[taxis], perm[last] = perm[last], perm[taxis] + moved = img.permute(perm).contiguous() + permuted = True - ks = [1] * dim - ks[spatial_axis] = K - w = torch.flip(k, dims=[0]).to(dtype=x.dtype, device=x.device) - w = w.view([1, 1] + ks).expand(C, 1, *ks).contiguous() + head_shape = moved.shape[:-1] + L = moved.shape[-1] - if dim == 3: - out = F.conv3d(x, w, groups=C) - else: - out = F.conv2d(x, w, groups=C) - return out.squeeze(0) + x = moved.reshape(-1, 1, L) + x = F.pad(x, (K - 1, 0), mode='constant', value=0.0) + + w = torch.flip(k, dims=[0]).to(dtype=x.dtype, device=x.device).view(1, 1, K) + y = F.conv1d(x, w) + y = y.reshape(*head_shape, L) + + if permuted: + y = y.permute(perm).contiguous() + return y def _density_weight(self, img: torch.Tensor, strength: float) -> torch.Tensor: if not self.density_modulated: