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..c6b798145b 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,43 @@ 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 + 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)}") + + 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 + + head_shape = moved.shape[:-1] + L = moved.shape[-1] + + 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: @@ -134,4 +162,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