batchgeneratorsv2 0.3.2__py3-none-any.whl

batchgeneratorsv2/transforms/utils/compose.py

@@ -0,0 +1,89 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+from time import perf_counter
+from typing import List, Optional, Dict, Any
+
+import numpy as np
+
+from batchgeneratorsv2.transforms.base.basic_transform import BasicTransform
+
+
+class ComposeTransforms(BasicTransform):
+    def __init__(self, transforms: List[BasicTransform]):
+        super().__init__()
+        self.transforms = transforms
+
+    def apply(self, data_dict, **params):
+        for t in self.transforms:
+            data_dict = t(**data_dict)
+        return data_dict
+
+
+@dataclass
+class _TimingStats:
+    total_s: float = 0.0
+    n: int = 0
+
+    def add(self, dt: float) -> None:
+        self.total_s += dt
+        self.n += 1
+
+    @property
+    def mean_s(self) -> float:
+        return self.total_s / self.n if self.n > 0 else 0.0
+
+
+class TimedComposeTransforms(BasicTransform):
+    """
+    ComposeTransforms variant that measures per-transform wall-clock time and prints
+    average times after every `print_every` calls to `apply`.
+
+    Notes:
+    - Measures wall clock time via perf_counter.
+    - For CPU-only pipelines, this is representative. For GPU, you'd need synchronization.
+    """
+
+    def __init__(self, transforms: List[BasicTransform], print_every: int = 100, name: Optional[str] = None, p_write: float = 1.0):
+        super().__init__()
+        if print_every <= 0:
+            raise ValueError("print_every must be >= 1")
+        self.transforms = transforms
+        self.print_every = int(print_every)
+        self.name = name or self.__class__.__name__
+
+        self._iter = 0
+        self.p_write = p_write
+        self._stats: Dict[int, _TimingStats] = {i: _TimingStats() for i in range(len(transforms))}
+
+    def reset_timings(self) -> None:
+        """Reset accumulated timing statistics and iteration counter."""
+        self._iter = 0
+        for s in self._stats.values():
+            s.total_s = 0.0
+            s.n = 0
+
+    def _transform_display_name(self, t: BasicTransform) -> str:
+        # Prefer explicit "name" attribute if present, otherwise class name
+        return getattr(t, "name", None) or t.__class__.__name__
+
+    def _print_report(self) -> None:
+        lines = [f"[{self.name}] Average transform times over last {self._iter} iterations:"]
+        # Print in pipeline order
+        for i, t in enumerate(self.transforms):
+            st = self._stats[i]
+            lines.append(f"  {i:02d}  {self._transform_display_name(t)}: {st.mean_s * 1e3:.3f} ms")
+        print("\n".join(lines), flush=True)
+
+    def apply(self, data_dict: Dict[str, Any], **params) -> Dict[str, Any]:
+        for i, t in enumerate(self.transforms):
+            t0 = perf_counter()
+            data_dict = t(**data_dict)
+            dt = perf_counter() - t0
+            self._stats[i].add(dt)
+
+        self._iter += 1
+        if self._iter % self.print_every == 0 and np.random.uniform() < self.p_write:
+            self._print_report()
+
+        return data_dict

batchgeneratorsv2/transforms/utils/cropping.py

@@ -0,0 +1,73 @@
+import torch
+
+
+def crop_tensor(input_tensor, center, crop_size, pad_mode='constant', pad_kwargs=None):
+    """
+    Crops and pads an input tensor based on the specified center and crop size. Padding can be customized.
+
+    Parameters:
+    - input_tensor (torch.Tensor): The input tensor with shape (c, x, y) or (c, x, y, z).
+    - center (tuple): The center coordinates of the crop (x, y) or (x, y, z).
+    - crop_size (tuple): The size of the crop (width, height) or (width, height, depth).
+    - pad_mode (str): The mode to use for padding (see torch.nn.functional.pad documentation).
+    - pad_kwargs (dict, optional): Additional keyword arguments for padding.
+
+    Returns:
+    - torch.Tensor: The cropped and possibly padded tensor.
+    """
+    if pad_kwargs is None:
+        pad_kwargs = {'value': 0}
+
+    # Calculate dimensions
+    dim = len(center)  # Spatial dimensions
+    assert len(crop_size) == dim, "Crop size and center must have the same number of dimensions"
+    assert input_tensor.ndim - 1 == dim, "Crop size and input_tensor must have the same number of spatial dimensions"
+
+    spatial_shape = input_tensor.shape[-dim:]
+    start = [max(0, cen - cs // 2) for cen, cs in zip(center, crop_size)]
+    end = [min(sh, st + cs) for sh, st, cs in zip(spatial_shape, start, crop_size)]
+
+    # Calculate padding
+    padding_needed = [(cs - (e - s)) for cs, s, e in zip(crop_size, start, end)]
+    pad_before = [max(0, - (cen - cs // 2)) for cen, cs in zip(center, crop_size)]
+    pad_after = [pn - pb for pn, pb in zip(padding_needed, pad_before)]
+
+    # Adjust start and end for the case where the crop is entirely outside the input tensor
+    start = [min(max(0, s), sh) for s, sh in zip(start, spatial_shape)]
+    end = [max(min(e, sh), 0) for e, sh in zip(end, spatial_shape)]
+
+    # Perform crop
+    slices = [slice(None)] + [slice(s, e) for s, e in zip(start, end)]
+    cropped = input_tensor[tuple(slices)]
+
+    # Pad
+    pad_width = sum([[b, a] for b, a in zip(pad_before[::-1], pad_after[::-1])], [])
+    if any(pad_width):
+        cropped = torch.nn.functional.pad(cropped, pad_width, mode=pad_mode, **pad_kwargs)
+
+    return cropped
+
+
+def center_crop(input_tensor, crop_size, pad_mode='constant', pad_kwargs=None):
+    """
+    Performs a center crop on the input tensor. If the crop extends beyond the borders of the tensor,
+    it will be padded according to the specified pad_mode and pad_kwargs.
+
+    Parameters:
+    - input_tensor (torch.Tensor): The input tensor with shape (c, x, y) or (c, x, y, z).
+    - crop_size (tuple): The size of the crop (width, height) or (width, height, depth).
+    - pad_mode (str): The mode to use for padding (see torch.nn.functional.pad documentation).
+    - pad_kwargs (dict, optional): Additional keyword arguments for padding.
+
+    Returns:
+    - torch.Tensor: The center-cropped and possibly padded tensor.
+    """
+    dim = len(input_tensor.shape) - 1  # Number of spatial dimensions (2 or 3)
+    spatial_shape = input_tensor.shape[-dim:]  # Spatial dimensions of the input tensor
+
+    # Calculate the center of the input tensor
+    center = tuple(s // 2 for s in spatial_shape)
+
+    # Use the previously defined function for cropping and padding
+    return crop_tensor(input_tensor, center, crop_size, pad_mode, pad_kwargs)
+

batchgeneratorsv2/transforms/utils/deep_supervision_downsampling.py

@@ -0,0 +1,59 @@
+from typing import Tuple, List, Union
+import torch
+
+from batchgeneratorsv2.transforms.base.basic_transform import SegOnlyTransform
+from torch.nn.functional import interpolate
+
+
+class DownsampleSegForDSTransform(SegOnlyTransform):
+    def __init__(self, ds_scales: Union[List, Tuple]):
+        super().__init__()
+        self.ds_scales = ds_scales
+
+    def _apply_to_segmentation(self, segmentation: torch.Tensor, **params) -> List[torch.Tensor]:
+        results = []
+        for s in self.ds_scales:
+            if not isinstance(s, (tuple, list)):
+                s = [s] * (segmentation.ndim - 1)
+            else:
+                assert len(s) == segmentation.ndim - 1
+
+            if all([i == 1 for i in s]):
+                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))
+        return results
+
+
+if __name__ == '__main__':
+    from time import time
+    import numpy as np
+    import os
+
+    os.environ['OMP_NUM_THREADS'] = '1'
+    torch.set_num_threads(1)
+
+    mbt = DownsampleSegForDSTransform((1, 0.5, 0.25))
+
+    times_torch = []
+    for _ in range(1):
+        data_dict = {'segmentation': torch.round(5 * torch.rand((2, 128, 192, 64)), decimals=0).to(torch.uint8)}
+        st = time()
+        out = mbt(**data_dict)
+        times_torch.append(time() - st)
+    print('torch', np.mean(times_torch))
+
+    from nnunetv2.training.data_augmentation.custom_transforms.deep_supervision_donwsampling import \
+        DownsampleSegForDSTransform2
+
+    gnt_bg = DownsampleSegForDSTransform2((1, 0.5, 0.25), order=0)
+    times_bg = []
+    for _ in range(1):
+        data_dict = {'seg': np.round(5 * np.random.uniform(size=(1, 2, 128, 192, 64)), decimals=0).astype(np.uint8)}
+        st = time()
+        out = gnt_bg(**data_dict)
+        times_bg.append(time() - st)
+    print('bg', np.mean(times_bg))

batchgeneratorsv2/transforms/utils/move_channels.py

@@ -0,0 +1,52 @@
+from typing import Tuple, Union
+import torch
+from batchgeneratorsv2.transforms.base.basic_transform import BasicTransform
+
+
+class MoveChannelsTransform(BasicTransform):
+    def __init__(self, channel_ids: Union[int, Tuple[int, ...]], source_key: str, target_key: str):
+        super().__init__()
+        if isinstance(channel_ids, int):
+            channel_ids = (channel_ids,)
+        self.channel_ids: Tuple[int, ...] = tuple(channel_ids)
+        self.source_key = source_key
+        self.target_key = target_key
+
+    def get_parameters(self, **data_dict) -> dict:
+        return {}
+
+    def apply(self, data_dict, **params):
+        src = data_dict[self.source_key]
+        assert src.ndim in (3, 4), f"Expected (C,X,Y) or (C,X,Y,Z), got {src.shape}"
+
+        if self.target_key in data_dict:
+            tgt = data_dict[self.target_key]
+            assert src.ndim == tgt.ndim, "source and target key must have the same number of dimensions"
+            assert src.shape[1:] == tgt.shape[1:], (
+                f"spatial dimensions must match. Got source: {src.shape} and target: {tgt.shape}"
+            )
+        else:
+            tgt = None
+
+        C = src.shape[0]
+        idx = torch.as_tensor(self.channel_ids, device=src.device, dtype=torch.long)
+
+        keep = torch.ones(C, device=src.device, dtype=torch.bool)
+        keep[idx] = False
+
+        move = src[~keep]      # channels to move
+        src_new = src[keep]    # remaining channels
+
+        # attach moved channels to target
+        if tgt is None:
+            data_dict[self.target_key] = move
+        else:
+            data_dict[self.target_key] = torch.cat((tgt, move), dim=0)
+
+        # update or remove source
+        if src_new.shape[0] == 0:
+            del data_dict[self.source_key]
+        else:
+            data_dict[self.source_key] = src_new
+
+        return data_dict

batchgeneratorsv2/transforms/utils/nnunet_masking.py

@@ -0,0 +1,24 @@
+from typing import List
+
+from batchgeneratorsv2.transforms.base.basic_transform import BasicTransform
+
+
+class MaskImageTransform(BasicTransform):
+    def __init__(self,
+                 apply_to_channels: List[int],
+                 channel_idx_in_seg: int = 0,
+                 set_outside_to: float = 0,
+                 ):
+        super().__init__()
+        self.apply_to_channels = apply_to_channels
+        self.channel_idx_in_seg = channel_idx_in_seg
+        self.set_outside_to = set_outside_to
+
+    def apply(self, data_dict, **params):
+        if len(self.apply_to_channels) == 0:
+            return data_dict
+        mask = data_dict['segmentation'][self.channel_idx_in_seg] < 0
+        for a in self.apply_to_channels:
+            data_dict['image'][a, mask] = self.set_outside_to
+        return data_dict
+

batchgeneratorsv2/transforms/utils/pseudo2d.py

@@ -0,0 +1,79 @@
+import torch
+
+from batchgeneratorsv2.transforms.base.basic_transform import BasicTransform
+
+
+class Convert3DTo2DTransform(BasicTransform):
+    def apply(self, data_dict, **params):
+        if 'image' in data_dict.keys():
+            data_dict['nchannels_img'] = data_dict['image'].shape[0]
+        if 'segmentation' in data_dict.keys():
+            data_dict['nchannels_seg'] = data_dict['segmentation'].shape[0]
+        if 'regression_target' in data_dict.keys():
+            data_dict['nchannels_regr_trg'] = data_dict['regression_target'].shape[0]
+        return super().apply(data_dict, **params)
+
+    def _apply_to_image(self, img: torch.Tensor, **params) -> torch.Tensor:
+        shp = img.shape
+        return img.reshape((shp[0] * shp[1], *shp[2:]))
+
+    def _apply_to_regr_target(self, regression_target, **params) -> torch.Tensor:
+        return self._apply_to_image(regression_target, **params)
+
+    def _apply_to_segmentation(self, segmentation: torch.Tensor, **params) -> torch.Tensor:
+        return self._apply_to_image(segmentation, **params)
+
+    def _apply_to_bbox(self, bbox, **params):
+        raise NotImplementedError
+
+    def _apply_to_keypoints(self, keypoints, **params):
+        raise NotImplementedError
+
+
+class Convert2DTo3DTransform(BasicTransform):
+    def get_parameters(self, **data_dict) -> dict:
+        return {i: data_dict[i] for i in
+                ['nchannels_img', 'nchannels_seg', 'nchannels_regr_trg']
+                if i in data_dict.keys()}
+
+    def apply(self, data_dict, **params):
+        data_dict = super().apply(data_dict, **params)
+        if 'nchannels_img' in data_dict.keys():
+            del data_dict['nchannels_img']
+        if 'nchannels_seg' in data_dict.keys():
+            del data_dict['nchannels_seg']
+        if 'nchannels_regr_trg' in data_dict.keys():
+            del data_dict['nchannels_regr_trg']
+        return data_dict
+
+    def _apply_to_image(self, img: torch.Tensor, **params) -> torch.Tensor:
+        return img.reshape((params['nchannels_img'], img.shape[0] // params['nchannels_img'], *img.shape[1:]))
+
+    def _apply_to_segmentation(self, segmentation: torch.Tensor, **params) -> torch.Tensor:
+        return segmentation.reshape(
+            (params['nchannels_seg'], segmentation.shape[0] // params['nchannels_seg'], *segmentation.shape[1:]))
+
+    def _apply_to_regr_target(self, regression_target, **params) -> torch.Tensor:
+        return regression_target.reshape(
+            (params['nchannels_regr_trg'], regression_target.shape[0] // params['nchannels_regr_trg'], *regression_target.shape[1:]))
+
+    def _apply_to_bbox(self, bbox, **params):
+        raise NotImplementedError
+
+    def _apply_to_keypoints(self, keypoints, **params):
+        raise NotImplementedError
+
+
+if __name__ == '__main__':
+    d = torch.rand((2, 32, 64, 128))
+    s = torch.ones((1, 32, 64, 128))
+
+    fwd = Convert3DTo2DTransform()
+    bwd = Convert2DTo3DTransform()
+
+    inp = {'image': d, 'segmentation': s}
+
+    tmp = fwd(**inp)
+    print(tmp['image'].shape, tmp['segmentation'].shape)
+    out = bwd(**tmp)
+    print(out['image'].shape, out['segmentation'].shape)

batchgeneratorsv2/transforms/utils/random.py

@@ -0,0 +1,46 @@
+from typing import List
+
+import torch
+
+from batchgeneratorsv2.transforms.base.basic_transform import BasicTransform
+import numpy as np
+
+
+class RandomTransform(BasicTransform):
+    def __init__(self, transform: BasicTransform, apply_probability: float = 1):
+        super().__init__()
+        self.transform = transform
+        self.apply_probability = apply_probability
+
+    def get_parameters(self, **data_dict) -> dict:
+        return {"apply_transform": torch.rand(1).item() < self.apply_probability}
+
+    def apply(self, data_dict: dict, **params) -> dict:
+        if params['apply_transform']:
+            return self.transform(**data_dict)
+        else:
+            return data_dict
+    
+    def __repr__(self):
+        ret_str = f"{type(self).__name__}(p={self.apply_probability}, transform={self.transform})"
+        return ret_str
+
+
+class OneOfTransform(BasicTransform):
+    """
+    Randomly selects and applies one transform from the provided list.
+
+    Each transform must be a callable (usually a BasicTransform subclass).
+    This does not override the internal probabilities of the transforms themselves.
+
+    Args:
+        list_of_transforms (List[BasicTransform]): A list of transform instances to choose from.
+    """
+
+    def __init__(self, list_of_transforms: List[BasicTransform]):
+        super().__init__()
+        self.list_of_transforms = list_of_transforms
+
+    def __call__(self, **data_dict) -> dict:
+        chosen_transform = np.random.choice(self.list_of_transforms)
+        return chosen_transform(**data_dict)

batchgeneratorsv2/transforms/utils/remove_label.py

@@ -0,0 +1,27 @@
+from typing import Union, Tuple, List
+
+import torch
+
+from batchgeneratorsv2.transforms.base.basic_transform import SegOnlyTransform
+
+
+class RemoveLabelTransform(SegOnlyTransform):
+    def __init__(self, label_value: int, set_to: int, segmentation_channels: Union[int, Tuple[int, ...], List[int]] = None):
+        if not isinstance(segmentation_channels, (list, tuple)) and segmentation_channels is not None:
+            segmentation_channels = [segmentation_channels]
+        self.segmentation_channels = segmentation_channels
+        self.label_value = label_value
+        self.set_to = set_to
+        super().__init__()
+
+    def _apply_to_segmentation(self, segmentation: torch.Tensor, **params) -> torch.Tensor:
+        if self.segmentation_channels is None:
+            channels = list(range(segmentation.shape[0]))
+        else:
+            channels = self.segmentation_channels
+        for s in channels:
+            segmentation[s][segmentation[s] == self.label_value] = self.set_to
+        return segmentation
+
+# uff.
+RemoveLabelTansform = RemoveLabelTransform

batchgeneratorsv2/transforms/utils/seg_to_regions.py

@@ -0,0 +1,24 @@
+from typing import Union, List, Tuple
+import torch
+
+from batchgeneratorsv2.transforms.base.basic_transform import SegOnlyTransform
+
+
+class ConvertSegmentationToRegionsTransform(SegOnlyTransform):
+    def __init__(self, regions: Union[List[Union[int, List[int]]], Tuple[Union[int, Tuple[int, ...]]]], channel_in_seg: int = 0):
+        super().__init__()
+        self.regions = [torch.Tensor(i) if not isinstance(i, int) else torch.Tensor([i]) for i in regions]
+        self.channel_in_seg = channel_in_seg
+
+    def _apply_to_segmentation(self, segmentation: torch.Tensor, **params) -> torch.Tensor:
+        num_regions = len(self.regions)
+        region_output = torch.zeros((num_regions, *segmentation.shape[1:]), dtype=torch.bool, device=segmentation.device)
+        for region_id, region_labels in enumerate(self.regions):
+            if len(region_labels) == 1:
+                region_output[region_id] = segmentation[self.channel_in_seg] == region_labels
+            else:
+                region_output[region_id] = torch.isin(segmentation[self.channel_in_seg], region_labels)
+        # we return bool here and leave it to the loss function to cast it to whatever it needs. Transferring bool to
+        # device followed by cast on device should be faster than having fp32 here and transferring that
+        return region_output
+