Stitching refactor + RAM-saving overlap + c_val option + trivial-memory warning

TPTBox/core/internal/nii_help.py

@@ -205,6 +205,7 @@ def _resample_from_to(
     mode: MODES = "nearest",
     align_corners: bool | Sentinel = Sentinel(),  # noqa: B008
     out_dtype: np.dtype | type | str | None = None,
+    c_val: float | None = None,
 ) -> tuple[np.ndarray, np.ndarray, object]:
     """Resample *from_img* into the voxel space defined by *to_img*.
 
@@ -323,7 +324,7 @@ def _resample_from_to(
         to_shape,
         order=order,
         mode=mode,
-        cval=from_img.get_c_val(),
+        cval=from_img.get_c_val(c_val),
         output=scipy_out,
     )
     if post_cast is not None:

TPTBox/core/nii_wrapper.py

@@ -999,7 +999,7 @@ def pad_to(self, target_shape: list[int] | tuple[int, int, int] | Self, mode: MO
             s = s.apply_crop(tuple(crop),inplace=inplace)
         return s.apply_pad(padding,inplace=inplace,mode=mode)
 
-    def apply_pad(self, padd: Sequence[tuple[int | None, int | None]] | int | None, mode: MODES = "constant", inplace=False, verbose: logging = True,) -> Self:
+    def apply_pad(self, padd: Sequence[tuple[int | None, int | None]] | int | None, mode: MODES = "constant", inplace=False, verbose: logging = True, c_val: float | None = None) -> Self:
         """Pads the image with explicit per-axis ``(before, after)`` amounts.
 
         The affine is updated so that the world-space origin is preserved (i.e. the
@@ -1064,7 +1064,7 @@ def apply_pad(self, padd: Sequence[tuple[int | None, int | None]] | int | None,
 
         args = {}
         if mode == "constant":
-            args["constant_values"] = self.get_c_val()
+            args["constant_values"] = self.get_c_val(c_val)
 
         if mode == "nearest":
             mode = "edge"
@@ -1249,7 +1249,7 @@ def resample_from_to(self, to_vox_map:Image_Reference|Has_Grid|tuple[SHAPE,AFFIN
                 pad_after = dst_shape - shift - src_shape
                 pad = tuple((int(b), int(a)) for b, a in zip(pad_before, pad_after))
                 try:
-                    ret = s.apply_pad(pad,mode=mode,inplace=inplace,verbose=verbose)
+                    ret = s.apply_pad(pad,mode=mode,inplace=inplace,verbose=verbose,c_val=c_val)
                     valid = ret.assert_affine(mapping,raise_error=False,origin_tolerance=0.0001,error_tolerance=0.0001,shape_tolerance=0)
                     if valid:
                         log.print(f"resample_from_to only needs padding/cropping {pad}",verbose=verbose)
@@ -1263,7 +1263,7 @@ def resample_from_to(self, to_vox_map:Image_Reference|Has_Grid|tuple[SHAPE,AFFIN
         log.print(f"resample_from_to: {self} to {mapping}",verbose=verbose)
         if order is None:
             order = 0 if self.seg else 3
-        nii = _resample_from_to(self, mapping,order=order, mode=mode,align_corners=align_corners, out_dtype=out_dtype)
+        nii = _resample_from_to(self, mapping,order=order, mode=mode,align_corners=align_corners, out_dtype=out_dtype, c_val=c_val)
 
 
         if inplace:

TPTBox/segmentation/VibeSeg/inference_nnunet.py

@@ -267,13 +267,6 @@ def run_inference_on_file(
         if "memory_factor" not in ds_info:
             missing_mem_keys.append("memory_factor")
         memory_factor = float(ds_info.get("memory_factor", 160))
-    if missing_mem_keys:
-        _suggest_memory_estimation_script(
-            idx,
-            model_path,
-            f"Memory parameter(s) {missing_mem_keys} not set in the model's dataset.json; falling back to defaults. {memory_base=}, {memory_factor=}",
-            logger=logger,
-        )
 
     use_folds_arg = tuple(folds) if len(folds) != 5 else None
     # Include every setting that changes the loaded predictor so a cache hit is always equivalent
@@ -362,11 +355,43 @@ def run_inference_on_file(
     if padd != 0:
         p = (padd, padd)
         input_nii = [i.apply_pad([p, p, p], mode="reflect") for i in input_nii]
+
+    def _defaults_fit_gpu(memory: float, gpu: int | None, safety_factor: float = 2.0) -> bool:
+        """Whether the fallback ``memory_base`` reservation is a small fraction of GPU memory.
+
+        Returns ``True`` when ``memory_base * safety_factor <= total_gpu_memory_mb`` —
+        i.e. the fallback reservation leaves ample headroom for tile scheduling and
+        the missing ``memory_base``/``memory_factor`` defaults are safe to use
+        without warning. Returns ``False`` on CUDA-unavailable systems or when GPU
+        memory can't be probed, so the load-time warning still fires on
+        memory-constrained setups (CPU-only, small GPUs, or driver errors).
+        """
+        try:
+            import torch
+
+            if not torch.cuda.is_available():
+                return False
+            device = torch.device(f"cuda:{gpu}") if gpu is not None else torch.device("cuda:0")
+            _, total_bytes = torch.cuda.mem_get_info(device)
+            total_mb = total_bytes / (1024**2)
+        except Exception:  # noqa: BLE001
+            return False
+        return memory * safety_factor <= total_mb
+
+    num_classes = int(nnunet.label_manager.num_segmentation_heads)
+    est_full_mb = estimate_peak_ram_mb(input_nii[0].shape, num_classes, len(input_nii))
+    # Only warn if memory requirement is non-trivial.
+    if missing_mem_keys and not _defaults_fit_gpu(est_full_mb, gpu):
+        _suggest_memory_estimation_script(
+            idx,
+            model_path,
+            f"Memory parameter(s) {missing_mem_keys} not set in the model's dataset.json; falling back to defaults. {memory_base=}, {memory_factor=}",
+            logger=logger,
+        )
+
     if _cpu_chunks is None or _cpu_chunks <= 1:
-        num_classes = int(nnunet.label_manager.num_segmentation_heads)
         total_ram_mb = _get_total_ram_mb()
         target_ram_mb = total_ram_mb * 0.5
-        est_full_mb = estimate_peak_ram_mb(input_nii[0].shape, num_classes, len(input_nii))
         if est_full_mb > target_ram_mb:
             shape = input_nii[0].shape
             split_axis = int(np.argmax(shape))

TPTBox/stitching/README.md

@@ -8,8 +8,8 @@ You can verify alignment by opening the images in ITKSnap with "open additional
 
 | Function | Description |
 |---|---|
-| `stitching(nii_list, out, ...)` | Stitch a list of `NII` objects; returns `(result_nii, ramp_nii)` |
-| `stitching_raw(paths, out, ...)` | Stitch from file paths directly |
+| `stitching(inputs, out, ...)` | High-level wrapper. `inputs` accepts any mix of `BIDS_FILE`, `NII`, `str`, or `Path`. Resolves the output path from a `BIDS_FILE` when one is passed. Returns `(result_nii, ramp_nii)` as `NII` objects. |
+| `stitching_raw(images, out, ...)` | Low-level driver. `images` accepts file paths, pre-loaded `NII` objects, or (fallback) `Nifti1Image` objects. Returns `(result_nii, ramp_nii)` as `NII` objects. |
 | `NAKO_stitch_T2w(HWS, BWS, LWS, n4_after_stitch=False)` | Stitch the three NAKO sagittal T2w spine stations (HWS cervical, BWS thoracic, LWS lumbar) into one volume |
 
 ![Example of a stitching](https://raw.githubusercontent.com/Hendrik-code/TPTBox/main/TPTBox/stitching/stitching.jpg "Example of a stitching")
@@ -24,7 +24,7 @@ stitching.py
 [-i IMAGES [IMAGES ...]] a list of input image paths
 [-o OUTPUT] The output image path
 [-v] verbose - if set, there will be more printouts.
-[-min_value MIN_VALUE] New pixels not present will get this value. Recommended 0 for MRI and for CT -1024 or the known min-value.
+[-min_value MIN_VALUE] Background fill used when resampling each chunk, and — when set explicitly — a hard floor applied to the stitched output. Pass 0 for MRI magnitude, -1024 for CT. Omitting it (the Python-API default `None`) uses 0 as the internal background and only applies a hard floor when the output dtype cannot represent negatives (unsigned integer types) — signed / float outputs then keep legitimate negatives (Philips-scaled fat-fraction, phase, B0 offsets).
 [-seg] This flag is required if you merge segmentation Niftis.
 Switches:
 [-no_bias] If set: Do not use n4_bias_field_correction. It speeds up the process, but n4_bias_field_correction helps in roughly aligning the histogram.
@@ -67,14 +67,15 @@ list_of_files = [
 # Call the stitching function
 # This will combine your images into a single NIfTI file
 stitching(
-    list_of_files,  # List of input files
-    out="out_path_stitched_image.nii.gz",  # Path to save stitched output
-    is_seg=False,  # Set True if these are segmentation masks
-    is_ct=False,  # True for CT, min_value will by -1024 instead of 0
+    list_of_files,  # BIDS_FILE / NII / str / Path (any mix)
+    out="out_path_stitched_image.nii.gz",  # Path or BIDS_FILE for the stitched output
+    is_seg=False,  # Set True for segmentation masks (forces integer dtype, nearest-neighbour resample)
+    is_ct=False,  # Sets min_value to -1024 (CT air) when min_value is not passed explicitly
     kick_out_fully_integrated_images=True,
-    dtype=float,  # Data type of the output image
+    dtype=float,  # Output dtype; "auto" picks the smallest lossless type from the inputs (float32 when any input has a non-trivial scl_slope/inter)
     match_histogram=False,  # Match intensity histograms across images
     store_ramp=False,  # Store blending ramp (optional)
+    min_value=None,  # Explicit background/floor. None (default) = clip only when the output dtype is unsigned integer (protects against negative-to-huge wraparound) and let signed / float outputs keep legitimate negatives. Pass 0 to floor magnitude MR at 0 so cubic-spline resample ringing doesn't leak small negatives into what should be a non-negative volume; pass -1024 for CT.
 )
 ```