From fd11e3ecd842ead64ad883f02c5f7a5fed8402cc Mon Sep 17 00:00:00 2001 From: gacou54 Date: Fri, 16 Feb 2024 13:12:14 -0500 Subject: [PATCH] Handling list of DICOM datasets reference when creating rtstruct --- rt_utils/ds_helper.py | 2 +- rt_utils/image_helper.py | 8 +++++--- rt_utils/rtstruct_builder.py | 4 ++-- tests/conftest.py | 24 +++++++++++++++++++----- tests/test_rtstruct_builder.py | 19 +++++++++++++------ 5 files changed, 40 insertions(+), 17 deletions(-) diff --git a/rt_utils/ds_helper.py b/rt_utils/ds_helper.py index 36c4db7..75e18b0 100644 --- a/rt_utils/ds_helper.py +++ b/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 diff --git a/rt_utils/image_helper.py b/rt_utils/image_helper.py index b030987..29b8745 100644 --- a/rt_utils/image_helper.py +++ b/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") diff --git a/rt_utils/rtstruct_builder.py b/rt_utils/rtstruct_builder.py index f92efb9..229e416 100644 --- a/rt_utils/rtstruct_builder.py +++ b/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 """ diff --git a/tests/conftest.py b/tests/conftest.py index ccd699e..25aa7bc 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -1,7 +1,11 @@ -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() @@ -9,13 +13,23 @@ 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") diff --git a/tests/test_rtstruct_builder.py b/tests/test_rtstruct_builder.py index e8cc856..3be6b90 100644 --- a/tests/test_rtstruct_builder.py +++ b/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")