Merge branch 'development_robert' of github.com:Hendrik-code/TPTBox into development_robert

Author: robert-graf

TPTBox/core/internal/elastic_deform.py

@@ -0,0 +1,139 @@
+import time
+
+import elasticdeform
+import numpy as np
+from numpy.typing import NDArray
+
+from TPTBox import NII
+
+
+def deformed_nii(
+    nii_dic: dict[str, NII],
+    sigma: float | None = None,
+    points=None,
+    deform_factor=1.0,
+    deform_padding=10,
+    normalize=True,
+    joint_normalize=False,
+) -> dict[str, NII]:
+    """
+    Deform a dictionary of NII objects using random grid deformation. Requires elasticdeform. 'pip install elasticdeform'
+
+    IMPORTANT: Normalize your image data to 0,1. The .seg property of NII shows if this is a segmentation. (NII is form our TPTBox and is a wrapper for nibable)
+
+    This function takes a dictionary of NII objects and applies random grid deformation to each object
+    using specified deformation parameters or, if not provided, random parameters generated based on
+    the `deform_factor`. The deformed objects are returned as a dictionary.
+
+    Args:
+        arr_dic (dict[str, NII]): A dictionary containing NII objects to be deformed.
+        sigma (float, optional): The standard deviation of the deformation field. If not provided,
+            it will be generated based on the `deform_factor`.
+        points (int, optional): The number of control points for the deformation grid. If not provided,
+            it will be generated based on the `deform_factor`.
+        deform_factor (float, optional): A factor used to determine the deformation parameters if
+            `sigma` and `points` are not specified. Larger values result in stronger deformations.
+        deform_padding (int, optional): The padding added to the deformed objects to avoid edge artifacts.
+        verbose (bool, optional): If True, enable verbose logging. Default is True.
+
+    Returns:
+        dict[str, NII]: A dictionary where keys correspond to the input dictionary keys, and values
+        correspond to the deformed NII objects.
+
+    Example:
+        # Deform a dictionary of NII objects using default deformation parameters
+        deformed_data = deformed_NII(arr_dic)
+
+        # Deform a dictionary of NII objects with specific deformation parameters
+        sigma = 1.0
+        points = 20
+        deformed_data = deformed_NII(arr_dic, sigma=sigma, points=points)
+    """
+    if sigma is None or points is None:
+        sigma, points = get_random_deform_parameter(deform_factor=deform_factor)
+
+    print("deformation parameter sigma = ", round(sigma, 4), "; n_points = ", points)
+    t = time.time()
+    values = list(nii_dic.values())
+    # Deform
+    if joint_normalize:
+        max_v = max([img.max() for img in nii_dic.values() if not img.seg])
+        nii_dic = {k: img if img.seg else img.set_dtype(np.float32) / max_v for k, img in nii_dic.items()}
+    elif normalize:
+        nii_dic = {k: img if img.seg else img.set_dtype(np.float32).normalize() for k, img in nii_dic.items()}
+    else:
+        nii_dic = {k: img if img.seg else img.set_dtype(np.float32) for k, img in nii_dic.items()}
+    assert sigma is not None
+    p = deform_padding
+    out: list[NDArray] = elasticdeform.deform_random_grid(
+        [pad(v.get_array(), p=p) for v in values],
+        sigma=sigma,  # type: ignore
+        points=points,
+        order=[0 if v.seg else 3 for v in values],  # type: ignore
+    )
+    out2: dict[str, NII] = {}
+    for (k, nii), arr in zip(nii_dic.items(), out, strict=True):
+        out2[k] = nii.set_array(arr[p:-p, p:-p, p:-p])
+    print("Deformation took", round(time.time() - t, 1), "Seconds")
+    return out2
+
+
+def pad(arr, p=10):
+    return np.pad(arr, p, mode="reflect")
+
+
+def get_random_deform_parameter(deform_factor: float = 1):
+    """
+    Generate random deformation parameters for use in 3D deformation.
+
+    This function generates random values for the deformation parameters, including 'sigma' and 'points',
+    based on the specified deformation factor. These parameters are used for 3D deformation operations.
+
+    Args:
+        deform_factor (float, optional): A factor to control the strength of deformation. Default is 1.
+
+    Returns:
+        tuple[float, int]: A tuple containing the generated 'sigma' (float) and 'points' (int) parameters.
+
+    Example:
+        # Generate random deformation parameters with a deformation factor of 1
+        sigma, points = get_random_deform_parameter()
+
+        # Generate random deformation parameters with a deformation factor of 2
+        sigma, points = get_random_deform_parameter(deform_factor=2)
+    """
+    sigma = 2 + np.random.uniform() * 2.5  # 1,5 - 4.5
+    min_points = 3
+    max_points = 17
+    if sigma < 2:
+        max_points = 17
+    elif sigma < 1.7:
+        max_points = 16
+    elif sigma < 2.1:
+        max_points = 15
+    elif sigma < 2.3:
+        max_points = 14
+    elif sigma < 2.5:
+        max_points = 13
+    elif sigma < 2.6:
+        max_points = 12
+    elif sigma < 2.7:
+        max_points = 11
+    elif sigma < 2.8:
+        max_points = 10
+    elif sigma < 3:
+        max_points = 9
+    elif sigma < 3.5:
+        max_points = 8
+    elif sigma < 4.0:
+        max_points = 7
+    elif sigma < 4.3:
+        max_points = 6
+    else:
+        max_points = 5
+    points = np.random.randint(max_points - min_points + 1) + min_points
+    # Stronger
+    sigma *= deform_factor
+    # points *= deform_factor
+    points = max(round(points), 1)
+    return (sigma, points)

TPTBox/core/nii_wrapper.py

@@ -4,6 +4,7 @@
 import traceback
 import warnings
 import zlib
+from collections import deque
 from collections.abc import Sequence
 from enum import Enum
 from math import ceil, floor
@@ -937,6 +938,34 @@ def resample_from_to(self, to_vox_map:Image_Reference|Has_Grid|tuple[SHAPE,AFFIN
         if isinstance(mapping,Has_Grid) and mapping.assert_affine(self,raise_error=False,origin_tolerance=0.000001,error_tolerance=0.000001,shape_tolerance=0):
             log.print(f"resample_from_to skipped; already in space: {self}",verbose=verbose)
             return self if inplace else self.copy()
+
+            #m1 = mapping.make_empty_POI().reorient(self.orientation)
+            #if m1.assert_affine(self,raise_error=False,origin_tolerance=0.000001,error_tolerance=0.000001,shape_tolerance=0):
+            #    log.print(f"resample_from_to only need reorientation; {self.orientation}",verbose=verbose)
+            #    return self.reorient(mapping.orientation,inplace=inplace)
+            #if self.orientation == mapping.orientation and self.zoom == mapping.zoom:
+            #    shift = (np.array(self.origin) - np.array(m1.origin)) / np.array(m1.zoom)
+            #    if np.allclose(shift, np.round(shift), atol=1e-6):
+            #        self = self.reorient(mapping.orientation,inplace=inplace)  # noqa: PLW0642
+            #        shift = (np.array(self.origin) - np.array(mapping.origin)) / np.array(mapping.zoom)
+            #        shift = np.round(shift).astype(int)
+            #        src_shape = np.array(mapping.shape)
+            #        dst_shape = np.array(self.shape)
+            #        # padding before = how much dst starts before src
+            #        pad_before = np.maximum(-shift, 0)
+            #
+            #        # where src ends inside dst
+            #        src_end_in_dst = shift + src_shape
+            #        # padding after = remaining dst size after src
+            #        pad_after = np.maximum(dst_shape - src_end_in_dst, 0)
+            #        pad = tuple((int(b), int(a)) for b, a in zip(pad_before, pad_after))
+            #        ret = self.apply_pad(pad, mode=mode)
+            #
+            #        log.print(f"resample_from_to only needs padding/cropping {pad}, ",verbose=verbose,)
+            #        ret.assert_affine(mapping,raise_error=False,origin_tolerance=0.000001,error_tolerance=0.000001,shape_tolerance=0)
+            #        return ret
+
+
         assert mapping is not None
         log.print(f"resample_from_to: {self} to {mapping}",verbose=verbose)
         if order is None:
@@ -1729,7 +1758,7 @@ def truncate_labels_beyond_reference(
     ):
         return self.truncate_labels_beyond_reference_(idx,not_beyond,fill,axis,inclusion,inplace=inplace)
 
-    def infect(self: NII, reference_mask: NII, inplace=False,verbose=True,axis:int|str|None=None,max_depth=None):
+    def infect(self: NII, reference_mask: NII, inplace=False,verbose=True,axis:int|str|None=None,max_depth=None, _do_crop=True):
         """
         Expands labels from self_mask into regions of reference_mask == 1 via breadth-first diffusion.
 
@@ -1742,8 +1771,14 @@ def infect(self: NII, reference_mask: NII, inplace=False,verbose=True,axis:int|s
             ndarray: Updated label mask.
         """
         self.assert_affine(reference_mask)
-        self_mask = self.compute_surface_mask().get_seg_array().copy()
-        self_mask_org = self.get_seg_array().copy()
+        if _do_crop:
+            crop = reference_mask.compute_crop(0,5)
+            s = self.apply_crop(crop)
+            reference_mask = reference_mask.apply_crop(crop)
+        else:
+            s = self
+        self_mask = s.compute_surface_mask().get_seg_array().copy()
+        self_mask_org = s.get_seg_array().copy()
         ref_mask = np.clip(reference_mask.get_seg_array(), 0, 1)
         ref_mask[self_mask_org != 0] = 0
         searched = np.clip(self_mask,0,1).astype(np.uint8)
@@ -1763,13 +1798,14 @@ def infect(self: NII, reference_mask: NII, inplace=False,verbose=True,axis:int|s
             else:
                 raise NotImplementedError(axis)
 
-        search = []
+        search = deque()
         coords = np.where(self_mask != 0)
         def _add_idx(x,y,z,v,d):
             for x1,y1,z1 in kernel:
                 a = x+x1
                 b = y+y1
                 c = z+z1
+
                 if a < 0 or b < 0 or c < 0:
                     continue
                 if a >= self_mask.shape[0] or b >= self_mask.shape[1] or c >= self_mask.shape[2]:
@@ -1782,28 +1818,37 @@ def _add_idx(x,y,z,v,d):
         def _infect(a,b,c,v,d):
             if d-1 == max_depth:
                 return
-            if searched[a,b,c] != 0:
+            if searched[x,y,z] != 0:
                 return
-            if ref_mask[a,b,c] == 0:
+            if ref_mask[x,y,z] == 0:
                 return
             #print(a,b,c)
             searched[a,b,c] = 1
             self_mask[a,b,c] = v
-            _add_idx(x,y,z,v,d)
+            _add_idx(a,b,c,v,d)
 
         from tqdm import tqdm
         for x,y,z in tqdm(zip(coords[0],coords[1],coords[2]),total=len(coords[0]),disable=not verbose,desc="Collecting Surface"):
             _add_idx(x,y,z,self_mask[x,y,z],0)
         while len(search) != 0:
             search2 = search
-            search = []
-            for x,y,z,v,d in tqdm(search2,disable=not verbose,desc="infect"):
+            search = deque()
+            for _ in tqdm(range(len(search2)),disable=not verbose,desc="infect"):
+                x,y,z,v,d = search2.popleft()
                 _infect(x,y,z,v,d+1)
         self_mask[self_mask == 0] = self_mask_org[self_mask == 0]
+        if _do_crop:
+            if inplace:
+                self[crop] = self_mask
+                return self
+            else:
+                arr = self.get_array()
+                arr[crop] = self_mask
+                self_mask = arr
         return self.set_array(self_mask,inplace=inplace)
 
-    def infect_(self: NII, reference_mask: NII,verbose=True,axis:int|str|None=None):
-        return self.infect(reference_mask, inplace=True,verbose=verbose,axis=axis)
+    def infect_(self: NII, reference_mask: NII,verbose=True,axis:int|str|None=None,_do_crop=True):
+        return self.infect(reference_mask, inplace=True,verbose=verbose,axis=axis,_do_crop=_do_crop)
 
     def map_labels(self, label_map:LABEL_MAP , verbose:logging=True, inplace=False):
         """

TPTBox/core/np_utils.py

@@ -598,16 +598,8 @@ def np_calc_crop_around_centerpoint(
 
     cutout_coords_slices = tuple([slice(cutout_coords[i], cutout_coords[i + 1]) for i in range(0, n_dim * 2, 2)])
     arr_cut = arr[cutout_coords_slices]
-    arr_cut = np.pad(
-        arr_cut,
-        tuple(padding),
-    )
-    return (
-        arr_cut,
-        cutout_coords_slices,
-        tuple(padding),
-        # tuple([slice(padding[i][0], padding[i][1]) for i in range(n_dim)]),
-    )
+    arr_cut = np.pad(arr_cut, tuple(padding))
+    return (arr_cut, cutout_coords_slices, tuple(padding))
 
 
 def np_bbox_binary(img: np.ndarray, px_dist: int | Sequence[int] | np.ndarray = 0, raise_error=True) -> tuple[slice, ...]: