diff --git a/rt_utils/rtstruct.py b/rt_utils/rtstruct.py index dfe82be..bad05b9 100644 --- a/rt_utils/rtstruct.py +++ b/rt_utils/rtstruct.py @@ -1,9 +1,12 @@ +from io import BytesIO from typing import List, Union import numpy as np +import pydicom from pydicom.dataset import FileDataset from rt_utils.utils import ROIData + from . import ds_helper, image_helper @@ -133,6 +136,15 @@ def save(self, file_path: str): file.close() except OSError: raise Exception(f"Cannot write to file path '{file_path}'") + + def save_to_memory(self): + """ + Saves the RTStruct to a BytesIO stream and returns it. + """ + buffer = BytesIO() + pydicom.dcmwrite(buffer, self.ds) + buffer.seek(0) # Rewind the buffer for reading + return buffer class ROIException(Exception): """ diff --git a/rt_utils/rtstruct_builder.py b/rt_utils/rtstruct_builder.py index f92efb9..78395ef 100644 --- a/rt_utils/rtstruct_builder.py +++ b/rt_utils/rtstruct_builder.py @@ -1,10 +1,11 @@ +import warnings from typing import List + from pydicom.dataset import Dataset from pydicom.filereader import dcmread -import warnings - from rt_utils.utils import SOPClassUID + from . import ds_helper, image_helper from .rtstruct import RTStruct @@ -23,6 +24,17 @@ def create_new(dicom_series_path: str) -> RTStruct: series_data = image_helper.load_sorted_image_series(dicom_series_path) ds = ds_helper.create_rtstruct_dataset(series_data) return RTStruct(series_data, ds) + + @staticmethod + def create_new_from_memory(dicom_datasets: List[Dataset]) -> RTStruct: + """ + Method to generate a new rt struct from a DICOM series in memory already loaded by pydicom + """ + # dicom_datasets are sorted by InstanceNumber + dicom_datasets.sort(key=lambda x: x.InstanceNumber if hasattr(x, 'InstanceNumber') else 0) + + ds = ds_helper.create_rtstruct_dataset(dicom_datasets) + return RTStruct(dicom_datasets, ds) @staticmethod def create_from(dicom_series_path: str, rt_struct_path: str, warn_only: bool = False) -> RTStruct: @@ -37,6 +49,17 @@ def create_from(dicom_series_path: str, rt_struct_path: str, warn_only: bool = F # TODO create new frame of reference? Right now we assume the last frame of reference created is suitable return RTStruct(series_data, ds) + + @staticmethod + def create_from_memory(dicom_datasets: List[Dataset], existing_rt_struct: Dataset, warn_only: bool = False) -> RTStruct : + """ + Method to generate a new rt struct from a DICOM series and a RTStruct in memory already loaded by pydicom + """ + dicom_datasets.sort(key=lambda x: x.InstanceNumber if hasattr(x, 'InstanceNumber') else 0) + + RTStructBuilder.validate_rtstruct(existing_rt_struct) + RTStructBuilder.validate_rtstruct_series_references(existing_rt_struct, dicom_datasets, warn_only) + return RTStruct(dicom_datasets, existing_rt_struct) @staticmethod def validate_rtstruct(ds: Dataset): diff --git a/tests/test_rtstruct_builder.py b/tests/test_rtstruct_builder.py index e8cc856..d707b22 100644 --- a/tests/test_rtstruct_builder.py +++ b/tests/test_rtstruct_builder.py @@ -6,7 +6,7 @@ from rt_utils import image_helper from pydicom.dataset import validate_file_meta import numpy as np - +import pydicom def test_create_from_empty_series_dir(): empty_dir_path = os.path.join(os.path.dirname(__file__), "empty") @@ -251,3 +251,27 @@ def get_empty_mask(rtstruct) -> np.ndarray: ) mask = np.zeros(mask_dims) return mask.astype(bool) + +def load_dicom_files_in_memory(file_paths): + dicom_datasets = [] + for path in file_paths: + with open(path, 'rb') as f: + ds = pydicom.dcmread(f) + dicom_datasets.append(ds) + return dicom_datasets + +def test_create_rtstruct_from_memory(): + print(os.getcwd()) + dicom_files = load_dicom_files_in_memory(['tests/mock_data/ct_1.dcm', 'tests/mock_data/ct_2.dcm']) + rtstruct = RTStructBuilder.create_new_from_memory(dicom_files) + assert rtstruct is not None + assert len(rtstruct.series_data) == len(dicom_files) + +def test_save_rtstruct_to_memory(): + dicom_files = load_dicom_files_in_memory(['tests/mock_data/ct_1.dcm', 'tests/mock_data/ct_2.dcm']) + rtstruct = RTStructBuilder.create_new_from_memory(dicom_files) + buffer = rtstruct.save_to_memory() + assert buffer is not None + buffer.seek(0) # Reset the buffer to the beginning for reading + loaded_ds = pydicom.dcmread(buffer) + assert loaded_ds is not None