Add memory efficient smoothing

rt_utils/rtstruct.py

@@ -36,8 +36,7 @@ def add_roi(
  approximate_contours: bool = True,
  roi_generation_algorithm: Union[str, int] = 0,
  apply_smoothing: Union[str, None] = None, # strings can be "2d" or "3d"
- smoothing_parameters: Dict = smoothing.default_smoothing_parameters
-
+ smoothing_parameters: Union[Dict, None] = None
  ):
  """
  Add a ROI to the rtstruct given a 3D binary mask for the ROI's at each slice
@@ -46,12 +45,8 @@ def add_roi(
  If approximate_contours is set to False, no approximation will be done when generating contour data, leading to much larger amount of contour data
  """
  if apply_smoothing:
- if apply_smoothing == "3d":
- mask = smoothing.pipeline_3d(mask, **smoothing_parameters)
- elif apply_smoothing == "2d":
- mask = smoothing.pipeline_2d(mask, **smoothing_parameters)
- else:
- print("Invalid input. Use '2d' or '3d'. Otherwise leave as None")
+ mask = smoothing.pipeline(mask=mask, apply_smoothing=apply_smoothing,
+ smoothing_parameters=smoothing_parameters)
 
  ## If upscaled coords are given, they should be adjusted accordingly
  rows = self.series_data[0][0x00280010].value

rt_utils/smoothing.py

@@ -1,58 +1,132 @@
 import SimpleITK as sitk
 import numpy as np
 from scipy import ndimage, signal
-from typing import Union, Tuple
+from typing import Union, Tuple, Dict
+import logging
+
 default_smoothing_parameters = {
- "scaling_factor": (3, 3, 1),
- "sigma": 2,
- "threshold": 0.4,
- "kernel_size": (3, 3, 1),
- "iterations": 2
+ "iterations": 3,
+ "np_kron": {"scaling_factor": 2},
+ "ndimage_gaussian_filter": {"sigma": 2,
+ "radius": 3},
+ "threshold": {"threshold": 0.4},
+ "ndimage_median_filter": {"size": 3,
+ "mode": "nearest"}
 }
-def kron_upscale(mask: np.ndarray, scaling_factor: Tuple[int, ...]):
+
+def kron_upscale(mask: np.ndarray, **kwargs):
+ scaling_factor = (kwargs["scaling_factor"], kwargs["scaling_factor"], 1)
  return np.kron(mask, np.ones(scaling_factor))
 
+def gaussian_blur(mask: np.ndarray, **kwargs):
+ return ndimage.gaussian_filter(mask, **kwargs)
 
-def gaussian_blur(mask: np.ndarray, sigma: float):
- return ndimage.gaussian_filter(mask, sigma=sigma)
+def binary_threshold(mask: np.ndarray, **kwargs):
+ return mask > kwargs["threshold"]
 
+def median_filter(mask: np.ndarray, **kwargs):
+ return ndimage.median_filter(mask.astype(float), **kwargs)
 
-def binary_threshold(mask: np.ndarray, threshold: float):
- return mask > threshold
+def crop_mask(mask: np.ndarray):
+ three_d = (len(mask.shape) == 3)
+ if three_d:
+ x, y, z = np.nonzero(mask)
+ x_max = (x.max() + 1) if x.max() < mask.shape[0] else x.max()
+ y_max = (y.max() + 1) if y.max() < mask.shape[1] else y.max()
+ z_max = (z.max() + 1) if z.max() < mask.shape[2] else z.max()
+ bbox = np.array([x.min(), x_max, y.min(), y_max, z.min(), z_max])
 
-def median_filter(mask: np.ndarray, kernel_size: Union[int, Tuple[int, ...]]):
- return ndimage.median_filter(mask.astype(float), size=kernel_size, mode="nearest")
+ return mask[bbox[0]: bbox[1],
+ bbox[2]: bbox[3],
+ bbox[4]: bbox[5]], bbox
+ else:
+ x, y = np.nonzero(mask)
+ x_max = (x.max() + 1) if x.max() < mask.shape[0] else x.max()
+ y_max = (y.max() + 1) if y.max() < mask.shape[1] else y.max()
+ bbox = [x.min(), x_max, y.min(), y_max]
 
-def pipeline_3d(mask: np.ndarray,
- iterations: int,
- scaling_factor: int,
- sigma: float,
- threshold: float,
- kernel_size: Union[int, Tuple[int, ...]]):
- scaling_factor = (scaling_factor, scaling_factor, 1)
- for i in range(iterations):
- mask = kron_upscale(mask=mask, scaling_factor=scaling_factor)
- mask = gaussian_blur(mask=mask, sigma=sigma)
- mask = binary_threshold(mask=mask, threshold=threshold)
- mask = median_filter(mask=mask, kernel_size=kernel_size)
- mask = mask.astype(bool)
- return mask
+ return mask[bbox[0]: bbox[1],
+ bbox[2]: bbox[3]], bbox
+
+def restore_mask_dimensions(cropped_mask: np.ndarray, new_shape, bbox):
+ new_mask = np.zeros(new_shape)
+
+ new_mask[bbox[0]: bbox[1], bbox[2]: bbox[3], bbox[4]: bbox[5]] = cropped_mask
+ return new_mask.astype(bool)
+
+def iteration_2d(mask: np.ndarray, np_kron, ndimage_gaussian_filter, threshold, ndimage_median_filter):
+ cropped_mask = kron_upscale(mask=cropped_mask, **np_kron)
+
+ for z_idx in range(cropped_mask.shape[2]):
+ slice = cropped_mask[:, :, z_idx]
+ slice = gaussian_blur(mask=slice, **ndimage_gaussian_filter)
+ slice = binary_threshold(mask=slice, **threshold)
+ slice = median_filter(mask=slice, **ndimage_median_filter)
+
+ cropped_mask[:, :, z_idx] = slice
+ return cropped_mask
+
+def iteration_3d(mask: np.ndarray, np_kron, ndimage_gaussian_filter, threshold, ndimage_median_filter):
+ cropped_mask = kron_upscale(mask=mask, **np_kron)
+ cropped_mask = gaussian_blur(mask=cropped_mask, **ndimage_gaussian_filter)
+ cropped_mask = binary_threshold(mask=cropped_mask, **threshold)
+ cropped_mask = median_filter(mask=cropped_mask, **ndimage_median_filter)
 
-def pipeline_2d(mask: np.ndarray,
- iterations: int,
- scaling_factor: int,
- sigma: float,
- threshold: float,
- kernel_size: Union[int, Tuple[int, ...]]):
- scaling_factor = (scaling_factor, scaling_factor, 1)
+ return cropped_mask
+
+def pipeline(mask: np.ndarray,
+ apply_smoothing: str,
+ smoothing_parameters: Union[Dict, None]):
+
+ if not smoothing_parameters:
+ smoothing_parameters = default_smoothing_parameters
+
+ iterations = smoothing_parameters["iterations"]
+ np_kron = smoothing_parameters["np_kron"]
+ ndimage_gaussian_filter = smoothing_parameters["ndimage_gaussian_filter"]
+ threshold = smoothing_parameters["threshold"]
+ ndimage_median_filter = smoothing_parameters["ndimage_median_filter"]
+
+ logging.info(f"Original mask shape {mask.shape}")
+ logging.info(f"Cropping mask to non-zero")
+ cropped_mask, bbox = crop_mask(mask)
+ final_shape, final_bbox = get_final_mask_shape_and_bbox(mask=mask,
+ scaling_factor=np_kron["scaling_factor"],
+ iterations=iterations,
+ bbox=bbox)
+ logging.info(f"Final scaling with factor: {np_kron['scaling_factor']} in {iterations} iterations")
  for i in range(iterations):
- mask = kron_upscale(mask=mask, scaling_factor=scaling_factor)
- for z in range(mask.shape[2]):
- slice = mask[:, : , z]
- slice = gaussian_blur(mask=slice, sigma=sigma)
- slice = binary_threshold(mask=slice, threshold=threshold)
- slice = median_filter(mask=slice, kernel_size=kernel_size)
- mask[:, :, z] = slice
- mask = mask.astype(bool)
+ logging.info(f"Iteration {i} out of {iterations}")
+ logging.info(f"Applying filters")
+
+ if apply_smoothing == "2d":
+ mask = iteration_2d(mask,
+ np_kron=np_kron,
+ ndimage_gaussian_filter=ndimage_gaussian_filter,
+ threshold=threshold,
+ ndimage_median_filter=ndimage_median_filter)
+ elif apply_smoothing == "3d":
+ mask = iteration_3d(mask,
+ np_kron=np_kron,
+ ndimage_gaussian_filter=ndimage_gaussian_filter,
+ threshold=threshold,
+ ndimage_median_filter=ndimage_median_filter)
+ else:
+ raise Exception("Wrong dimension parameter. Use '2d' or '3d'.")
+
+ # Restore dimensions
+ logging.info("Restoring original mask shape")
+ mask = restore_mask_dimensions(cropped_mask, final_shape, final_bbox)
  return mask
 
+def get_final_mask_shape_and_bbox(mask, bbox, scaling_factor, iterations):
+ final_scaling_factor = pow(scaling_factor, iterations)
+
+ final_shape = np.array(mask.shape)
+ final_shape[:2] *= final_scaling_factor
+ bbox[:4] *= final_scaling_factor
+ logging.info("Final shape: ", final_shape)
+ logging.info("Final bbox: ", bbox)
+ return final_shape, bbox
+
+