Skip to content

Commit 72c00d7

Browse files
authored
Merge pull request #157 from Hendrik-code/robert-pc-pr-stitching-refactor
Stitching refactor + RAM-saving overlap + c_val option + trivial-memory warning
2 parents 7f39708 + 424df86 commit 72c00d7

7 files changed

Lines changed: 402 additions & 414 deletions

File tree

‎TPTBox/core/internal/nii_help.py‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -205,6 +205,7 @@ def _resample_from_to(
205205
mode: MODES = "nearest",
206206
align_corners: bool | Sentinel = Sentinel(), # noqa: B008
207207
out_dtype: np.dtype | type | str | None = None,
208+
c_val: float | None = None,
208209
) -> tuple[np.ndarray, np.ndarray, object]:
209210
"""Resample *from_img* into the voxel space defined by *to_img*.
210211
@@ -323,7 +324,7 @@ def _resample_from_to(
323324
to_shape,
324325
order=order,
325326
mode=mode,
326-
cval=from_img.get_c_val(),
327+
cval=from_img.get_c_val(c_val),
327328
output=scipy_out,
328329
)
329330
if post_cast is not None:

‎TPTBox/core/nii_wrapper.py‎

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -1000,7 +1000,7 @@ def pad_to(self, target_shape: list[int] | tuple[int, int, int] | Self, mode: MO
10001000
s = s.apply_crop(tuple(crop),inplace=inplace)
10011001
return s.apply_pad(padding,inplace=inplace,mode=mode)
10021002

1003-
def apply_pad(self, padd: Sequence[tuple[int | None, int | None]] | int | None, mode: MODES = "constant", inplace=False, verbose: logging = True,) -> Self:
1003+
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:
10041004
"""Pads the image with explicit per-axis ``(before, after)`` amounts.
10051005
10061006
The affine is updated so that the world-space origin is preserved (i.e. the
@@ -1065,7 +1065,7 @@ def apply_pad(self, padd: Sequence[tuple[int | None, int | None]] | int | None,
10651065

10661066
args = {}
10671067
if mode == "constant":
1068-
args["constant_values"] = self.get_c_val()
1068+
args["constant_values"] = self.get_c_val(c_val)
10691069

10701070
if mode == "nearest":
10711071
mode = "edge"
@@ -1250,7 +1250,7 @@ def resample_from_to(self, to_vox_map:Image_Reference|Has_Grid|tuple[SHAPE,AFFIN
12501250
pad_after = dst_shape - shift - src_shape
12511251
pad = tuple((int(b), int(a)) for b, a in zip(pad_before, pad_after))
12521252
try:
1253-
ret = s.apply_pad(pad,mode=mode,inplace=inplace,verbose=verbose)
1253+
ret = s.apply_pad(pad,mode=mode,inplace=inplace,verbose=verbose,c_val=c_val)
12541254
valid = ret.assert_affine(mapping,raise_error=False,origin_tolerance=0.0001,error_tolerance=0.0001,shape_tolerance=0)
12551255
if valid:
12561256
log.print(f"resample_from_to only needs padding/cropping {pad}",verbose=verbose)
@@ -1264,7 +1264,7 @@ def resample_from_to(self, to_vox_map:Image_Reference|Has_Grid|tuple[SHAPE,AFFIN
12641264
log.print(f"resample_from_to: {self} to {mapping}",verbose=verbose)
12651265
if order is None:
12661266
order = 0 if self.seg else 3
1267-
nii = _resample_from_to(self, mapping,order=order, mode=mode,align_corners=align_corners, out_dtype=out_dtype)
1267+
nii = _resample_from_to(self, mapping,order=order, mode=mode,align_corners=align_corners, out_dtype=out_dtype, c_val=c_val)
12681268

12691269

12701270
if inplace:

‎TPTBox/segmentation/VibeSeg/inference_nnunet.py‎

Lines changed: 34 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -267,13 +267,6 @@ def run_inference_on_file(
267267
if "memory_factor" not in ds_info:
268268
missing_mem_keys.append("memory_factor")
269269
memory_factor = float(ds_info.get("memory_factor", 160))
270-
if missing_mem_keys:
271-
_suggest_memory_estimation_script(
272-
idx,
273-
model_path,
274-
f"Memory parameter(s) {missing_mem_keys} not set in the model's dataset.json; falling back to defaults. {memory_base=}, {memory_factor=}",
275-
logger=logger,
276-
)
277270

278271
use_folds_arg = tuple(folds) if len(folds) != 5 else None
279272
# Include every setting that changes the loaded predictor so a cache hit is always equivalent
@@ -362,11 +355,43 @@ def run_inference_on_file(
362355
if padd != 0:
363356
p = (padd, padd)
364357
input_nii = [i.apply_pad([p, p, p], mode="reflect") for i in input_nii]
358+
359+
def _defaults_fit_gpu(memory: float, gpu: int | None, safety_factor: float = 2.0) -> bool:
360+
"""Whether the fallback ``memory_base`` reservation is a small fraction of GPU memory.
361+
362+
Returns ``True`` when ``memory_base * safety_factor <= total_gpu_memory_mb`` —
363+
i.e. the fallback reservation leaves ample headroom for tile scheduling and
364+
the missing ``memory_base``/``memory_factor`` defaults are safe to use
365+
without warning. Returns ``False`` on CUDA-unavailable systems or when GPU
366+
memory can't be probed, so the load-time warning still fires on
367+
memory-constrained setups (CPU-only, small GPUs, or driver errors).
368+
"""
369+
try:
370+
import torch
371+
372+
if not torch.cuda.is_available():
373+
return False
374+
device = torch.device(f"cuda:{gpu}") if gpu is not None else torch.device("cuda:0")
375+
_, total_bytes = torch.cuda.mem_get_info(device)
376+
total_mb = total_bytes / (1024**2)
377+
except Exception: # noqa: BLE001
378+
return False
379+
return memory * safety_factor <= total_mb
380+
381+
num_classes = int(nnunet.label_manager.num_segmentation_heads)
382+
est_full_mb = estimate_peak_ram_mb(input_nii[0].shape, num_classes, len(input_nii))
383+
# Only warn if memory requirement is non-trivial.
384+
if missing_mem_keys and not _defaults_fit_gpu(est_full_mb, gpu):
385+
_suggest_memory_estimation_script(
386+
idx,
387+
model_path,
388+
f"Memory parameter(s) {missing_mem_keys} not set in the model's dataset.json; falling back to defaults. {memory_base=}, {memory_factor=}",
389+
logger=logger,
390+
)
391+
365392
if _cpu_chunks is None or _cpu_chunks <= 1:
366-
num_classes = int(nnunet.label_manager.num_segmentation_heads)
367393
total_ram_mb = _get_total_ram_mb()
368394
target_ram_mb = total_ram_mb * 0.5
369-
est_full_mb = estimate_peak_ram_mb(input_nii[0].shape, num_classes, len(input_nii))
370395
if est_full_mb > target_ram_mb:
371396
shape = input_nii[0].shape
372397
split_axis = int(np.argmax(shape))

‎TPTBox/stitching/README.md‎

Lines changed: 9 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -8,8 +8,8 @@ You can verify alignment by opening the images in ITKSnap with "open additional
88

99
| Function | Description |
1010
|---|---|
11-
| `stitching(nii_list, out, ...)` | Stitch a list of `NII` objects; returns `(result_nii, ramp_nii)` |
12-
| `stitching_raw(paths, out, ...)` | Stitch from file paths directly |
11+
| `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. |
12+
| `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. |
1313
| `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 |
1414

1515
![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
2424
[-i IMAGES [IMAGES ...]] a list of input image paths
2525
[-o OUTPUT] The output image path
2626
[-v] verbose - if set, there will be more printouts.
27-
[-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.
27+
[-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).
2828
[-seg] This flag is required if you merge segmentation Niftis.
2929
Switches:
3030
[-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 = [
6767
# Call the stitching function
6868
# This will combine your images into a single NIfTI file
6969
stitching(
70-
list_of_files, # List of input files
71-
out="out_path_stitched_image.nii.gz", # Path to save stitched output
72-
is_seg=False, # Set True if these are segmentation masks
73-
is_ct=False, # True for CT, min_value will by -1024 instead of 0
70+
list_of_files, # BIDS_FILE / NII / str / Path (any mix)
71+
out="out_path_stitched_image.nii.gz", # Path or BIDS_FILE for the stitched output
72+
is_seg=False, # Set True for segmentation masks (forces integer dtype, nearest-neighbour resample)
73+
is_ct=False, # Sets min_value to -1024 (CT air) when min_value is not passed explicitly
7474
kick_out_fully_integrated_images=True,
75-
dtype=float, # Data type of the output image
75+
dtype=float, # Output dtype; "auto" picks the smallest lossless type from the inputs (float32 when any input has a non-trivial scl_slope/inter)
7676
match_histogram=False, # Match intensity histograms across images
7777
store_ramp=False, # Store blending ramp (optional)
78+
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.
7879
)
7980
```
8081

0 commit comments

Comments
 (0)