Merge pull request #13 from qurit/11_colors

11 colors

Author: awtkns

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

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)]
+
+    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')

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:

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

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)