Merge pull request #91 from Hendrik-code/snapshot_readability

small bugfixes

Author: robert-graf

TPTBox/core/bids_files.py

@@ -186,28 +186,40 @@ def save_buffer(f: Path, buffer_name):
 
                 age = today - file_mod_time
                 if age.days >= int(max_age_days):
-                    print(
-                        "[ ] Delete Buffer - to old:",
-                        (folder / buffer_name),
-                        f"{' ':20}",
-                    ) if verbose else None
+                    (
+                        print(
+                            "[ ] Delete Buffer - to old:",
+                            (folder / buffer_name),
+                            f"{' ':20}",
+                        )
+                        if verbose
+                        else None
+                    )
                     (folder / buffer_name).unlink()
             if (folder / buffer_name).exists() and parent not in recompute_parents:
                 with open((folder / buffer_name), "rb") as b:
                     l = pickle.load(b)
+                    (
+                        print(
+                            f"[{len(l):8}] Read Buffer:",
+                            (folder / buffer_name),
+                            f"{' ':20}",
+                        )
+                        if verbose
+                        else None
+                    )
+                    files[dataset] += l
+            else:
+                (
                     print(
-                        f"[{len(l):8}] Read Buffer:",
+                        f"[{_cont:8}] Create new Buffer:",
                         (folder / buffer_name),
                         f"{' ':20}",
-                    ) if verbose else None
-                    files[dataset] += l
-            else:
-                print(
-                    f"[{_cont:8}] Create new Buffer:",
-                    (folder / buffer_name),
-                    f"{' ':20}",
-                    end="\r",
-                ) if verbose else None
+                        end="\r",
+                    )
+                    if verbose
+                    else None
+                )
                 files[dataset] += save_buffer((folder), buffer_name)
     if filter_file is not None:
         files: dict[Path | str, list[Path]] = {d: [g for g in f if filter_file(g)] for d, f in files.items()}
@@ -353,10 +365,14 @@ def add_file_2_subject(self, bids: BIDS_FILE | Path, ds=None) -> None:
         if subject not in self.subjects:
             self.subjects[subject] = Subject_Container(subject, self.sequence_splitting_keys)
         self.count_file += 1
-        print(
-            f"Found: {subject}, total file keys {(self.count_file)},  total subjects = {len(self.subjects)}    ",
-            end="\r",
-        ) if self.verbose else None
+        (
+            print(
+                f"Found: {subject}, total file keys {(self.count_file)},  total subjects = {len(self.subjects)}    ",
+                end="\r",
+            )
+            if self.verbose
+            else None
+        )
         self.subjects[subject].add(bids)
 
     def enumerate_subjects(self, sort=False, shuffle=False) -> list[tuple[str, Subject_Container]]:
@@ -729,6 +745,9 @@ def get_changed_path(  # noqa: C901
             info = {}
         if non_strict_mode and not self.BIDS_key.startswith("sub"):
             info["sub"] = self.BIDS_key.replace("_", "-").replace(".", "-")
+        else:
+            # replace _ with - in all info
+            self.info = {k: v.replace("_", "-") for k, v in self.info.items()}
         if isinstance(file_type, str) and file_type.startswith("."):
             file_type = file_type[1:]
         path = self.insert_info_into_path(path)

TPTBox/core/nii_wrapper.py

@@ -16,6 +16,7 @@
 from nibabel import Nifti1Header, Nifti1Image  # type: ignore
 from typing_extensions import Self
 
+from TPTBox.core import bids_files
 from TPTBox.core.compat import zip_strict
 from TPTBox.core.internal.nii_help import _resample_from_to, secure_save
 from TPTBox.core.nii_poi_abstract import Has_Grid
@@ -47,10 +48,7 @@
     np_unique_withoutzero,
     np_volume,
 )
-from TPTBox.logger.log_file import Log_Type
-
-from . import bids_files
-from .vert_constants import (
+from TPTBox.core.vert_constants import (
     AFFINE,
     AX_CODES,
     COORDINATE,
@@ -65,6 +63,7 @@
     logging,
     v_name2idx,
 )
+from TPTBox.logger.log_file import Log_Type
 
 if TYPE_CHECKING:
     from torch import device
@@ -1065,7 +1064,7 @@ def normalize_to_range_(self, min_value: int = 0, max_value: int = 1500, verbose
         mi, ma = self.min(), self.max()
         self += -mi + min_value  # min = 0
         self_dtype = self.dtype
-        max_value2 = ma
+        max_value2 = self.max() # this is a new value if min got shifted
         if max_value2 > max_value:
             self *= max_value / max_value2
             self.set_dtype_(self_dtype)
@@ -1125,7 +1124,9 @@ def smooth_gaussian_labelwise(
         boundary_mode: str = "nearest",
         dilate_prior: int = 0,
         dilate_connectivity: int = 1,
+        dilate_channelwise: bool = False,
         smooth_background: bool = True,
+        background_threshold: float | None = None,
         inplace: bool = False,
     ):
         """Smoothes the segmentation mask by applying a gaussian filter label-wise and then using argmax to derive the smoothed segmentation labels again.
@@ -1145,8 +1146,20 @@ def smooth_gaussian_labelwise(
             NII: The smoothed NII object.
         """
         assert self.seg, "You cannot use this on a non-segmentation NII"
-        smoothed = np_smooth_gaussian_labelwise(self.get_seg_array(), label_to_smooth=label_to_smooth, sigma=sigma, radius=radius, truncate=truncate, boundary_mode=boundary_mode, dilate_prior=dilate_prior, dilate_connectivity=dilate_connectivity,smooth_background=smooth_background,)
-        return self.set_array(smoothed,inplace,verbose=False)
+        smoothed = np_smooth_gaussian_labelwise(
+            self.get_seg_array(),
+            label_to_smooth=label_to_smooth,
+            sigma=sigma,
+            radius=radius,
+            truncate=truncate,
+            boundary_mode=boundary_mode,
+            dilate_prior=dilate_prior,
+            dilate_connectivity=dilate_connectivity,
+            smooth_background=smooth_background,
+            background_threshold=background_threshold,
+            dilate_channelwise=dilate_channelwise,
+        )
+        return self.set_array(smoothed, inplace, verbose=False)
 
     def smooth_gaussian_labelwise_(
         self,
@@ -1157,9 +1170,23 @@ def smooth_gaussian_labelwise_(
         boundary_mode: str = "nearest",
         dilate_prior: int = 1,
         dilate_connectivity: int = 1,
-        smooth_background: bool = True
+        dilate_channelwise: bool = False,
+        smooth_background: bool = True,
+        background_threshold: float | None = None,
     ):
-        return self.smooth_gaussian_labelwise(label_to_smooth=label_to_smooth, sigma=sigma, radius=radius, truncate=truncate, boundary_mode=boundary_mode, dilate_prior=dilate_prior, dilate_connectivity=dilate_connectivity, smooth_background=smooth_background, inplace=True,)
+        return self.smooth_gaussian_labelwise(
+            label_to_smooth=label_to_smooth,
+            sigma=sigma,
+            radius=radius,
+            truncate=truncate,
+            boundary_mode=boundary_mode,
+            dilate_prior=dilate_prior,
+            dilate_connectivity=dilate_connectivity,
+            smooth_background=smooth_background,
+            inplace=True,
+            background_threshold=background_threshold,
+            dilate_channelwise=dilate_channelwise,
+        )
 
     def to_ants(self):
         try:
@@ -1402,7 +1429,7 @@ def filter_connected_components(self, labels: int |list[int]|None=None,min_volum
         #print("filter",nii.unique())
         #assert max_count_component is None or nii.max() <= max_count_component, nii.unique()
         return self.set_array(arr, inplace=inplace)
-    def filter_connected_components_(self, labels: int |list[int]|None,min_volume:int=0,max_volume:int|None=None, max_count_component = None, connectivity: int = 3,keep_label=False):
+    def filter_connected_components_(self, labels: int |list[int]|None=None,min_volume:int=0,max_volume:int|None=None, max_count_component = None, connectivity: int = 3,keep_label=False):
         return self.filter_connected_components(labels,min_volume=min_volume,max_volume=max_volume, max_count_component = max_count_component, connectivity = connectivity,keep_label=keep_label,inplace=True)
 
     def get_segmentation_connected_components_center_of_mass(self, label: int, connectivity: int = 3, sort_by_axis: int | None = None) -> list[COORDINATE]:

TPTBox/core/np_utils.py

@@ -645,9 +645,14 @@ def np_compute_surface(arr: UINTARRAY, connectivity: int = 3, dilated_surface: b
     """
     assert 1 <= connectivity <= 3, f"expected connectivity in [1,3], but got {connectivity}"
     if dilated_surface:
-        return np_dilate_msk(arr.copy(), n_pixel=1, connectivity=connectivity) - arr
+        dil = np_dilate_msk(arr.copy(), n_pixel=1, connectivity=connectivity)
+        dil[arr != 0] = 0  # remove all non-zero entries
+        return dil
     else:
-        return arr - np_erode_msk(arr.copy(), n_pixel=1, connectivity=connectivity)
+        ero = np_erode_msk(arr.copy(), n_pixel=1, connectivity=connectivity)
+        arr = arr.copy()
+        arr[ero != 0] = 0  # remove all non-zero entries
+        return arr
 
 
 def np_point_coordinates(
@@ -969,7 +974,9 @@ def np_smooth_gaussian_labelwise(
     boundary_mode: str = "nearest",
     dilate_prior: int = 0,
     dilate_connectivity: int = 3,
+    dilate_channelwise: bool = False,
     smooth_background: bool = True,
+    background_threshold: float | None = None,
 ) -> UINTARRAY:
     """Smoothes selected labels in a segmentation mask using Gaussian filtering,
     while keeping other labels unaffected.
@@ -1010,7 +1017,7 @@ def np_smooth_gaussian_labelwise(
     for l in label_to_smooth:
         assert l in sem_labels, f"You want to smooth label {l} but it is not present in the given segmentation mask"
 
-    if dilate_prior > 0:
+    if dilate_prior > 0 and not dilate_channelwise:
         arr = np_dilate_msk(
             arr,
             n_pixel=dilate_prior,
@@ -1023,6 +1030,13 @@ def np_smooth_gaussian_labelwise(
     sem_labels_plus_background.append(0)
     for l in sem_labels_plus_background[:-1]:
         arr_l = (arr == l).astype(float)
+        if dilate_prior > 0 and dilate_channelwise:
+            arr_l = np_dilate_msk(
+                arr_l,
+                n_pixel=dilate_prior,
+                label_ref=1,
+                connectivity=dilate_connectivity,
+            )
         if l in label_to_smooth:
             arr_l = gaussian_filter(
                 arr_l,
@@ -1053,6 +1067,9 @@ def np_smooth_gaussian_labelwise(
     seg_arr_smoothed = np.argmax(arr_stack, axis=0)
     seg_arr_s = seg_arr_smoothed.copy()
 
+    if background_threshold is not None:
+        seg_arr_smoothed[seg_arr_smoothed < background_threshold] = len(sem_labels_plus_background) - 1  # background label
+
     for idx, l in enumerate(sem_labels_plus_background):
         seg_arr_s[seg_arr_smoothed == idx] = l
 

TPTBox/core/vert_constants.py

@@ -368,6 +368,7 @@ def __init__(
         self._rib = None
         self._ivd = None
         self._endplate = None
+        self.has_rib = has_rib
         if has_rib:
             self._rib = (
                 vertebra_label + VERTEBRA_INSTANCE_RIB_LABEL_OFFSET if vertebra_label != 28 else 21 + VERTEBRA_INSTANCE_RIB_LABEL_OFFSET
@@ -487,7 +488,7 @@ def get_previous_poi(self, poi: POI | NII | list[int]):
     C3 = 3
     C4 = 4
     C5 = 5
-    C6 = 6
+    C6 = 6, True, True
     C7 = 7, True, True
     T1 = 8, True, True
     T2 = 9, True, True

TPTBox/mesh3D/mesh.py

@@ -31,14 +31,17 @@ def __init__(self, mesh: pv.PolyData) -> None:
     def save(self, filepath: str | Path, mode: MeshOutputType = MeshOutputType.PLY, verbose: logging = True):
         filepath = str(filepath)
         if not filepath.endswith(mode.value):
-            filepath += mode.value
+            filepath += "." + mode.value
 
         filepath = Path(filepath)
-        if not filepath.exists():
-            raise FileNotFoundError(filepath)
+        if not filepath.parent.exists():
+            raise FileNotFoundError(filepath.parent)
 
         if mode == MeshOutputType.PLY:
-            self.mesh.export_obj(filepath)
+            try:
+                self.mesh.export_obj(filepath)
+            except AttributeError:
+                self.mesh.save(filepath)
         else:
             raise NotImplementedError(f"save with mode {mode}")
         log.print(f"Saved mesh: {filepath}", Log_Type.SAVE, verbose=verbose)
@@ -61,6 +64,18 @@ def show(self):
         pl.add_mesh(self.mesh)
         pl.show()
 
+    def save_to_html(self, file_output: str | Path):
+        pv.start_xvfb()
+        pl = pv.Plotter()
+        pl.set_background("black", top=None)
+        pl.add_axes()
+        pv.global_theme.axes.show = True
+        pv.global_theme.edge_color = "white"
+        pv.global_theme.interactive = True
+
+        pl.add_mesh(self.mesh)
+        pl.export_html(file_output)
+
 
 class SegmentationMesh(Mesh3D):
     def __init__(self, int_arr: np.ndarray | Image_Reference) -> None:

TPTBox/mesh3D/snapshot3D.py

@@ -34,7 +34,8 @@ def make_snapshot3D(
     ids_list: list[Sequence[int]] | None = None,
     smoothing=20,
     resolution: float | None = None,
-    width_factor=1.0,
+    width_factor: float = 1.0,
+    scale_factor: int = 1,
     verbose=True,
     crop=True,
 ) -> Image.Image:
@@ -80,7 +81,7 @@ def make_snapshot3D(
     nii = to_nii_seg(img)
     if crop:
         try:
-            nii.apply_crop_(nii.compute_crop())
+            nii.apply_crop_(nii.compute_crop(dist=2))
         except ValueError:
             pass
     if resolution is None:
@@ -98,7 +99,7 @@ def make_snapshot3D(
         ids_list = ids_list2
 
     # TOP : ("A", "I", "R")
-    nii = nii.reorient(("A", "S", "L")).rescale_((resolution, resolution, resolution))
+    nii = nii.reorient(("A", "S", "L")).rescale_((resolution, resolution, resolution), mode="constant")
     width = int(max(nii.shape[0], nii.shape[2]) * width_factor)
     window_size = (width * len(ids_list), nii.shape[1])
     with Xvfb():
@@ -110,7 +111,7 @@ def make_snapshot3D(
             _plot_sub_seg(scene, nii.extract_label(ids, keep_label=True), x, 0, smoothing, view[i % len(view)])
         scene.projection(proj_type="parallel")
         scene.reset_camera_tight(margin_factor=1.02)
-        window.record(scene, size=window_size, out_path=output_path, reset_camera=False)
+        window.record(scene, size=window_size, out_path=output_path, reset_camera=False, magnification=scale_factor)
         scene.clear()
     if not is_tmp:
         logger.on_save("Save Snapshot3D:", output_path, verbose=verbose)
@@ -129,6 +130,7 @@ def make_snapshot3D_parallel(
     resolution: float = 2,
     cpus=10,
     width_factor=1.0,
+    scale_factor: int = 1,
     override=True,
 ):
     ress = []
@@ -146,6 +148,7 @@ def make_snapshot3D_parallel(
                     "smoothing": smoothing,
                     "resolution": resolution,
                     "width_factor": width_factor,
+                    "scale_factor": scale_factor,
                 },
             )
             ress.append(res)

TPTBox/registration/deepali/deepali_model.py

@@ -173,8 +173,9 @@ def __init__(
         fixed_mask: Image_Reference | None = None,
         moving_mask: Image_Reference | None = None,
         # normalize
-        normalize_strategy: Literal["auto", "CT", "MRI"]
-        | None = "auto",  # Override on_normalize for finer normalization schema or normalize before and set to None. auto: [min,max] -> [0,1]; None: Do noting
+        normalize_strategy: (
+            Literal["auto", "CT", "MRI"] | None
+        ) = "auto",  # Override on_normalize for finer normalization schema or normalize before and set to None. auto: [min,max] -> [0,1]; None: Do noting
         # Pyramid
         pyramid_levels: int | None = None,  # 1/None = no pyramid; int: number of stacks, tuple from to (0 is finest)
         finest_level: int = 0,
@@ -188,6 +189,7 @@ def __init__(
         transform_init: PathStr | None = None,  # reload initial flowfield from file
         optim_name="Adam",  # Optimizer name defined in torch.optim. or override on_optimizer finer control
         lr: float | Sequence[float] = 0.01,  # Learning rate
+        lr_end_factor: float | None = None,  # if set, will use a LinearLR scheduler to reduce the learning rate to this factor * lr
         optim_args=None,  # args of Optimizer with out lr
         smooth_grad=0.0,
         verbose=99,
@@ -245,6 +247,7 @@ def __init__(
             transform_init=transform_init,
             optim_name=optim_name,
             lr=lr,
+            lr_end_factor=lr_end_factor,
             optim_args=optim_args,
             smooth_grad=smooth_grad,
             verbose=verbose,