Add smoothing pipeline

rt_utils/image_helper.py

@@ -1,7 +1,8 @@
 import os
-from typing import List
+from typing import List, Union
 from enum import IntEnum
 
+import SimpleITK
 import cv2 as cv
 import numpy as np
 from pydicom import dcmread
@@ -62,8 +63,9 @@ def get_contours_coords(roi_data: ROIData, series_data):
  # Get contours from mask
  contours, _ = find_mask_contours(mask_slice,
  roi_data.approximate_contours,
- smoothing_factor=roi_data.smoothing_factor,
  scaling_factor=roi_data.scaling_factor)
+ if not contours:
+ continue
  validate_contours(contours)
 
  # Format for DICOM
@@ -85,7 +87,8 @@ def get_contours_coords(roi_data: ROIData, series_data):
  return series_contours
 
 
-def find_mask_contours(mask: np.ndarray, approximate_contours: bool, smoothing_factor=1, scaling_factor=1):
+
+def find_mask_contours(mask: np.ndarray, approximate_contours: bool, scaling_factor: int):
  approximation_method = (
  cv.CHAIN_APPROX_SIMPLE if approximate_contours else cv.CHAIN_APPROX_NONE
  )
@@ -97,10 +100,10 @@ def find_mask_contours(mask: np.ndarray, approximate_contours: bool, smoothing_f
  contours
  ) # Open-CV updated contours to be a tuple so we convert it back into a list here
 
- assert smoothing_factor > 0
  for i, contour in enumerate(contours):
- contours[i] = [[(contour[i][0][0]/scaling_factor), (contour[i][0][1]/scaling_factor)] for i in range(0, len(contour), smoothing_factor)]
- hierarchy = hierarchy[0] # Format extra array out of data
+ contours[i] = [[(contour[i][0][0] / scaling_factor), (contour[i][0][1] / scaling_factor)] for i in
+ range(0, len(contour))]
+ # hierarchy = hierarchy[0] # Format extra array out of data
 
  return contours, hierarchy
 
@@ -131,7 +134,7 @@ def create_pin_hole_mask(mask: np.ndarray, approximate_contours: bool):
 
 
 def draw_line_upwards_from_point(
- mask: np.ndarray, start, fill_value: int
+  mask: np.ndarray, start, fill_value: int
 ) -> np.ndarray:
  line_width = 2
  end = (start[0], start[1] - 1)
@@ -201,7 +204,7 @@ def get_patient_to_pixel_transformation_matrix(series_data):
 
 
 def apply_transformation_to_3d_points(
- points: np.ndarray, transformation_matrix: np.ndarray
+  points: np.ndarray, transformation_matrix: np.ndarray
 ):
  """
  * Augment each point with a '1' as the fourth coordinate to allow translation
@@ -224,7 +227,7 @@ def get_slice_directions(series_slice: Dataset):
  slice_direction = np.cross(row_direction, column_direction)
 
  if not np.allclose(
- np.dot(row_direction, column_direction), 0.0, atol=1e-3
+  np.dot(row_direction, column_direction), 0.0, atol=1e-3
  ) or not np.allclose(np.linalg.norm(slice_direction), 1.0, atol=1e-3):
  raise Exception("Invalid Image Orientation (Patient) attribute")
 
@@ -268,7 +271,7 @@ def get_slice_contour_data(series_slice: Dataset, contour_sequence: Sequence):
 
 
 def get_slice_mask_from_slice_contour_data(
- series_slice: Dataset, slice_contour_data, transformation_matrix: np.ndarray
+  series_slice: Dataset, slice_contour_data, transformation_matrix: np.ndarray
 ):
  # Go through all contours in a slice, create polygons in correct space and with a correct format 
  # and append to polygons array (appropriate for fillPoly) 
@@ -280,9 +283,10 @@ def get_slice_mask_from_slice_contour_data(
  polygon = np.array(polygon).squeeze()
  polygons.append(polygon)
  slice_mask = create_empty_slice_mask(series_slice).astype(np.uint8)
- cv.fillPoly(img=slice_mask, pts = polygons, color = 1)
+ cv.fillPoly(img=slice_mask, pts=polygons, color=1)
  return slice_mask
 
+
 def create_empty_series_mask(series_data):
  ref_dicom_image = series_data[0]
  mask_dims = (

rt_utils/rtstruct.py

@@ -1,11 +1,11 @@
-from typing import List, Union
+from typing import List, Union, Dict
 
 import numpy as np
 from pydicom.dataset import FileDataset
 
 from rt_utils.utils import ROIData
-from . import ds_helper, image_helper
-
+from . import ds_helper, image_helper, smoothing
+from typing import Tuple
 
 class RTStruct:
  """
@@ -35,14 +35,23 @@ def add_roi(
  use_pin_hole: bool = False,
  approximate_contours: bool = True,
  roi_generation_algorithm: Union[str, int] = 0,
- smoothing_factor: int = 1,
+ apply_smoothing: Union[str, None] = None, # strings can be "2d" or "3d"
+ smoothing_parameters: Dict = smoothing.default_smoothing_parameters
+
  ):
  """
  Add a ROI to the rtstruct given a 3D binary mask for the ROI's at each slice
  Optionally input a color or name for the ROI
  If use_pin_hole is set to true, will cut a pinhole through ROI's with holes in them so that they are represented with one contour
  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")
 
  ## If upscaled coords are given, they should be adjusted accordingly
  rows = self.series_data[0][0x00280010].value
@@ -61,8 +70,7 @@ def add_roi(
  use_pin_hole,
  approximate_contours,
  roi_generation_algorithm,
- smoothing_factor,
- scaling_factor,
+ scaling_factor
  )
 
  self.ds.ROIContourSequence.append(

rt_utils/smoothing.py

@@ -0,0 +1,58 @@
+import SimpleITK as sitk
+import numpy as np
+from scipy import ndimage, signal
+from typing import Union, Tuple
+default_smoothing_parameters = {
+ "scaling_factor": (3, 3, 1),
+ "sigma": 2,
+ "threshold": 0.4,
+ "kernel_size": (3, 3, 1),
+ "iterations": 2
+}
+def kron_upscale(mask: np.ndarray, scaling_factor: Tuple[int, ...]):
+ return np.kron(mask, np.ones(scaling_factor))
+
+
+def gaussian_blur(mask: np.ndarray, sigma: float):
+ return ndimage.gaussian_filter(mask, sigma=sigma)
+
+
+def binary_threshold(mask: np.ndarray, threshold: float):
+ return mask > threshold
+
+def median_filter(mask: np.ndarray, kernel_size: Union[int, Tuple[int, ...]]):
+ return ndimage.median_filter(mask.astype(float), size=kernel_size, mode="nearest")
+
+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
+
+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)
+ 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)
+ return mask
+

rt_utils/utils.py

@@ -1,4 +1,4 @@
-from typing import List, Union
+from typing import List, Union, Tuple
 from random import randrange
 from pydicom.uid import PYDICOM_IMPLEMENTATION_UID
 from dataclasses import dataclass
@@ -50,10 +50,9 @@ class ROIData:
  use_pin_hole: bool = False
  approximate_contours: bool = True
  roi_generation_algorithm: Union[str, int] = 0
- smoothing_factor: int = 1
  scaling_factor: int = 1
-
-
+ smooth_radius: Union[int, None] = None
+ smooth_scale: Union[int, None] = None
 
 
  def __post_init__(self):