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")