Merge pull request #1 from paradim-platform/feature/from-coordinates

Add roi from coordinates

rt_utils/ds_helper.py

@@ -1,4 +1,6 @@
 import datetime
+from typing import List, Union
+
 from rt_utils.image_helper import get_contours_coords
 from rt_utils.utils import ROIData, SOPClassUID
 import numpy as np
@@ -7,6 +9,8 @@
 from pydicom.sequence import Sequence
 from pydicom.uid import ImplicitVRLittleEndian
 
+from rt_utils.utils import _flatten_lists
+
 """
 File contains helper methods that handles DICOM header creation/formatting
 """
@@ -98,7 +102,7 @@ def add_patient_information(ds: FileDataset, series_data):
 
 def add_refd_frame_of_ref_sequence(ds: FileDataset, series_data):
  refd_frame_of_ref = Dataset()
- refd_frame_of_ref.FrameOfReferenceUID =  getattr(series_data[0], 'FrameOfReferenceUID', generate_uid())
+ refd_frame_of_ref.FrameOfReferenceUID = getattr(series_data[0], 'FrameOfReferenceUID', generate_uid())
  refd_frame_of_ref.RTReferencedStudySequence = create_frame_of_ref_study_sequence(series_data)
 
  # Add to sequence
@@ -157,7 +161,7 @@ def create_roi_contour(roi_data: ROIData, series_data) -> Dataset:
  return roi_contour
 
 
-def create_contour_sequence(roi_data: ROIData, series_data) -> Sequence:
+def create_contour_sequence(roi_data: ROIData, series_data: List[Dataset]) -> Sequence:
  """
  Iterate through each slice of the mask
  For each connected segment within a slice, create a contour
@@ -175,7 +179,36 @@ def create_contour_sequence(roi_data: ROIData, series_data) -> Sequence:
  return contour_sequence
 
 
-def create_contour(series_slice: Dataset, contour_data: np.ndarray) -> 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)
+ roi_contour.ReferencedROINumber = str(roi_data.number)
+
+ return roi_contour
+
+
+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
+ """
+ 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_data = _flatten_lists(contours_coords)
+
+ slice_contour = create_contour(closest_slice, contour_data)
+ contour_sequence.append(slice_contour)
+
+ return contour_sequence
+
+
+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
@@ -222,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/image_helper.py

@@ -8,15 +8,17 @@
 from pydicom.dataset import Dataset
 from pydicom.sequence import Sequence
 
-from rt_utils.utils import ROIData, SOPClassUID
+from rt_utils.utils import ROIData
 
 
-def load_sorted_image_series(dicom_series_path: str):
+def load_sorted_image_series(dicom_series_path: str | List[Dataset]) -> List[Dataset]:
  """
  File contains helper methods for loading / formatting DICOM images and contours
  """
-
- series_data = load_dcm_images_from_path(dicom_series_path)
+ if isinstance(dicom_series_path, str):
+ series_data = load_dcm_images_from_path(dicom_series_path)
+ else:
+ series_data = dicom_series_path
 
  if len(series_data) == 0:
  raise Exception("No DICOM Images found in input path")

rt_utils/rtstruct.py

@@ -68,6 +68,39 @@ def add_roi(
  ds_helper.create_rtroi_observation(roi_data)
  )
 
+ def add_roi_from_coordinates(
+ self,
+ coordinates: List[List[List[float]]],
+ color: Union[str, List[int]] = None,
+ name: str = None,
+ description: str = "",
+ use_pin_hole: bool = False,
+ approximate_contours: bool = True,
+ roi_generation_algorithm: Union[str, int] = 0,
+ ):
+ roi_number = len(self.ds.StructureSetROISequence) + 1
+ roi_data = ROIData(
+ np.zeros(1), # Fake mask since we do not use it because we use coordinates
+ color,
+ roi_number,
+ name,
+ self.frame_of_reference_uid,
+ description,
+ use_pin_hole,
+ approximate_contours,
+ roi_generation_algorithm,
+ )
+
+ self.ds.ROIContourSequence.append(
+ ds_helper.create_roi_contour_from_coordinates(coordinates, roi_data, self.series_data)
+ )
+ self.ds.StructureSetROISequence.append(
+ ds_helper.create_structure_set_roi(roi_data)
+ )
+ self.ds.RTROIObservationsSequence.append(
+ ds_helper.create_rtroi_observation(roi_data)
+ )
+
  def validate_mask(self, mask: np.ndarray) -> bool:
  if mask.dtype != bool:
  raise RTStruct.ROIException(

rt_utils/rtstruct_builder.py

@@ -15,7 +15,7 @@ class RTStructBuilder:
  """
 
  @staticmethod
- def create_new(dicom_series_path: str) -> RTStruct:
+ def create_new(dicom_series_path: str | List[Dataset]) -> RTStruct:
  """
  Method to generate a new rt struct from a DICOM series
  """
@@ -25,7 +25,7 @@ def create_new(dicom_series_path: str) -> RTStruct:
  return RTStruct(series_data, ds)
 
  @staticmethod
- def create_from(dicom_series_path: str, rt_struct_path: str, warn_only: bool = False) -> RTStruct:
+ def create_from(dicom_series_path: str | List[Dataset], rt_struct_path: str, warn_only: bool = False) -> RTStruct:
  """
  Method to load an existing rt struct, given related DICOM series and existing rt struct
  """

rt_utils/utils.py

@@ -1,5 +1,6 @@
 from typing import List, Union
-from random import randrange
+
+import numpy as np
 from pydicom.uid import PYDICOM_IMPLEMENTATION_UID
 from dataclasses import dataclass
 
@@ -41,8 +42,7 @@ class SOPClassUID:
 @dataclass
 class ROIData:
  """Data class to easily pass ROI data to helper methods."""
-
- mask: str
+ mask: np.ndarray
  color: Union[str, List[int]]
  number: int
  name: str
@@ -125,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

setup.py

@@ -15,7 +15,7 @@
  long_description=long_description,
  long_description_content_type="text/markdown",
  url="https://github.com/qurit/rtutils",
- package_dir={'':"rt_utils"},
+ package_dir={'rt_utils': "rt_utils"},
  packages=setuptools.find_packages("rt_utils", exclude="tests"),
  keywords=["RTStruct", "Dicom", "Pydicom"],
  classifiers=[

tests/conftest.py

@@ -1,21 +1,35 @@
-from rt_utils.rtstruct import RTStruct
-import pytest
 import os
-from rt_utils import RTStructBuilder
+from typing import List
+
+import pydicom
+import pytest
+
+from rt_utils import RTStructBuilder, image_helper
+from rt_utils.rtstruct import RTStruct
 
 
 @pytest.fixture()
 def series_path() -> str:
  return get_and_test_series_path("mock_data")
 
 
+@pytest.fixture()
+def series_datasets(series_path) -> List[pydicom.Dataset]:
+ return image_helper.load_dcm_images_from_path(series_path)
+
+
+@pytest.fixture()
+def rtstruct_path(series_path) -> str:
+ return os.path.join(series_path, "rt.dcm")
+
+
 @pytest.fixture()
 def new_rtstruct() -> RTStruct:
  return get_rtstruct("mock_data")
 
 
 @pytest.fixture()
-def oriented_series_path() -> RTStruct:
+def oriented_series_path() -> str:
  return get_and_test_series_path("oriented_data")
 
 
@@ -25,7 +39,7 @@ def oriented_rtstruct() -> RTStruct:
 
 
 @pytest.fixture()
-def one_slice_series_path() -> RTStruct:
+def one_slice_series_path() -> str:
  return get_and_test_series_path("one_slice_data")
 
 

tests/test_rtstruct_builder.py

@@ -1,9 +1,9 @@
+from _pytest.fixtures import FixtureRequest
+
 from rt_utils.rtstruct import RTStruct
 import pytest
 import os
 from rt_utils import RTStructBuilder
-from rt_utils.utils import SOPClassUID
-from rt_utils import image_helper
 from pydicom.dataset import validate_file_meta
 import numpy as np
 
@@ -15,6 +15,13 @@ def test_create_from_empty_series_dir():
  RTStructBuilder.create_new(empty_dir_path)
 
 
+@pytest.mark.parametrize('series_fixture', ['series_path', 'series_datasets'])
+def test_create_new_datasets(series_fixture, request: FixtureRequest):
+ rtstruct = RTStructBuilder.create_new(request.getfixturevalue(series_fixture))
+
+ assert len(rtstruct.ds.ReferencedFrameOfReferenceSequence[0].RTReferencedStudySequence) != 0
+
+
 def test_only_images_loaded_into_series_data(new_rtstruct: RTStruct):
  assert len(new_rtstruct.series_data) > 0
  for ds in new_rtstruct.series_data:
@@ -74,6 +81,41 @@ def test_add_valid_roi(new_rtstruct: RTStruct):
  assert new_rtstruct.get_roi_names() == [NAME]
 
 
+def test_add_valid_roi_from_coordinates(new_rtstruct: RTStruct):
+ assert new_rtstruct.get_roi_names() == []
+ assert len(new_rtstruct.ds.ROIContourSequence) == 0
+ assert len(new_rtstruct.ds.StructureSetROISequence) == 0
+ assert len(new_rtstruct.ds.RTROIObservationsSequence) == 0
+
+ NAME = "Test ROI"
+ COLOR = [123, 123, 232]
+ coordinates = [
+ [
+ [-100, -100, 60],
+ [-100, -75, 60],
+ [-75, -75, 60],
+ [-75, -100, 60]
+ ],
+ [
+ [-90, -90, 65],
+ [-90, -65, 65],
+ [-65, -65, 65],
+ [-65, -90, 65],
+ ]
+ ]
+
+ new_rtstruct.add_roi_from_coordinates(coordinates, color=COLOR, name=NAME)
+
+ assert len(new_rtstruct.ds.ROIContourSequence) == 1
+ assert (
+ len(new_rtstruct.ds.ROIContourSequence[0].ContourSequence) == 2
+ ) # 2 contour on to slice were added
+ assert len(new_rtstruct.ds.StructureSetROISequence) == 1
+ assert len(new_rtstruct.ds.RTROIObservationsSequence) == 1
+ assert new_rtstruct.ds.ROIContourSequence[0].ROIDisplayColor == COLOR
+ assert new_rtstruct.get_roi_names() == [NAME]
+
+
 def test_get_invalid_roi_mask_by_name(new_rtstruct: RTStruct):
  assert new_rtstruct.get_roi_names() == []
  with pytest.raises(RTStruct.ROIException):
@@ -112,10 +154,10 @@ def test_non_existant_referenced_study_sequence(series_path):
  )
 
 
-def test_loading_valid_rt_struct(series_path):
- valid_rt_struct_path = os.path.join(series_path, "rt.dcm")
- assert os.path.exists(valid_rt_struct_path)
- rtstruct = RTStructBuilder.create_from(series_path, valid_rt_struct_path)
+@pytest.mark.parametrize('series_fixture', ['series_path', 'series_datasets'])
+def test_loading_valid_rt_struct(rtstruct_path, series_fixture, request: FixtureRequest):
+ assert os.path.exists(rtstruct_path)
+ rtstruct = RTStructBuilder.create_from(request.getfixturevalue(series_fixture), rtstruct_path)
 
  # Tests existing values predefined in the file are found
  assert hasattr(rtstruct.ds, "ROIContourSequence")