Simplify algorithm to create contour from coordinates

rt_utils/ds_helper.py

@@ -1,5 +1,5 @@
 import datetime
-from typing import List
+from typing import List, Union
 
 from rt_utils.image_helper import get_contours_coords
 from rt_utils.utils import ROIData, SOPClassUID
@@ -9,6 +9,8 @@
 from pydicom.sequence import Sequence
 from pydicom.uid import ImplicitVRLittleEndian
 
+from utils import _flatten_lists
+
 """
 File contains helper methods that handles DICOM header creation/formatting
 """
@@ -177,7 +179,7 @@ def create_contour_sequence(roi_data: ROIData, series_data: List[Dataset]) -> Se
  return contour_sequence
 
 
-def create_roi_contour_from_coordinates(coordinates: List[List[float]], roi_data: ROIData, series_data) -> Dataset:
+def create_roi_contour_from_coordinates(coordinates: List[List[List[float]]], roi_data: ROIData, series_data) -> Dataset:
  roi_contour = Dataset()
  roi_contour.ROIDisplayColor = roi_data.color
  roi_contour.ContourSequence = create_contour_sequence_from_coordinates(coordinates, series_data)
@@ -186,54 +188,27 @@ def create_roi_contour_from_coordinates(coordinates: List[List[float]], roi_data
  return roi_contour
 
 
-def _find_closest_slice(series_slices: List[Dataset], z_coord: float) -> Dataset:
- return min(series_slices, key=lambda series_slice: abs(series_slice.SliceLocation - z_coord))
-
-
-def _flatten_lists(lists: List[List[float]]) -> List[float]:
- """Flatten the list [[1, 2, 3], [1, 2, 3] -> [1, 2, 3, 1, 2, 3]"""
- flatten_list = []
- for l in lists:
- flatten_list += l
-
- return flatten_list
-
-
 def create_contour_sequence_from_coordinates(coordinates: List[List[List[float]]], series_data: List[Dataset]) -> Sequence:
  """
  Iterate through each slice of the mask
  For each connected segment within a slice, create a contour
  """
- slice_position_to_contours_coords = {}
+ contour_sequence = Sequence()
 
  for contours_coords in coordinates:
  # Find the closest slice from the provided z coordinates
  closest_slice = _find_closest_slice(series_slices=series_data, z_coord=contours_coords[0][2])
 
  # Format contour coordinates in DICOM format [x1,y1,z1,x2,y2,z2,x3,y3,z3]
- contour = _flatten_lists(contours_coords)
+ contour_data = _flatten_lists(contours_coords)
 
- # Add contour to
- if closest_slice.SliceLocation not in slice_position_to_contours_coords:
- slice_position_to_contours_coords[closest_slice.SliceLocation] = []
-
- slice_position_to_contours_coords[closest_slice.SliceLocation].append(contour)
-
- # Making the DICOM contour sequence
- contour_sequence = Sequence()
-
- for z_position, slice_contours in slice_position_to_contours_coords.items():
- for series_slice in series_data:
- if series_slice.SliceLocation == z_position:
- for contour_data in slice_contours:
- contour = create_contour(series_slice, contour_data)
- contour_sequence.append(contour)
- continue
+ slice_contour = create_contour(closest_slice, contour_data)
+ contour_sequence.append(slice_contour)
 
  return contour_sequence
 
 
-def create_contour(series_slice: Dataset, contour_data: np.ndarray) -> Dataset:
+def create_contour(series_slice: Dataset, contour_data: Union[np.ndarray, List[float]]) -> Dataset:
  contour_image = Dataset()
  contour_image.ReferencedSOPClassUID = series_slice.SOPClassUID
  contour_image.ReferencedSOPInstanceUID = series_slice.SOPInstanceUID
@@ -280,3 +255,7 @@ def get_contour_sequence_by_roi_number(ds, roi_number):
  return Sequence()
 
  raise Exception(f"Referenced ROI number '{roi_number}' not found")
+
+
+def _find_closest_slice(series_slices: List[Dataset], z_coord: float) -> Dataset:
+ return min(series_slices, key=lambda series_slice: abs(series_slice.SliceLocation - z_coord))

rt_utils/utils.py

@@ -1,5 +1,4 @@
 from typing import List, Union
-from random import randrange
 
 import numpy as np
 from pydicom.uid import PYDICOM_IMPLEMENTATION_UID
@@ -126,3 +125,12 @@ def validate_roi_generation_algoirthm(self):
  type(self.roi_generation_algorithm)
  )
  )
+
+
+def _flatten_lists(lists: List[List[float]]) -> List[float]:
+ """Flatten the list [[1, 2, 3], [1, 2, 3] -> [1, 2, 3, 1, 2, 3]"""
+ flatten_list = []
+ for l in lists:
+ flatten_list += l
+
+ return flatten_list