diff --git a/README.md b/README.md index 7920fb5..c7cfc45 100644 --- a/README.md +++ b/README.md @@ -55,6 +55,25 @@ rtstruct.add_roi( name="RT-Utils ROI!" ) +# Add another ROI from coordinates +rtstruct.add_roi_from_coordinates( + coordinates=[ + [ + # Example of One contour on one slice + [-20.0, -170.0, -559.0], + [30.0, -170.0, -559.0], + [30.0, -110.0, -559.0], + [-20.0, -110.0, -559.0], + ], + [ + [-20.0, -170.0, -562.4], + [30.0, -170.0, -562.4], + [30.0, -110.0, -562.4], + [-20.0, -110.0, -562.4], + ] + ] +) + rtstruct.save('new-rt-struct') ``` diff --git a/rt_utils/ds_helper.py b/rt_utils/ds_helper.py index 36c4db7..0654233 100644 --- a/rt_utils/ds_helper.py +++ b/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.ImagePositionPatient[2] - z_coord)) diff --git a/rt_utils/image_helper.py b/rt_utils/image_helper.py index b030987..e02df28 100644 --- a/rt_utils/image_helper.py +++ b/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") diff --git a/rt_utils/rtstruct.py b/rt_utils/rtstruct.py index dfe82be..4ce4749 100644 --- a/rt_utils/rtstruct.py +++ b/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( 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/rt_utils/utils.py b/rt_utils/utils.py index f04089e..e2ca56e 100644 --- a/rt_utils/utils.py +++ b/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 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..8835f6c 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: @@ -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")