batchgeneratorsv2 0.3.3__tar.gz → 0.3.5__tar.gz

{batchgeneratorsv2-0.3.3 → batchgeneratorsv2-0.3.5}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: batchgeneratorsv2
-Version: 0.3.3
+Version: 0.3.5
 Summary: Batchgenerators but better
 Author: Helmholtz Imaging Applied Computer Vision Lab
 Author-email: Fabian Isensee <f.isensee@dkfz-heidelberg.de>

{batchgeneratorsv2-0.3.3 → batchgeneratorsv2-0.3.5}/batchgeneratorsv2/transforms/intensity/brightness.py

@@ -30,7 +30,7 @@ class MultiplicativeBrightnessTransform(ImageOnlyTransform):
         # even though this is array notation it's a lot slower. Shame shame
         # img[params['apply_to_channel']] *= params['multipliers'].view(-1, *[1]*(img.ndim - 1))
         for c, m in zip(params['apply_to_channel'], params['multipliers']):
-            img[c] *= m
+            img[c].mul_(m)  # in-place on the channel view; avoids the __setitem__ write-back copy of `img[c] *= m`
         return img
 
 

{batchgeneratorsv2-0.3.3 → batchgeneratorsv2-0.3.5}/batchgeneratorsv2/transforms/intensity/contrast.py

@@ -73,8 +73,7 @@ class ContrastTransform(ImageOnlyTransform):
 
                 x = img[c]
                 mean = x.mean()
-                minm = x.min()
-                maxm = x.max()
+                minm, maxm = torch.aminmax(x)  # single reduction pass instead of separate min()/max()
 
                 x.sub_(mean)
                 x.mul_(m)

{batchgeneratorsv2-0.3.3 → batchgeneratorsv2-0.3.5}/batchgeneratorsv2/transforms/intensity/gamma.py

@@ -79,8 +79,7 @@ class GammaTransform(ImageOnlyTransform):
                 mean = x.mean()
                 std = x.std()
 
-            minm = x.min()
-            maxm = x.max()
+            minm, maxm = torch.aminmax(x)  # single reduction pass instead of separate min()/max()
             rnge = maxm - minm
             denom = torch.clamp(rnge, min=eps)
 

{batchgeneratorsv2-0.3.3 → batchgeneratorsv2-0.3.5}/batchgeneratorsv2/transforms/intensity/gaussian_noise.py

@@ -66,9 +66,10 @@ class GaussianNoiseTransform(ImageOnlyTransform):
             sigma = sigmas
             noise = torch.empty((n, *spatial), device=device, dtype=dtype).normal_(mean=0.0, std=float(sigma))
 
-        # Advanced indexing (img[idx]) returns a copy, so use indexed assignment
-        # to make sure modifications are written back to img.
-        img[idx] += noise
+        # index_add_ adds noise[i] into channel img[idx[i]] fully in place. idx is unique (from nonzero),
+        # so there is no accumulation ambiguity. This avoids the full (n, *spatial) gather copy + scatter that
+        # advanced-indexed `img[idx] += noise` would allocate.
+        img.index_add_(0, idx, noise)
         return img
     
 

{batchgeneratorsv2-0.3.3 → batchgeneratorsv2-0.3.5}/batchgeneratorsv2/transforms/intensity/inversion.py

@@ -31,10 +31,11 @@ class InvertImageTransform(ImageOnlyTransform):
             return img
         else:
             for ch in params['apply_to_channel']:
-                mn = img[ch].mean()
-                img[ch] -= mn
-                img[ch] *= -1
-                img[ch] += mn
+                x = img[ch]  # channel view; operate in place to drop the per-op __setitem__ write-backs
+                mn = x.mean()
+                x.sub_(mn)
+                x.mul_(-1)
+                x.add_(mn)
         return img
 
 

{batchgeneratorsv2-0.3.3 → batchgeneratorsv2-0.3.5}/batchgeneratorsv2/transforms/intensity/random_clip.py

@@ -67,18 +67,18 @@ class CutOffOutliersTransform(ImageOnlyTransform):
             if retain_std[c]:
                 orig_std = img_c.std()
 
-            # Percentiles in torch to avoid numpy roundtrip
+            # Percentiles in torch to avoid numpy roundtrip; must run on the pre-clamp values
             q = torch.tensor([perc[0] / 100.0, perc[1] / 100.0], device=img_c.device, dtype=torch.float32)
             lower_val, upper_val = torch.quantile(img_c.float(), q)
 
-            img_c_clipped = img_c.clamp(min=lower_val.item(), max=upper_val.item())
+            # clamp in place on the channel view
+            img_c.clamp_(min=lower_val.item(), max=upper_val.item())
 
             if retain_std[c]:
-                clipped_std = img_c_clipped.std()
+                clipped_std = img_c.std()
                 if clipped_std > 1e-8:
-                    img_c_clipped = (img_c_clipped - img_c_clipped.mean()) / clipped_std * orig_std + img_c_clipped.mean()
-
-            img[c] = img_c_clipped
+                    m = img_c.mean()  # reused for both centering steps
+                    img[c] = (img_c - m) / clipped_std * orig_std + m
 
         return img
 

{batchgeneratorsv2-0.3.3 → batchgeneratorsv2-0.3.5}/batchgeneratorsv2/transforms/local/brightness_gradient.py

@@ -140,6 +140,10 @@ class BrightnessGradientAdditiveTransform(ImageOnlyTransform, LocalTransform):
         return {'kernels': kernels}
 
     def _apply_to_image(self, img: torch.Tensor, **params) -> torch.Tensor:
+        if self.clip_intensities:
+            original_min = img.min()
+            original_max = img.max()
+
         for c, kernel in enumerate(params['kernels']):
             if kernel is None:
                 continue
@@ -147,7 +151,7 @@ class BrightnessGradientAdditiveTransform(ImageOnlyTransform, LocalTransform):
             img[c].add_(kernel_tensor)
 
         if self.clip_intensities:
-            img.clamp_(min=img.min(), max=img.max())
+            img.clamp_(min=original_min, max=original_max)
 
         return img
 

{batchgeneratorsv2-0.3.3 → batchgeneratorsv2-0.3.5}/batchgeneratorsv2/transforms/local/local_contrast.py

@@ -40,7 +40,7 @@ class LocalContrastTransform(ImageOnlyTransform, LocalTransform):
             return {'kernels': [None] * C, 'contrasts': [None] * C}
 
         if self.same_for_all_channels:
-            kernel = self._generate_kernel(spatial).astype(np.float32)
+            kernel = self._generate_kernel(spatial)
             contrast = sample_scalar(self.new_contrast)
 
             kernels = [kernel if apply else None for apply in apply_channel]
@@ -52,7 +52,7 @@ class LocalContrastTransform(ImageOnlyTransform, LocalTransform):
                     kernels.append(None)
                     contrasts.append(None)
                     continue
-                kernel = self._generate_kernel(spatial).astype(np.float32)
+                kernel = self._generate_kernel(spatial)
                 contrast = sample_scalar(self.new_contrast)
                 kernels.append(kernel)
                 contrasts.append(contrast)
@@ -67,9 +67,11 @@ class LocalContrastTransform(ImageOnlyTransform, LocalTransform):
                 continue
 
             channel = img_np[c]
-            mean = (channel * kernel).sum() / (kernel.sum() + 1e-8)
-            modified = (channel - mean) * contrast + mean
-            img_np[c] = self.run_interpolation(channel, modified, kernel)
+            mean = (channel * kernel).sum() / (kernel.sum() + 1e-8)  # keep exact summation order
+            modified = channel - mean          # fresh owned array
+            modified *= contrast
+            modified += mean                   # (channel - mean) * contrast + mean
+            img_np[c] = self.run_interpolation_inplace(channel, modified, kernel)
 
         return torch.from_numpy(img_np).to(img.device, dtype=img.dtype)
 

{batchgeneratorsv2-0.3.3 → batchgeneratorsv2-0.3.5}/batchgeneratorsv2/transforms/local/local_gamma.py

@@ -42,7 +42,7 @@ class LocalGammaTransform(ImageOnlyTransform, LocalTransform):
             return {'kernels': [None] * C, 'gammas': [None] * C}
 
         if self.same_for_all_channels:
-            kernel = self._generate_kernel(spatial).astype(np.float32)
+            kernel = self._generate_kernel(spatial)
             gamma = sample_scalar(self.gamma)
 
             kernels = [kernel if apply else None for apply in apply_channel]
@@ -54,7 +54,7 @@ class LocalGammaTransform(ImageOnlyTransform, LocalTransform):
                     kernels.append(None)
                     gammas.append(None)
                     continue
-                kernel = self._generate_kernel(spatial).astype(np.float32)
+                kernel = self._generate_kernel(spatial)
                 gamma = sample_scalar(self.gamma)
                 kernels.append(kernel)
                 gammas.append(gamma)
@@ -72,15 +72,15 @@ class LocalGammaTransform(ImageOnlyTransform, LocalTransform):
             min_val, max_val = channel.min(), channel.max()
             denom = max(max_val - min_val, 1e-8)
 
-            # Normalize to [0, 1]
-            norm = (channel - min_val) / denom
+            # Normalize to [0, 1] (norm is a fresh owned array; channel/kernel are never mutated)
+            norm = channel - min_val
+            norm /= denom
             gamma_corrected = np.power(norm, gamma)
 
-            # Blend using kernel
-            blended = self.run_interpolation(norm, gamma_corrected, kernel)
-
-            # Rescale to original range
-            img_np[c] = blended * denom + min_val
+            blended = self.run_interpolation_inplace(norm, gamma_corrected, kernel)
+            blended *= denom            # rescale to original range
+            blended += min_val
+            img_np[c] = blended
 
         return torch.from_numpy(img_np).to(img.device, dtype=img.dtype)
 

{batchgeneratorsv2-0.3.3 → batchgeneratorsv2-0.3.5}/batchgeneratorsv2/transforms/local/local_smoothing.py

@@ -47,7 +47,7 @@ class LocalSmoothingTransform(ImageOnlyTransform, LocalTransform):
         sigma = sample_scalar(self.kernel_size)
 
         if self.same_for_all_channels:
-            kernel = self._generate_kernel(spatial).astype(np.float32)
+            kernel = self._generate_kernel(spatial)
             strength = sample_scalar(self.smoothing_strength)
 
             kernels = [kernel if apply else None for apply in apply_channel]
@@ -59,7 +59,7 @@ class LocalSmoothingTransform(ImageOnlyTransform, LocalTransform):
                     kernels.append(None)
                     strengths.append(None)
                     continue
-                kernel = self._generate_kernel(spatial).astype(np.float32)
+                kernel = self._generate_kernel(spatial)
                 strength = sample_scalar(self.smoothing_strength)
                 kernels.append(kernel)
                 strengths.append(strength)
@@ -74,9 +74,9 @@ class LocalSmoothingTransform(ImageOnlyTransform, LocalTransform):
             if kernel is None:
                 continue
 
-            kernel = kernel * strength  # scale kernel by smoothing strength
+            kernel = kernel * strength  # fresh owned array (shared param kernel is never mutated)
             smoothed = gaussian_filter(img_np[c], sigma=sigma)
-            img_np[c] = self.run_interpolation(img_np[c], smoothed, kernel)
+            img_np[c] = self.run_interpolation_inplace(img_np[c], smoothed, kernel)
 
         return torch.from_numpy(img_np).to(img.device, dtype=img.dtype)
 

{batchgeneratorsv2-0.3.3 → batchgeneratorsv2-0.3.5}/batchgeneratorsv2/transforms/local/local_transform.py

@@ -11,6 +11,9 @@ class LocalTransform(ABC):
         self.scale = scale
 
     def _generate_kernel(self, img_shp: Tuple[int, ...]) -> np.ndarray:
+        """
+        Always returns a float32 array in [0, 1]; callers rely on this and do not re-cast.
+        """
         ndim = len(img_shp)
         x_grids = [np.arange(-0.5, s + 0.5, dtype=np.float32) for s in img_shp]
         kernels = []
@@ -75,6 +78,8 @@ class LocalTransform(ABC):
         """
         Blends original and modified images using the given kernel as a per-pixel weight map.
 
+        Reference implementation; prefer run_interpolation_inplace in per-sample code paths.
+
         Parameters:
             original_image (np.ndarray): Unmodified input image
             modified_image (np.ndarray): Modified version (e.g., gamma-corrected)
@@ -83,4 +88,25 @@ class LocalTransform(ABC):
         Returns:
             np.ndarray: Blended result
         """
-        return original_image * (1.0 - kernel_image) + modified_image * kernel_image
+        return original_image * (1.0 - kernel_image) + modified_image * kernel_image
+
+    @staticmethod
+    def run_interpolation_inplace(original_image: np.ndarray,
+                                  modified_image: np.ndarray,
+                                  kernel_image: np.ndarray) -> np.ndarray:
+        """
+        Same blend as run_interpolation, computed into modified_image's buffer with a single extra
+        temporary. Uses the identical individual * and + ops with only IEEE-commutative operand swaps
+        (no reassociation), so the result is bit-identical to run_interpolation.
+
+        Contract: modified_image must be a fresh buffer owned by the caller (it is overwritten);
+        original_image and kernel_image are only read. original_image may alias the eventual output
+        slot as long as that slot is written only after this call returns.
+
+        Returns modified_image (now holding the blended result).
+        """
+        tmp = 1.0 - kernel_image
+        tmp *= original_image       # original * (1 - kernel)
+        modified_image *= kernel_image
+        modified_image += tmp
+        return modified_image

{batchgeneratorsv2-0.3.3 → batchgeneratorsv2-0.3.5}/batchgeneratorsv2/transforms/nnunet/seg_to_onehot.py

@@ -22,10 +22,17 @@ class MoveSegAsOneHotToDataTransform(BasicTransform):
 
     def apply(self, data_dict, **params):
         seg = data_dict['segmentation'][self.source_channel_idx]
-        seg_onehot = torch.zeros((len(self.all_labels), *seg.shape), dtype=data_dict['image'].dtype)
+        img = data_dict['image']
+        C = img.shape[0]
+        # Allocate the concatenated result once and write the one-hot channels straight into their final slots,
+        # instead of building a separate seg_onehot tensor and then torch.cat-copying everything again. Only the
+        # one-hot tail needs zeroing; the image channels are fully overwritten anyway.
+        out = torch.empty((C + len(self.all_labels), *seg.shape), dtype=img.dtype)
+        out[:C] = img
+        out[C:].zero_()
         for i, l in enumerate(self.all_labels):
-            seg_onehot[i][seg == l] = 1
-        data_dict['image'] = torch.cat((data_dict['image'], seg_onehot))
+            out[C + i][seg == l] = 1
+        data_dict['image'] = out
         if self.remove_channel_from_source:
             remaining_channels = [i for i in range(data_dict['segmentation'].shape[0]) if i != self.source_channel_idx]
             data_dict['segmentation'] = data_dict['segmentation'][remaining_channels]

{batchgeneratorsv2-0.3.3 → batchgeneratorsv2-0.3.5}/batchgeneratorsv2/transforms/noise/gaussian_blur.py

@@ -20,7 +20,8 @@ def blur_dimension(img: torch.Tensor, sigma: float, dim_to_blur: int, force_use_
                 where C is the channel dimension and X, Y, Z are spatial dimensions.
     :param sigma: The standard deviation of the Gaussian kernel.
     :param dim_to_blur: The dimension along which to apply the Gaussian blur (0 for X, 1 for Y, 2 for Z).
-    :return: The blurred image tensor.
+    :return: The blurred image tensor. `img` is never modified; a new tensor is returned (callers such as
+             GaussianBlurTransform chain these calls and rely on this non-mutation contract).
     """
     assert img.ndim - 1 > dim_to_blur, "dim_to_blur must be a valid spatial dimension of the input image."
     # Adjustments for kernel based on image dimensions
@@ -122,22 +123,28 @@ class GaussianBlurTransform(ImageOnlyTransform):
         # print(params['sigmas'])
         if self.synchronize_channels:
             # we can compute that in one go as the conv implementation supports arbitrary input channels (with expanded kernel)
+            # Gather the selected channels once (advanced indexing copies), chain the per-dim blurs, scatter back once.
+            # blur_dimension never mutates its input, so this is the same composition as writing back each dim.
+            mask = params['apply_to_channel']
+            sub = img[mask]
             for d in range(dim):
-                # print(d, params['sigmas'][d])
                 if not self.benchmark:
-                    img[params['apply_to_channel']] = blur_dimension(img[params['apply_to_channel']], params['sigmas'][d], d)
+                    sub = blur_dimension(sub, params['sigmas'][d], d)
                 else:
-                    img[params['apply_to_channel']] = self._benchmark_wrapper(img[params['apply_to_channel']], params['sigmas'][d], d)
+                    sub = self._benchmark_wrapper(sub, params['sigmas'][d], d)
+            img[mask] = sub
         else:
             # we have to go through all the channels, build the kernel for each channel etc
             idx = np.where(params['apply_to_channel'])[0]
             for j, i in enumerate(idx):
+                sub = img[i:i+1]  # view for the first read; blur_dimension returns fresh tensors thereafter
                 for d in range(dim):
                     # print(i, d, params['sigmas'][i][d])
                     if not self.benchmark:
-                        img[i:i+1] = blur_dimension(img[i:i+1], params['sigmas'][j][d], d)
+                        sub = blur_dimension(sub, params['sigmas'][j][d], d)
                     else:
-                        img[i:i+1] = self._benchmark_wrapper(img[i:i+1], params['sigmas'][j][d], d)
+                        sub = self._benchmark_wrapper(sub, params['sigmas'][j][d], d)
+                img[i:i+1] = sub
         return img
 
     def _benchmark_wrapper(self, img: torch.Tensor, sigma: float, dim_to_blur: int):

{batchgeneratorsv2-0.3.3 → batchgeneratorsv2-0.3.5}/batchgeneratorsv2/transforms/noise/rician.py

@@ -1,3 +1,4 @@
+import numpy as np
 import torch
 from typing import Tuple
 from batchgeneratorsv2.transforms.base.basic_transform import ImageOnlyTransform
@@ -25,19 +26,21 @@ class RicianNoiseTransform(ImageOnlyTransform):
         noise_imag = torch.empty_like(img).normal_(mean=0.0, std=var)
 
         min_val = img.min()
-        shifted = img - min_val
+        shifted = img - min_val  # fresh owned buffer (img itself is never mutated below)
 
-        rician = torch.sqrt((shifted + noise_real).pow_(2).add_(noise_imag.pow_(2)))
-        rician = rician + min_val
+        # Run the whole Rician pipeline in `shifted`'s buffer. Same elementwise ops in the same order as
+        # sqrt((shifted+noise_real)^2 + noise_imag^2) + min_val, just in place -> bit-identical.
+        rician = shifted.add_(noise_real).pow_(2).add_(noise_imag.pow_(2)).sqrt_()
+        rician.add_(min_val)
 
-        # Normalize to match original mean and std
+        # Normalize to match original mean and std (scalars captured before the in-place normalization)
         input_mean, input_std = img.mean(), img.std()
         rician_mean, rician_std = rician.mean(), rician.std()
 
         if rician_std > 0:
-            rician = (rician - rician_mean) / rician_std * input_std + input_mean
+            rician.sub_(rician_mean).div_(rician_std).mul_(input_std).add_(input_mean)
         else:
-            rician = rician * 0 + input_mean  # fallback if std is zero (flat image)
+            rician.mul_(0).add_(input_mean)  # keep mul_(0) (not fill_) so inf/nan propagate as before
 
         return rician
 

{batchgeneratorsv2-0.3.3 → batchgeneratorsv2-0.3.5}/batchgeneratorsv2/transforms/noise/sharpen.py

@@ -81,7 +81,7 @@ class SharpeningTransform(ImageOnlyTransform):
                 continue
 
             if clamp:
-                mn, mx = torch.min(img[c]), torch.max(img[c])
+                mn, mx = torch.aminmax(img[c])  # single reduction pass instead of separate min()/max()
 
             x = img[c].unsqueeze(0).unsqueeze(0)  # (1, 1, H, W) or (1, 1, D, H, W)
             padded = F.pad(x, pad, mode='replicate')
@@ -91,8 +91,9 @@ class SharpeningTransform(ImageOnlyTransform):
             else:
                 laplace = F.conv3d(padded, kernel)
 
-            sharpened = x + strength * laplace
-            out[c] = sharpened.squeeze()
+            # laplace is a fresh conv output (owned, not reused). Combine in place: strength*laplace + x == x +
+            # strength*laplace (IEEE add commutes). x is read here before out[c] (== img[c]) is written -> no hazard.
+            out[c] = laplace.mul_(strength).add_(x).squeeze()
 
             if clamp:
                 out[c].clamp_(mn, mx)

{batchgeneratorsv2-0.3.3 → batchgeneratorsv2-0.3.5}/batchgeneratorsv2/transforms/spatial/channel_misalignment.py

@@ -201,24 +201,74 @@ class ChannelMisalignmentTransform(ImageOnlyTransform):
             if params['elastic_offsets'] is not None:
                 grid += params['elastic_offsets']
             if params['affine'] is not None:
-                grid = torch.matmul(grid, torch.from_numpy(params['affine']).float())
-
-            # we center the grid around the center_location_in_pixels. We should center the mean of the grid, not the center position
-            # only do this if we elastic deform
+                # grid stores spatial vectors as row vectors (shape [..., dim]), whereas affine is built to act on
+                # column vectors (x' = affine @ x). We therefore multiply by affine.T so that each grid point is
+                # transformed by affine rather than its transpose (see issue #24).
+                grid = torch.matmul(grid, torch.from_numpy(params['affine'].T).float())
+
+            # re-center the mean of the grid (not the center position). new_center is all-zeros here (translation is
+            # handled via the shifted crop center below, not via the grid), so in the non-deform case mn is also 0 and
+            # `grid += (0 - 0)` would be a wasted full-grid pass -> skipped. This skip is only valid while new_center
+            # stays all-zeros; if a nonzero new_center is ever reintroduced, the shift must apply in both paths again.
             if params['elastic_offsets'] is not None:
                 mn = grid.mean(dim=list(range(img.ndim - 1)))
-            else:
-                mn = 0
+                # new_center = torch.Tensor([c - s / 2 for c, s in zip(params['center_location_in_pixels'], img.shape[1:])])
+                new_center = torch.zeros(len(im_shape))
+                grid += (new_center - mn)
 
-            # new_center = torch.Tensor([c - s / 2 for c, s in zip(params['center_location_in_pixels'], img.shape[1:])])
-            new_center = torch.Tensor([0, 0, 0])
-            grid += (new_center - mn)
+            # compute the grid_sample grid once instead of once per channel
+            grid_sample_grid = _convert_my_grid_to_grid_sample_grid(grid, img.shape[1:])[None]
 
             for ch in self.im_channels_2_misalign:
-                img[ch, ...] = grid_sample(img[ch, ...].unsqueeze(0).unsqueeze(0),
-                                           _convert_my_grid_to_grid_sample_grid(grid, img.shape[1:])[None],
-                                           mode='bilinear', padding_mode="zeros", align_corners=False)[0]
-                img[ch, ...] = crop_tensor(img[ch, ...].unsqueeze(0),
+                # crop the grid_sample result directly instead of writing it into img[ch] and reading it back
+                resampled = grid_sample(img[ch, ...].unsqueeze(0).unsqueeze(0),
+                                        grid_sample_grid,
+                                        mode='bilinear', padding_mode="zeros", align_corners=False)[0]
+                img[ch, ...] = crop_tensor(resampled,
                                            [math.floor(i) for i in params['center_location_in_pixels']], im_shape,
                                            pad_mode='constant', pad_kwargs={'value': 0})
             return img
+
+
+# from copy import deepcopy
+# import SimpleITK as sitk
+# from pathlib import Path
+
+# if __name__ == '__main__':
+
+#     source_folder = Path('./test_images/')
+
+#     img0 = sitk.GetArrayFromImage(sitk.ReadImage(source_folder / 'case001_0000.nii.gz'))
+#     img1 = sitk.GetArrayFromImage(sitk.ReadImage(source_folder / 'case001_0001.nii.gz'))
+#     img2 = sitk.GetArrayFromImage(sitk.ReadImage(source_folder / 'case001_0002.nii.gz'))
+#     seg = sitk.GetArrayFromImage(sitk.ReadImage(source_folder / 'case001.nii.gz'))
+
+#     sp = ChannelMisalignmentTransform(
+#         im_channels_2_misalign=(0,),
+
+#         squeezing_zyx=(0.0, 0.1, 0.0),
+#         p_squeeze=1.0,
+
+#         rotation_ax_cor_sag=(30, 0, 0),
+#         rad_or_deg="deg",
+#         p_rotation=1.0,
+
+#         shift_zyx=(5, 100, 20),
+#         p_shift = 1.0,
+
+#     )
+
+#     data_dict = {'image': torch.cat([torch.from_numpy(deepcopy(img0[None])).float(),
+#                                      torch.from_numpy(deepcopy(img1[None])).float(),
+#                                      torch.from_numpy(deepcopy(img2[None])).float()], dim=0)}
+#     if seg is not None:
+#         data_dict['segmentation'] = torch.from_numpy(deepcopy(seg[None]))
+
+#     out = sp(**data_dict)
+
+#     sitk.WriteImage(sitk.GetImageFromArray(img0), source_folder / 'original_ch0.nii.gz')
+#     sitk.WriteImage(sitk.GetImageFromArray(img1), source_folder / 'original_ch1.nii.gz')
+#     sitk.WriteImage(sitk.GetImageFromArray(img2), source_folder / 'original_ch2.nii.gz')
+#     sitk.WriteImage(sitk.GetImageFromArray(out['image'][0].numpy()), source_folder / 'transformed_ch0.nii.gz')
+#     sitk.WriteImage(sitk.GetImageFromArray(out['image'][1].numpy()), source_folder / 'transformed_ch1.nii.gz')
+#     sitk.WriteImage(sitk.GetImageFromArray(out['image'][2].numpy()), source_folder / 'transformed_ch2.nii.gz')

{batchgeneratorsv2-0.3.3 → batchgeneratorsv2-0.3.5}/batchgeneratorsv2/transforms/spatial/spatial.py

@@ -219,7 +219,10 @@ class SpatialTransform(BasicTransform):
             if offsets is not None:
                 grid += offsets
             if affine is not None:
-                grid = torch.matmul(grid, torch.from_numpy(affine).float())
+                # grid stores spatial vectors as row vectors (shape [..., dim]), whereas affine is built to act on
+                # column vectors (x' = affine @ x). We therefore multiply by affine.T so that each grid point is
+                # transformed by affine rather than its transpose (see issue #24).
+                grid = torch.matmul(grid, torch.from_numpy(affine.T).float())
 
             # we center the grid around the center_location_in_pixels. We should center the mean of the grid, not the center position
             # only do this if we elastic deform
@@ -300,6 +303,11 @@ class SpatialTransform(BasicTransform):
             if self.bg_style_seg_sampling:
                 for c in range(segmentation.shape[0]):
                     labels = torch.from_numpy(np.sort(pd.unique(segmentation[c].numpy().ravel())))
+                    # result_seg is zero-initialized, so when the lowest label is the 0 background we never have to
+                    # write it: any voxel it would claim is already 0, and being the lowest label it can only be
+                    # overwritten by (never overwrite) the foreground labels processed after it. Skipping it saves a
+                    # full grid_sample per channel. Output is bit-identical.
+                    bg_is_zero = bool(labels[0] == 0)
                     # if we only have 2 labels then we can save compute time
                     if len(labels) == 2:
                         out = grid_sample(
@@ -310,9 +318,10 @@ class SpatialTransform(BasicTransform):
                             align_corners=self.align_corners
                         )[0][0] >= 0.5
                         result_seg[c][out] = labels[1]
-                        result_seg[c][~out] = labels[0]
+                        if not bg_is_zero:
+                            result_seg[c][~out] = labels[0]
                     else:
-                        for i, u in enumerate(labels):
+                        for u in (labels[1:] if bg_is_zero else labels):
                             result_seg[c][
                                 grid_sample(
                                     ((segmentation[c] == u).float())[None, None],
@@ -322,26 +331,35 @@ class SpatialTransform(BasicTransform):
                                     align_corners=self.align_corners
                                 )[0][0] >= 0.5] = u
             else:
+                # Per-voxel argmax over the interpolated one-hot label channels, computed incrementally so the
+                # (num_labels, *patch_size) stack is never materialized: only a running best value and the winning
+                # label (written straight into result_seg) are kept. This equals threshold-then-argmax sampling
+                # because the interpolated channels form a convex combination (they sum to <= scale_factor per
+                # voxel, so at most one label can exceed 0.7 * scale_factor). Load-bearing for bit-identical
+                # results: the running comparison is done in float16, and the strict > keeps the lowest label
+                # index on ties, exactly like torch.argmax(0) over a float16 stack.
+                scale_factor = 1000
                 for c in range(segmentation.shape[0]):
                     labels = torch.from_numpy(np.sort(pd.unique(segmentation[c].numpy().ravel())))
-                    # torch.where(torch.bincount(segmentation.ravel()) > 0)[0].to(segmentation.dtype)
-                    tmp = torch.zeros((len(labels), *self.patch_size), dtype=torch.float16)
-                    scale_factor = 1000
-                    done_mask = torch.zeros(*self.patch_size, dtype=torch.bool)
-                    for i, u in enumerate(labels):
-                        tmp[i] = grid_sample(
-                            ((segmentation[c] == u).float() * scale_factor)[None, None],
+                    best_val = None
+                    for u in labels:
+                        onehot = (segmentation[c] == u).float()
+                        onehot *= scale_factor
+                        cur = grid_sample(
+                            onehot[None, None],
                             grid[None],
                             mode=self.mode_seg,
                             padding_mode=grid_sample_padding_mode,
                             align_corners=self.align_corners
-                        )[0][0]
-                        mask = tmp[i] > (0.7 * scale_factor)
-                        result_seg[c][mask] = u
-                        done_mask = done_mask | mask
-                    if not torch.all(done_mask):
-                        result_seg[c][~done_mask] = labels[tmp[:, ~done_mask].argmax(0)]
-                    del tmp
+                        )[0][0].to(torch.float16)
+                        if best_val is None:
+                            best_val = cur
+                            if u != 0:  # result_seg is zero-initialized, so filling with 0 would be a no-op
+                                result_seg[c] = u
+                        else:
+                            better = cur > best_val
+                            torch.maximum(best_val, cur, out=best_val)
+                            result_seg[c][better] = u
 
         if self._requires_constant_padding_fixup(self.border_mode_seg, self.padding_value_seg):
             out_of_bounds_mask = self._compute_out_of_bounds_mask(grid, segmentation.shape[1:])

{batchgeneratorsv2-0.3.3 → batchgeneratorsv2-0.3.5}/batchgeneratorsv2/transforms/utils/deep_supervision_downsampling.py

@@ -12,6 +12,8 @@ class DownsampleSegForDSTransform(SegOnlyTransform):
 
     def _apply_to_segmentation(self, segmentation: torch.Tensor, **params) -> List[torch.Tensor]:
         results = []
+        dtype = segmentation.dtype
+        seg_float = None  # full-res float copy, built once and shared across scales (interpolate only reads it)
         for s in self.ds_scales:
             if not isinstance(s, (tuple, list)):
                 s = [s] * (segmentation.ndim - 1)
@@ -22,9 +24,10 @@ class DownsampleSegForDSTransform(SegOnlyTransform):
                 results.append(segmentation)
             else:
                 new_shape = [round(i * j) for i, j in zip(segmentation.shape[1:], s)]
-                dtype = segmentation.dtype
                 # interpolate is not defined for short etc
-                results.append(interpolate(segmentation[None].float(), new_shape, mode='nearest-exact')[0].to(dtype))
+                if seg_float is None:
+                    seg_float = segmentation[None].float()
+                results.append(interpolate(seg_float, new_shape, mode='nearest-exact')[0].to(dtype))
         return results
 
 

{batchgeneratorsv2-0.3.3 → batchgeneratorsv2-0.3.5}/batchgeneratorsv2.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: batchgeneratorsv2
-Version: 0.3.3
+Version: 0.3.5
 Summary: Batchgenerators but better
 Author: Helmholtz Imaging Applied Computer Vision Lab
 Author-email: Fabian Isensee <f.isensee@dkfz-heidelberg.de>

{batchgeneratorsv2-0.3.3 → batchgeneratorsv2-0.3.5}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "batchgeneratorsv2"
-version = "0.3.3"
+version = "0.3.5"
 requires-python = ">=3.9"
 description = "Batchgenerators but better"
 readme = "readme.md"