improve nnunet inference for large files

Author: robert-graf

TPTBox/segmentation/TotalVibeSeg/inference_nnunet.py

@@ -57,12 +57,13 @@ def run_inference_on_file(
     padd: int = 0,
     ddevice: Literal["cpu", "cuda", "mps"] = "cuda",
     _model_path=None,
+    step_size=0.5,
 ) -> tuple[Image_Reference, np.ndarray | None]:
     global model_path  # noqa: PLW0603
     if _model_path is not None:
         _model_path = Path(_model_path)
         model_path = _model_path / "nnUNet_results"
-        assert model_path.exists(), _model_path
+        assert model_path.exists(), model_path
     if out_file is not None and Path(out_file).exists() and not override:
         return out_file, None
 
@@ -78,7 +79,7 @@ def run_inference_on_file(
         nnunet_path = next(next(iter(model_path.glob(f"*{idx:03}*"))).glob("*__nnUNetPlans*"))
     folds = [int(f.name.split("fold_")[-1]) for f in nnunet_path.glob("fold*")]
     if max_folds is not None:
-        folds = folds[:max_folds]
+        folds = max_folds if isinstance(max_folds, list) else folds[:max_folds]
 
     # if idx in _unets:
     #    nnunet = _unets[idx]
@@ -90,6 +91,7 @@ def run_inference_on_file(
         use_folds=tuple(folds) if len(folds) != 5 else None,
         gpu=gpu,
         ddevice=ddevice,
+        step_size=step_size,
     )
 
     #    _unets[idx] = nnunet
@@ -118,6 +120,7 @@ def run_inference_on_file(
     if zoom is not None:
         input_nii = [i.rescale_(zoom, mode=mode) for i in input_nii]
     input_nii = [squash_so_it_fits_in_float16(i) for i in input_nii]
+
     if crop:
         crop = input_nii[0].compute_crop(minimum=20)
         input_nii = [i.apply_crop(crop) for i in input_nii]
@@ -158,8 +161,13 @@ def run_total_seg(
     fill_holes=False,
     crop=False,
     max_folds: int | None = None,
+    _model_path=None,
+    step_size=0.5,
     **_kargs,
 ):
+    global model_path
+    if _model_path is not None:
+        model_path = _model_path
     if dataset_id is None:
         for idx in known_idx:
             download_weights(idx)
@@ -210,4 +218,5 @@ def run_total_seg(
         fill_holes=fill_holes,
         crop=crop,
         max_folds=max_folds,
+        step_size=step_size,
     )[0]

TPTBox/segmentation/nnUnet_utils/export_prediction.py

@@ -20,32 +20,35 @@ def convert_predicted_logits_to_segmentation_with_correct_shape(
     properties_dict: dict,
     return_probabilities: bool = False,
     num_threads_torch: int = 8,
+    rescale=True,
 ):
     old_threads = torch.get_num_threads()
     torch.set_num_threads(num_threads_torch)
-
-    # resample to original shape
-    current_spacing = (
-        configuration_manager.spacing
-        if len(configuration_manager.spacing) == len(properties_dict["shape_after_cropping_and_before_resampling"])
-        else [properties_dict["spacing"][0], *configuration_manager.spacing]
-    )
-    predicted_logits = configuration_manager.resampling_fn_probabilities(
-        predicted_logits,
-        properties_dict["shape_after_cropping_and_before_resampling"],
-        current_spacing,
-        properties_dict["spacing"],
-    )
+    if rescale:
+        # resample to original shape
+        current_spacing = (
+            configuration_manager.spacing
+            if len(configuration_manager.spacing) == len(properties_dict["shape_after_cropping_and_before_resampling"])
+            else [properties_dict["spacing"][0], *configuration_manager.spacing]
+        )
+        predicted_logits = configuration_manager.resampling_fn_probabilities(
+            predicted_logits,
+            properties_dict["shape_after_cropping_and_before_resampling"],
+            current_spacing,
+            properties_dict["spacing"],
+        )
     # return value of resampling_fn_probabilities can be ndarray or Tensor but that doesnt matter because
     # apply_inference_nonlin will covnert to torch
-    predicted_probabilities = label_manager.apply_inference_nonlin(predicted_logits)
-    del predicted_logits
-    segmentation = label_manager.convert_probabilities_to_segmentation(predicted_probabilities)
-
+    # And this is stupid because convert_probabilities_to_segmentation transforms it back to a numpy...
+    if label_manager.has_regions:
+        # Softmax does not change when we use argmax in the next step
+        predicted_logits = label_manager.apply_inference_nonlin(predicted_logits)
     # segmentation may be torch.Tensor but we continue with numpy
-    if isinstance(segmentation, torch.Tensor):
-        segmentation = segmentation.cpu().numpy()
+    if isinstance(predicted_logits, torch.Tensor):
+        predicted_logits = predicted_logits.cpu().numpy()
 
+    segmentation = label_manager.convert_probabilities_to_segmentation(predicted_logits)
+    del predicted_logits
     # put segmentation in bbox (revert cropping)
     segmentation_reverted_cropping = np.zeros(
         properties_dict["shape_before_cropping"],

TPTBox/segmentation/nnUnet_utils/inference_api.py

@@ -29,6 +29,7 @@ def load_inf_model(
     init_threads: bool = True,
     allow_non_final: bool = True,
     inference_augmentation: bool = False,
+    use_gaussian=True,
     verbose: bool = False,
     gpu=None,
 ) -> nnUNetPredictor:
@@ -69,11 +70,11 @@ def load_inf_model(
 
     predictor = nnUNetPredictor(
         tile_step_size=step_size,
-        use_gaussian=True,
+        use_gaussian=use_gaussian,
         use_mirroring=inference_augmentation,  # <- mirroring augmentation!
         perform_everything_on_gpu=ddevice != "cpu",
         device=device,
-        verbose=False,
+        verbose=verbose,
         verbose_preprocessing=False,
         cuda_id=0 if gpu is None else gpu,
     )
@@ -116,6 +117,8 @@ def run_inference(
     Returns:
         Segmentation (NII), Uncertainty Map (NII), Softmax Logits (numpy arr)
     """
+    if logits:
+        raise NotImplementedError("logits=True")
     if isinstance(input_nii, str):
         assert input_nii.endswith(".nii.gz"), f"input file is not a .nii.gz! Got {input_nii}"
         input_nii = NII.load(input_nii, seg=False)
@@ -124,52 +127,80 @@ def run_inference(
     if isinstance(input_nii, NII):
         input_nii = [input_nii]
     orientation = input_nii[0].orientation
-    zoom = input_nii[0].zoom
 
     img_arrs = []
     # Prepare for nnUNet behavior
     for i in input_nii:
         if reorient_PIR:
             i.reorient_()
-        sitk_nii = sitk_utils.nii_to_sitk(i)
-        nii_img_converted = sitk.GetArrayFromImage(sitk_nii).astype(np.float16)[np.newaxis, :]
-        # nii_img_converted = i.get_array()
-        # nii_img_converted = np.pad(nii_img_converted, pad_width=pad_size, mode="edge")
-        # nii_img_converted = np.swapaxes(nii_img_converted, 0, 2)[np.newaxis, :].astype(np.float16)
+        a = i.get_array().astype(np.float16)
+        nii_img_converted = np.transpose(a, axes=a.ndim - 1 - np.arange(a.ndim))[np.newaxis, :]
         img_arrs.append(nii_img_converted)
     try:
         img = np.vstack(img_arrs)
     except Exception:
-        print([a.shape for a in img_arrs])
+        print("could not stack images; shapes=", [a.shape for a in img_arrs])
         raise
     props = {
-        "sitk_stuff": {
-            # this saves the sitk geometry information. This part is NOT used by nnU-Net!
-            "spacing": sitk_nii.GetSpacing(),  # type:ignore
-            "origin": sitk_nii.GetOrigin(),  # type:ignore
-            "direction": sitk_nii.GetDirection(),  # type:ignore
-        },
-        "spacing": zoom[::-1],  # PIR
+        "spacing": i.zoom[::-1],  # PIR
     }
-    out = predictor.predict_single_npy_array(img, props, save_or_return_probabilities=logits)
-    if logits:
-        segmentation, _, softmax_logits = out  # type: ignore
-        softmax_logits = np.expand_dims(softmax_logits.astype(np.float16), 0)
-        # softmax_logits = np.swapaxes(softmax_logits, 0, 3)
-        # PRI label
-        # softmax_logits = np.swapaxes(softmax_logits, 1, 2)
-    else:
-        segmentation, _ = out  # type: ignore
-        softmax_logits = None
-    itk_image = sitk.GetImageFromArray(segmentation.astype(np.uint8))
-    itk_image.SetSpacing(sitk_nii.GetSpacing())
-    itk_image.SetOrigin(sitk_nii.GetOrigin())
-    itk_image.SetDirection(sitk_nii.GetDirection())
-    seg_nii = sitk_utils.sitk_to_nii(itk_image, True)
-
-    # segmentation = np.swapaxes(segmentation, 0, 2)
-    # assert isinstance(segmentation, np.ndarray)
-    # seg_nii = NII(nib.ni1.Nifti1Image(segmentation, affine=affine, header=header), seg=True)
-
+    out = predictor.predict_single_npy_array(img, props, logits=False, rescale=False)
+    segmentation: np.ndarray = out  # type: ignore
+    softmax_logits = None
+    segmentation = np.transpose(segmentation.astype(np.uint8), axes=segmentation.ndim - 1 - np.arange(segmentation.ndim))
+    assert segmentation.shape == input_nii[0].shape
+    seg_nii = input_nii[0].set_array(segmentation.astype(np.uint8), seg=True)
     seg_nii.reorient_(orientation, verbose=False)
     return seg_nii, None, softmax_logits
+
+
+# def predict_single_npy_array(predictor: nnUNetPredictor, img, props, logits, rescale):
+#    return predictor.predict_single_npy_array(img, props, save_or_return_probabilities=logits, rescale=rescale)
+#
+#    def fun(x):
+#        return predictor.predict_single_npy_array(x, props, save_or_return_probabilities=False)[0][None]
+#
+#    p = 750 if max_v % 700 > max_v % 800 else 800
+#    patch_size = tuple(p for _ in img.shape)
+#    overlap = min(50, max(predictor.configuration_manager.patch_size) // 2)
+#    print(f"image very large ({img.shape}>1000); use sliding window", f"{patch_size=}", predictor.configuration_manager.patch_size)
+#
+#    return sliding_nd_slices(img, patch_size=patch_size, overlap=overlap, fun=fun)[0], None
+
+
+def sliding_nd_slices(arr: np.ndarray, patch_size, overlap, fun):
+    print("sliding window")
+    step = tuple(p - overlap for p in patch_size)
+    half_overlap = overlap // 2
+    shape = arr.shape
+
+    # Compute number of steps in each dimension
+    ranges = [range(0, max(s, 1), st) if s != 1 else [0] for s, st in zip(shape, step)]
+    result = np.zeros_like(arr)
+    for starts in np.ndindex(*[len(r) for r in ranges]):
+        # Compute actual start and end indices for this patch
+        idx_start = [ranges[dim][i] for dim, i in enumerate(starts)]
+        idx_start2 = [ranges[dim][i] + half_overlap if ranges[dim][i] != 0 else 0 for dim, i in enumerate(starts)]
+        idx_start3 = [half_overlap if ranges[dim][i] != 0 else 0 for dim, i in enumerate(starts)]
+        idx_end = [min(start + size, shape[dim]) for start, size, dim in zip(idx_start, patch_size, range(len(shape)))]
+        idx_end2 = [
+            (start + size - half_overlap if start + size < shape[dim] else shape[dim])
+            for start, size, dim in zip(idx_start, patch_size, range(len(shape)))
+        ]
+        idx_end3 = [(-half_overlap if a != shape[dim] else None) for a, dim in zip(idx_end2, range(len(shape)))]
+
+        slices = tuple(slice(s, e) for s, e in zip(idx_start, idx_end))
+        slices2 = tuple(slice(s, e) for s, e in zip(idx_start2, idx_end2))
+        slices3 = tuple(slice(s, e) for s, e in zip(idx_start3, idx_end3))
+        print("sliding window", slices)
+        patch = arr[slices]
+        patch = fun(patch)
+        result[slices2] = patch[slices3]
+    return result
+
+
+# if __name__ == "__main__":
+# np.zeros((1, 2243, 472, 622))
+# x = sliding_nd_slices()
+# max_v=2243, (1, 2243, 472, 622)
+# image very large ((1, 2243, 472, 622)>1000); use sliding window patch_size=<generator object predict_single_npy_array.<locals>.<genexpr> at 0x7f89c12dfac0> [160, 192, 192]