diff --git a/README.md b/README.md index 52c1dc6..98dd14c 100644 --- a/README.md +++ b/README.md @@ -66,6 +66,14 @@ rtstruct.add_roi( name="RT-Utils ROI!" ) +# By default contours pass through the centres of boundary voxels, preserving +# historical behavior. To place contours on the outer voxel edges instead: +rtstruct.add_roi( + mask=MASK_FROM_ML_MODEL, + name="Voxel-edge ROI", + contour_mode="voxel_edge" +) + rtstruct.save('new-rt-struct') ``` @@ -183,4 +191,3 @@ If you are incorporating RT-Utils into your projects, kindly include the followi [![DOI](https://joss.theoj.org/papers/10.21105/joss.07361/status.svg)](https://doi.org/10.21105/joss.07361) Read the full paper: Asim Shrestha, Adam Watkins, Fereshteh Yousefirizi, Arman Rahmim, and Carlos Uribe [RT-utils: A Minimal Python Library for RT-struct Manipulation](https://joss.theoj.org/papers/10.21105/joss.07361) - diff --git a/rt_utils/image_helper.py b/rt_utils/image_helper.py index 5770b42..1ac27c0 100644 --- a/rt_utils/image_helper.py +++ b/rt_utils/image_helper.py @@ -60,7 +60,11 @@ def get_contours_coords(roi_data: ROIData, series_data): mask_slice = create_pin_hole_mask(mask_slice, roi_data.approximate_contours) # Get contours from mask - contours, _ = find_mask_contours(mask_slice, roi_data.approximate_contours) + contours, _ = find_mask_contours( + mask_slice, + roi_data.approximate_contours, + contour_mode=roi_data.contour_mode, + ) validate_contours(contours) # Format for DICOM @@ -82,7 +86,21 @@ def get_contours_coords(roi_data: ROIData, series_data): return series_contours -def find_mask_contours(mask: np.ndarray, approximate_contours: bool): +def find_mask_contours( + mask: np.ndarray, + approximate_contours: bool, + contour_mode: str = "voxel_center", +): + if contour_mode == "voxel_edge": + contours = [polygon.tolist() for polygon in mask_to_edge_polygons(mask)] + hierarchy = np.full((len(contours), 4), -1, dtype=np.int32) + return contours, hierarchy + if contour_mode != "voxel_center": + raise ValueError( + f"Invalid contour mode '{contour_mode}'. Expected one of " + "['voxel_center', 'voxel_edge']." + ) + approximation_method = ( cv.CHAIN_APPROX_SIMPLE if approximate_contours else cv.CHAIN_APPROX_NONE ) @@ -100,6 +118,87 @@ def find_mask_contours(mask: np.ndarray, approximate_contours: bool): return contours, hierarchy +def mask_to_edge_polygons(mask: np.ndarray) -> List[np.ndarray]: + """Trace polygons on the outer edges of foreground mask voxels. + + Returned vertices use OpenCV's ``(x, y)`` ordering and half-integer + coordinates. Each foreground pixel is treated as a unit square centred on + its integer index, so the result describes the complete voxel footprint. + """ + foreground = np.ascontiguousarray(mask, dtype=bool) + if foreground.ndim != 2: + raise ValueError("Mask must be two dimensional") + + rows, columns = foreground.shape + padded = np.zeros((rows + 2, columns + 2), dtype=bool) + padded[1:-1, 1:-1] = foreground + cells = padded[1:-1, 1:-1] + + neighbours = ( + padded[:-2, 1:-1], + padded[1:-1, 2:], + padded[2:, 1:-1], + padded[1:-1, :-2], + ) + offsets = ( + (-0.5, -0.5, 0.5, -0.5), + (0.5, -0.5, 0.5, 0.5), + (0.5, 0.5, -0.5, 0.5), + (-0.5, 0.5, -0.5, -0.5), + ) + + edges = {} + for neighbour, (start_x, start_y, end_x, end_y) in zip(neighbours, offsets): + for row, column in zip(*np.nonzero(cells & ~neighbour)): + start = (column + start_x, row + start_y) + end = (column + end_x, row + end_y) + edges.setdefault(start, []).append(end) + + polygons = [] + while edges: + start = next(iter(edges)) + current = start + points = [start] + while True: + ends = edges.get(current) + if not ends: + break + if len(ends) > 1: + ends.sort( + key=lambda end: np.arctan2( + end[1] - current[1], end[0] - current[0] + ) + ) + current = ends.pop(0) + if not ends: + del edges[points[-1]] + if current == start: + break + points.append(current) + + polygon = remove_collinear_points(np.asarray(points, dtype=float)) + if len(polygon) >= 3: + polygons.append(polygon) + return polygons + + +def remove_collinear_points(polygon: np.ndarray) -> np.ndarray: + """Remove intermediate vertices from straight polygon segments.""" + if len(polygon) < 3: + return polygon + + keep = [] + for index, point in enumerate(polygon): + previous = polygon[(index - 1) % len(polygon)] + following = polygon[(index + 1) % len(polygon)] + first = point - previous + second = following - point + cross_product = first[0] * second[1] - first[1] * second[0] + if abs(cross_product) > 1e-9 or float(first @ second) < 0: + keep.append(point) + return np.asarray(keep, dtype=float) + + def create_pin_hole_mask(mask: np.ndarray, approximate_contours: bool): """ Creates masks with pin holes added to contour regions with holes. diff --git a/rt_utils/rtstruct.py b/rt_utils/rtstruct.py index c52d816..35cb723 100644 --- a/rt_utils/rtstruct.py +++ b/rt_utils/rtstruct.py @@ -32,6 +32,7 @@ def add_roi( use_pin_hole: bool = False, approximate_contours: bool = True, roi_generation_algorithm: Union[str, int] = 0, + contour_mode: str = "voxel_center", ): """ Add a Region of Interest (ROI) to the RTStruct given a 3D binary mask for each slice. @@ -63,6 +64,10 @@ def add_roi( If False, skips approximation during contour generation, leading to larger contour data. Defaults to True. roi_generation_algorithm : str or int, optional Identifier for the algorithm used to generate the ROI. Defaults to 0. + contour_mode : {"voxel_center", "voxel_edge"}, optional + Placement of generated contours. ``"voxel_center"`` preserves the + original OpenCV behavior. ``"voxel_edge"`` traces the outer edges + of foreground voxels. Defaults to ``"voxel_center"``. Raises ------ @@ -89,6 +94,7 @@ def add_roi( use_pin_hole, approximate_contours, roi_generation_algorithm, + contour_mode, ) self.ds.ROIContourSequence.append( diff --git a/rt_utils/utils.py b/rt_utils/utils.py index f04089e..32b9e63 100644 --- a/rt_utils/utils.py +++ b/rt_utils/utils.py @@ -28,6 +28,7 @@ ] ROI_GENERATION_ALGORITHMS = ["AUTOMATIC", "SEMIAUTOMATIC", "MANUAL"] +CONTOUR_MODES = ["voxel_center", "voxel_edge"] class SOPClassUID: @@ -51,11 +52,20 @@ class ROIData: use_pin_hole: bool = False approximate_contours: bool = True roi_generation_algorithm: Union[str, int] = 0 + contour_mode: str = "voxel_center" def __post_init__(self): self.validate_color() self.add_default_values() self.validate_roi_generation_algoirthm() + self.validate_contour_mode() + + def validate_contour_mode(self): + if self.contour_mode not in CONTOUR_MODES: + raise ValueError( + f"Invalid contour mode '{self.contour_mode}'. " + f"Expected one of {CONTOUR_MODES}." + ) def add_default_values(self): if self.color is None: diff --git a/tests/test_rtstruct_builder.py b/tests/test_rtstruct_builder.py index e8cc856..15e1d21 100644 --- a/tests/test_rtstruct_builder.py +++ b/tests/test_rtstruct_builder.py @@ -181,6 +181,37 @@ def test_no_approximation_iou(new_rtstruct: RTStruct): run_mask_iou_test(new_rtstruct, mask, IOU_threshold, approximate_contours=False) +def test_voxel_edge_contour_coordinates(new_rtstruct: RTStruct): + mask = get_empty_mask(new_rtstruct) + mask[50:100, 60:120, 0] = True + + new_rtstruct.add_roi(mask, contour_mode="voxel_edge") + + contour_data = np.asarray( + new_rtstruct.ds.ROIContourSequence[-1].ContourSequence[0].ContourData, + dtype=float, + ).reshape(-1, 3) + pixel_to_patient = image_helper.get_pixel_to_patient_transformation_matrix( + new_rtstruct.series_data + ) + patient_to_pixel = np.linalg.inv(pixel_to_patient) + pixel_points = image_helper.apply_transformation_to_3d_points( + contour_data, patient_to_pixel + ) + assert np.isclose(pixel_points[:, 0].min(), 59.5) + assert np.isclose(pixel_points[:, 0].max(), 119.5) + assert np.isclose(pixel_points[:, 1].min(), 49.5) + assert np.isclose(pixel_points[:, 1].max(), 99.5) + + +def test_invalid_contour_mode(new_rtstruct: RTStruct): + mask = get_empty_mask(new_rtstruct) + mask[50:100, 50:100, 0] = True + + with pytest.raises(ValueError, match="Invalid contour mode"): + new_rtstruct.add_roi(mask, contour_mode="not-a-mode") + + def test_contour_data_sizes(new_rtstruct: RTStruct): mask = get_empty_mask(new_rtstruct) mask[50:100, 50:100, 0] = 1 diff --git a/tests/test_utils.py b/tests/test_utils.py index 1a6bec9..bac9d69 100644 --- a/tests/test_utils.py +++ b/tests/test_utils.py @@ -1,9 +1,54 @@ +import numpy as np import pytest +from rt_utils.image_helper import mask_to_edge_polygons from rt_utils.utils import COLOR_PALETTE from tests.test_rtstruct_builder import get_empty_mask +def test_single_voxel_edge_polygon(): + mask = np.zeros((3, 4), dtype=bool) + mask[1, 2] = True + + polygons = mask_to_edge_polygons(mask) + + assert len(polygons) == 1 + assert set(map(tuple, polygons[0])) == { + (1.5, 0.5), + (2.5, 0.5), + (2.5, 1.5), + (1.5, 1.5), + } + + +def test_rectangular_edge_polygon_merges_collinear_vertices(): + mask = np.zeros((6, 7), dtype=bool) + mask[1:4, 2:6] = True + + polygon = mask_to_edge_polygons(mask)[0] + + assert len(polygon) == 4 + assert (polygon[:, 0].min(), polygon[:, 0].max()) == (1.5, 5.5) + assert (polygon[:, 1].min(), polygon[:, 1].max()) == (0.5, 3.5) + + +def test_edge_polygons_preserve_holes_as_separate_loops(): + mask = np.ones((5, 5), dtype=bool) + mask[2, 2] = False + + polygons = mask_to_edge_polygons(mask) + + assert len(polygons) == 2 + + +def test_edge_polygons_handle_diagonally_touching_voxels(): + mask = np.eye(2, dtype=bool) + + polygons = mask_to_edge_polygons(mask) + + assert sum(len(polygon) for polygon in polygons) >= 6 + + VALID_COLORS = [ ("fff", [255, 255, 255]), ("#fff", [255, 255, 255]),