diff --git a/rt_utils/rtstruct.py b/rt_utils/rtstruct.py index f63a39d..cdf1296 100644 --- a/rt_utils/rtstruct.py +++ b/rt_utils/rtstruct.py @@ -1,4 +1,4 @@ -from typing import List +from typing import List, Union import numpy as np from pydicom.dataset import FileDataset @@ -24,7 +24,7 @@ def set_series_description(self, description: str): self.ds.SeriesDescription = description - def add_roi(self, mask: np.ndarray, color=None, name=None, description='', use_pin_hole=False): + def add_roi(self, mask: np.ndarray, color: Union[str, List[int]] = None, name: str = None, description: str = '', use_pin_hole: bool = False): """ Add a ROI to the rtstruct given a 3D binary mask for the ROI's at each slice Optionally input a color or name for the ROI diff --git a/rt_utils/utils.py b/rt_utils/utils.py index 48fdd4f..6392cf6 100644 --- a/rt_utils/utils.py +++ b/rt_utils/utils.py @@ -2,12 +2,38 @@ from pydicom.uid import PYDICOM_ROOT_UID from dataclasses import dataclass -class SOPClassUID(): - RTSTRUCT_IMPLEMENTATION_CLASS = PYDICOM_ROOT_UID # TODO find out if this is ok +COLOR_PALETTE= [ + [255, 0, 255], + [0, 235, 235], + [255, 255, 0], + [255, 0, 0], + [0, 132, 255], + [0, 240, 0], + [255, 175, 0], + [0, 208, 255], + [180, 255, 105], + [255, 20, 147], + [160, 32, 240], + [0, 255, 127], + [255, 114, 0], + [64, 224, 208], + [0, 178, 47], + [220, 20, 60], + [238, 130, 238], + [218, 165, 32], + [255, 140, 190], + [0, 0, 255], + [255, 225, 0] +] + + +class SOPClassUID: + RTSTRUCT_IMPLEMENTATION_CLASS = PYDICOM_ROOT_UID # TODO find out if this is ok CT_IMAGE_STORAGE = '1.2.840.10008.5.1.4.1.1.2' DETACHED_STUDY_MANAGEMENT = '1.2.840.10008.3.1.2.3.1' RTSTRUCT = '1.2.840.10008.5.1.4.1.1.481.3' + @dataclass class ROIData: """Data class to easily pass ROI data to helper methods.""" @@ -20,15 +46,42 @@ class ROIData: use_pin_hole: bool = False def __post_init__(self): + self.validate_color() self.add_default_values() def add_default_values(self): - if self.color == None: - self.color = self.get_random_colour() + if self.color is None: + self.color = COLOR_PALETTE[(self.number - 1) % len(COLOR_PALETTE)] - if self.name == None: + if self.name is None: self.name = f"ROI-{self.number}" - - def get_random_colour(self): - max = 256 - return [randrange(max), randrange(max), randrange(max)] \ No newline at end of file + + def validate_color(self): + if self.color is None: + return + + # Validating list eg: [0, 0, 0] + if type(self.color) is list: + if len(self.color) != 3: + raise ValueError(f'{self.color} is an invalid color for an ROI') + for c in self.color: + try: + assert 0 <= c <= 255 + except: + raise ValueError(f'{self.color} is an invalid color for an ROI') + + else: + self.color: str = str(self.color) + self.color = self.color.strip('#') + + # fff -> ffffff + if len(self.color) == 3: + self.color = ''.join([x*2 for x in self.color]) + + if not len(self.color) == 6: + raise ValueError(f'{self.color} is an invalid color for an ROI') + + try: + self.color = [int(self.color[i:i+2], 16) for i in (0, 2, 4)] + except Exception as e: + raise ValueError(f'{self.color} is an invalid color for an ROI') diff --git a/setup.py b/setup.py index 7a84a10..3860570 100644 --- a/setup.py +++ b/setup.py @@ -1,6 +1,6 @@ import setuptools -VERSION = '1.0.5' +VERSION = '1.1.0' with open("README.md", "r", encoding="utf-8") as fh: long_description = fh.read() with open('requirements.txt') as f: diff --git a/tests/test_rtstruct_builder.py b/tests/test_rtstruct_builder.py index bd67656..d86bd42 100644 --- a/tests/test_rtstruct_builder.py +++ b/tests/test_rtstruct_builder.py @@ -55,7 +55,7 @@ def test_add_valid_roi(new_rtstruct: RTStruct): assert len(new_rtstruct.ds.RTROIObservationsSequence) == 0 NAME = "Test ROI" - COLOR = [123, 321, 456] + COLOR = [123, 123, 232] mask = get_empty_mask(new_rtstruct) mask[50:100, 50:100, 0] = 1 @@ -163,4 +163,3 @@ def get_empty_mask(rtstruct) -> np.ndarray: mask_dims = (int(ref_dicom_image.Columns), int(ref_dicom_image.Rows), len(rtstruct.series_data)) mask = np.zeros(mask_dims) return mask.astype(bool) - diff --git a/tests/test_utils.py b/tests/test_utils.py new file mode 100644 index 0000000..768fa8d --- /dev/null +++ b/tests/test_utils.py @@ -0,0 +1,53 @@ +import pytest + +from rt_utils.utils import COLOR_PALETTE +from tests.test_rtstruct_builder import get_empty_mask + + +VALID_COLORS = [ + ('fff', [255, 255, 255]), + ('#fff', [255, 255, 255]), + (None, COLOR_PALETTE[0]), + (COLOR_PALETTE[1], COLOR_PALETTE[1]), + ('#696969', [105, 105, 105]), + ('a81414', [168, 20, 20]), + ('#000', [0, 0, 0]), +] + +INVALID_COLORS = [ + ('GGG', ValueError), + ('red', ValueError), + ('22', ValueError), + ('[]', ValueError), + ([], ValueError), + ([24, 34], ValueError), + ([24, 34, 454], ValueError), + ([0, 344, 0], ValueError), + ('a8141', ValueError), + ('a814111', ValueError), + (KeyboardInterrupt, ValueError), +] + + +@pytest.mark.parametrize('color', VALID_COLORS) +def test_mask_colors(new_rtstruct, color): + color_in, color_out = color + + name = "Test ROI" + mask = get_empty_mask(new_rtstruct) + mask[50:100, 50:100, 0] = 1 + + new_rtstruct.add_roi(mask, color=color_in, name=name) + assert new_rtstruct.ds.ROIContourSequence[0].ROIDisplayColor == color_out + + +@pytest.mark.parametrize('color', INVALID_COLORS) +def test_mask_colors_fail(new_rtstruct, color): + color_in, err = color + + name = "Test ROI" + mask = get_empty_mask(new_rtstruct) + mask[50:100, 50:100, 0] = 1 + + with pytest.raises(err): + new_rtstruct.add_roi(mask, color=color_in, name=name)