Handling list of DICOM datasets reference when creating rtstruct

rt_utils/ds_helper.py

@@ -98,7 +98,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

rt_utils/image_helper.py

@@ -11,12 +11,14 @@
 from rt_utils.utils import ROIData, SOPClassUID
 
 
-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_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
  """

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:
@@ -112,10 +119,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")