Optimize DecohesionTransform using F.conv1d (#201)

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

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.")