"""
PyPlanAnalysis.io
==================
DICOM discovery and loading, CT/dose grid resampling, and RT Struct
mask extraction (binary and fractional).
"""
import warnings
import numpy as np
from pathlib import Path
from typing import Union
from matplotlib.path import Path as MplPath
import pydicom
from scipy.interpolate import splprep, splev
from collections import defaultdict
from skimage.draw import polygon as sk_polygon
import SimpleITK as sitk
from dataclasses import dataclass, field
from typing import Optional
# We read RT Struct contours directly via pydicom for reliability.
"""
DICOM auto-discovery for an RT (proton/photon) dataset folder.
Given a folder that may contain a mix of CT slices, one or more RTSTRUCT,
RTPLAN and RTDOSE (physical dose + LET) files — possibly several
candidates of each, possibly nested in sub-folders, possibly with some
files that don't actually belong together — this finds the one
self-consistent set by walking the standard DICOM cross-reference chain:
RTPLAN
ReferencedStructureSetSequence[0]
(0008,1150) ReferencedSOPClassUID
(0008,1155) ReferencedSOPInstanceUID -> RTSTRUCT
(0008,0016) SOPClassUID
(0008,0018) SOPInstanceUID
RTSTRUCT
ReferencedFrameOfReferenceSequence
-> RTReferencedStudySequence
-> RTReferencedSeriesSequence
-> (0020,000E) SeriesInstanceUID -> CT series
RTDOSE
ReferencedRTPlanSequence[0]
(0008,1155) ReferencedSOPInstanceUID -> RTPLAN
(0008,0018) SOPInstanceUID
If the first candidate of each type doesn't line up, instead of failing
the function tries other candidates present in the folder until it finds
a combination where every link checks out. If no fully-verified
combination exists, it falls back to the best partial match it can build
and explains exactly what couldn't be confirmed via `link_warnings`.
"""
def _collect_candidates(folder: Path):
@dataclass
class CTSeries:
series_uid: str
directory: Path
patient_id: Optional[str] = None
frame_of_ref_uid: Optional[str] = None
@dataclass
class RTStructCandidate:
path: Path
sop_class: Optional[str]
sop_uid: Optional[str]
patient_id: Optional[str]
ref_series_uids: set = field(default_factory=set)
frame_of_ref_uid: Optional[str] = None
@dataclass
class RTPlanCandidate:
path: Path
sop_class: Optional[str]
sop_uid: Optional[str]
patient_id: Optional[str]
ref_struct: Optional[tuple] = None # (ReferencedSOPClassUID, ReferencedSOPInstanceUID)
@dataclass
class RTDoseCandidate:
path: Path
sop_uid: Optional[str]
patient_id: Optional[str]
dose_kind: str # "PHYSICAL", "LET", or "EFFECTIVE"
dose_SumType: str #"PLAN" or "BEAM"
ref_plan_sop: Optional[str] = None
frame_of_ref_uid: Optional[str] = None
"""One recursive pass over every *.dcm file in `folder` (including
sub-folders), sorting each into a lightweight candidate record keyed
by modality. Searching recursively means files don't need to live
directly in `folder` for this to find them."""
ct_series: dict = {}
rtstructs: list = []
rtplans: list = []
rtdoses: list = []
for f in folder.rglob("*.dcm"):
try:
ds = pydicom.dcmread(str(f), stop_before_pixels=True)
except Exception:
continue
modality = getattr(ds, "Modality", "")
patient_id = getattr(ds, "PatientID", None)
if modality == "CT":
series_uid = getattr(ds, "SeriesInstanceUID", None)
if series_uid and series_uid not in ct_series:
ct_series[series_uid] = CTSeries(series_uid, f.parent, patient_id,
frame_of_ref_uid=getattr(ds, "FrameOfReferenceUID", None))
elif modality == "RTSTRUCT":
ref_series_uids = set()
for frame in getattr(ds, "ReferencedFrameOfReferenceSequence", []):
struct_frame_uid = getattr(frame, "FrameOfReferenceUID", None)
for study in getattr(frame, "RTReferencedStudySequence", []):
for series in getattr(study, "RTReferencedSeriesSequence", []):
uid = getattr(series, "SeriesInstanceUID", None)
if uid:
ref_series_uids.add(uid)
rtstructs.append(RTStructCandidate(
path=f,
sop_class=getattr(ds, "SOPClassUID", None),
sop_uid=getattr(ds, "SOPInstanceUID", None),
patient_id=patient_id,
ref_series_uids=ref_series_uids,
frame_of_ref_uid=struct_frame_uid,
))
elif modality == "RTPLAN":
ref_struct = None
ref_seq = getattr(ds, "ReferencedStructureSetSequence", None)
if ref_seq:
ref = ref_seq[0]
ref_struct = (
getattr(ref, "ReferencedSOPClassUID", None),
getattr(ref, "ReferencedSOPInstanceUID", None),
)
rtplans.append(RTPlanCandidate(
path=f,
sop_class=getattr(ds, "SOPClassUID", None),
sop_uid=getattr(ds, "SOPInstanceUID", None),
patient_id=patient_id,
ref_struct=ref_struct,
))
elif modality == "RTDOSE":
dose_type = getattr(ds, "DoseType", "").upper()
label = getattr(ds, "DoseComment", "").upper()
dose_SumType = getattr(ds, "DoseSummationType", "").upper()
fname = f.name.upper()
is_let = ("LET" in label or "LET" in fname or dose_type == "LET")
if is_let:
dose_kind = "LET"
elif dose_type in ("PHYSICAL"):
dose_kind = "PHYSICAL"
elif dose_type in ("EFFECTIVE"):
dose_kind = "EFFECTIVE"
else:
dose_kind = "UNKNOWN"
ref_plan_sop = None
ref_plan_seq = getattr(ds, "ReferencedRTPlanSequence", None)
if ref_plan_seq:
ref_plan_sop = getattr(ref_plan_seq[0], "ReferencedSOPInstanceUID", None)
rtdoses.append(RTDoseCandidate(
path=f,
sop_uid=getattr(ds, "SOPInstanceUID", None),
patient_id=patient_id,
dose_kind=dose_kind,
dose_SumType=dose_SumType,
ref_plan_sop=ref_plan_sop,
frame_of_ref_uid=getattr(ds, "FrameOfReferenceUID", None)
))
return ct_series, rtstructs, rtplans, rtdoses
def _best_dose(doses_of_kind, plan, fallback_pool, kind_label,
frame_of_ref_uid, link_warnings):
"""
Pick the best RTDOSE candidate of one kind ("physical" or "LET"),
trying progressively weaker (but still verifiable) links instead of
ever guessing. If nothing can be verified, return a sentinel object
that will cause a clear crash downstream rather than silently
proceeding with a wrong file.
"""
# if no candidates of this kind AND no fallback candidates, nothing to pick
if not doses_of_kind and not fallback_pool:
return None
# --- opt 1: RTPLAN is available and a candidate references it ---
# this is the strongest possible link: RTDOSE -> RTPLAN -> RTSTRUCT
if plan is not None and doses_of_kind:
match = next((d for d in doses_of_kind if d.ref_plan_sop == plan.sop_uid), None)
if match is not None:
return match
# --- opt 2: no RTPLAN match (or no RTPLAN at all), but a candidate
# shares the FrameOfReferenceUID with the matched CT/RTSTRUCT ---
# this works even when RTPLAN is completely missing from the folder
if frame_of_ref_uid and doses_of_kind:
match = next((d for d in doses_of_kind
if d.frame_of_ref_uid == frame_of_ref_uid), None)
if match is not None:
if plan is not None:
link_warnings.append(
f"No {kind_label} RTDOSE references the matched RTPLAN, "
"but one shares the same FrameOfReferenceUID as the "
"matched CT/RTSTRUCT; using that instead."
)
else:
link_warnings.append(
f"No RTPLAN available, but one {kind_label} RTDOSE shares the same FrameOfReferenceUID as the "
"matched CT/RTSTRUCT; using that instead."
)
return match
# --- no candidates of the requested kind at all: repeat the same
# three tiers against the fallback pool (e.g. EFFECTIVE in case of physical dose) ---
# RTPLAN reference
if plan is not None and fallback_pool:
match = next((d for d in fallback_pool if d.ref_plan_sop == plan.sop_uid), None)
if match is not None:
link_warnings.append(
f"No RTDOSE was tagged or named as {kind_label}; using an "
"EFFECTIVE RTDOSE because it references the matched RTPLAN."
)
return match
# FrameOfReferenceUID match
if frame_of_ref_uid and fallback_pool:
match = next((d for d in fallback_pool
if d.frame_of_ref_uid == frame_of_ref_uid), None)
if match is not None:
if plan is not None:
link_warnings.append(
f"No fallback RTDOSE references the matched RTPLAN, "
"but one shares the same FrameOfReferenceUID as the "
"matched CT/RTSTRUCT; using that instead."
)
else:
link_warnings.append(
f"No RTPLAN available, but one fallback RTDOSE shares the same FrameOfReferenceUID as the "
"matched CT/RTSTRUCT; using that instead."
)
return match
# If no matches are found, and exactly one unlinked candidate of the selected kind exists, pick it ---
# weak evidence, but in single-plan-per-folder layouts this is
# usually correct; flag it so the caller can decide whether to trust it
if len(doses_of_kind) == 1:
link_warnings.append(
f"Only one {kind_label} RTDOSE file found; USING IT WITHOUT MATCH to the RTPLAN or CT/RTSTRUCT."
)
return doses_of_kind[0]
# Multiple candidates of the selected kind, none verifiable ---
# force a hard failure
if len(doses_of_kind) > 1:
link_warnings.append(
f"Multiple {kind_label} RTDOSE files found and none verifiably "
"link to the matched RTPLAN or CT/RTSTRUCT — refusing to guess."
)
return _AMBIGUOUS_DOSE
# If no matches are found, and exactly one unlinked candidate of the fallback dose type exists, pick it
if len(fallback_pool) == 1:
link_warnings.append(
f"No RTDOSE was tagged or named as {kind_label}; using the only "
"EFFECTIVE RTDOSE found, without a verified link."
)
return fallback_pool[0]
# multiple unlinked fallback candidates — force a hard failure
if len(fallback_pool) > 1:
link_warnings.append(
f"No RTDOSE was tagged or named as {kind_label}, and multiple "
"EFFECTIVE RTDOSE files exist with no verifiable link — "
"refusing to guess."
)
return _AMBIGUOUS_DOSE
return None
def _strict_chain(ct_series, rtstructs, rtplans):
"""Try every RTPLAN/RTSTRUCT pairing (or RTSTRUCT alone, if no RTPLAN
exists) until one is found whose reference tags fully check out
against a CT series actually present in the folder. Returns
(plan, struct, series_uid, fully_verified) or (None, None, None, False)
if nothing checks out."""
plan_candidates = rtplans if rtplans else [None]
for plan in plan_candidates:
if plan is not None:
struct = next(
(s for s in rtstructs if plan.ref_struct == (s.sop_class, s.sop_uid)),
None,
)
else:
struct = next((s for s in rtstructs if s.ref_series_uids & set(ct_series)), None)
if struct is None:
continue
series_uid = next((uid for uid in struct.ref_series_uids if uid in ct_series), None)
if series_uid is None:
continue
return plan, struct, series_uid, (plan is not None)
return None, None, None, False
def _fallback_chain(ct_series, rtstructs, rtplans, link_warnings):
"""No fully cross-referenced combination exists. Build the best
available guess one piece at a time, logging exactly what had to be
assumed instead of confirmed."""
struct = next((s for s in rtstructs if s.ref_series_uids & set(ct_series)), None)
if struct is not None:
series_uid = next(uid for uid in struct.ref_series_uids if uid in ct_series)
else:
series_uid = next(iter(ct_series), None)
if rtstructs:
link_warnings.append(
"No RTSTRUCT references any of the discovered CT series; using "
"the first RTSTRUCT found without a verified CT link."
)
struct = rtstructs[0]
if series_uid is not None:
link_warnings.append(
f"Assuming CT series {series_uid} since it could not be "
"confirmed via RTSTRUCT reference tags."
)
plan = None
if struct is not None:
plan = next(
(p for p in rtplans if p.ref_struct == (struct.sop_class, struct.sop_uid)),
None,
)
if plan is None and rtplans:
link_warnings.append(
"No RTPLAN references the selected RTSTRUCT; using the first "
"RTPLAN found without a verified link."
)
plan = rtplans[0]
return plan, struct, series_uid
class _AmbiguousDose:
"""
Sentinel returned when multiple RTDOSE candidates exist and none can
be verifiably linked to the matched plan/structure. Any attempt to use
this as a real candidate (accessing .path, .sop_uid, etc.) raises
AttributeError immediately, so the ambiguity surfaces as a hard crash
instead of a silently wrong file being picked.
"""
def __getattr__(self, name):
raise AttributeError(
f"Ambiguous RTDOSE match: cannot access '.{name}' — multiple "
"unlinked candidates were found and none could be picked safely. "
"Resolve manually (check link_warnings) before proceeding."
)
def __bool__(self):
# so `if dose:` style checks still behave like "something is there"
# forcing any downstream .path access to be the point of failure
return True
_AMBIGUOUS_DOSE = _AmbiguousDose()
[docs]
def find_dicom_files(folder: Path) -> dict:
"""
Auto-discover RT Dose (physical dose, LET), RT Struct, RT Plan and CT
files belonging to the same plan, by inspecting DICOM modality tags
and cross-reference (Referenced UID) tags. The folder (and its
sub-folders) may contain extra or unrelated files of any of these
types; this function searches through all of them for the one set that is
actually linked together.
Returns
-------
dict with keys:
"dose", "let", "rtstruct", "rtplan" : Path or None
"CT" : Path to the directory holding the matched CT
series, or None
"Patient_ID" : str or None
"linked" : True if RTPLAN -> RTSTRUCT -> CT was fully
confirmed via DICOM reference tags; False if a
fallback/best-guess match had to be used; None
if there wasn't enough data to even attempt the
check (e.g. no RTSTRUCT and no CT found at all).
"link_warnings": list[str], one entry per fallback
that was needed, explaining what couldn't be
confirmed and what was used instead.
"""
folder = Path(folder)
link_warnings: list = []
ct_series, rtstructs, rtplans, rtdoses = _collect_candidates(folder)
plan, struct, series_uid, verified = _strict_chain(ct_series, rtstructs, rtplans)
if struct is None and (rtstructs or ct_series):
# The strict search found nothing usable at all; fall back.
msg = "Could not find an RTPLAN/RTSTRUCT/CT combination that fully "
"cross-references; falling back to best-effort matching."
link_warnings.append(msg)
print(msg)
plan, struct, series_uid = _fallback_chain(ct_series, rtstructs, rtplans, link_warnings)
verified = False
elif plan is None and rtplans:
# A struct/CT-only chain was found (e.g. no RTPLAN references it),
# but there are RTPLAN files sitting in the folder we never matched.
msg = f"{len(rtplans)} RTPLAN file(s) found but none reference the "
"matched RTSTRUCT; proceeding without a confirmed RTPLAN."
link_warnings.append(msg)
physical_doses = [d for d in rtdoses if (d.dose_kind == "PHYSICAL" and d.dose_SumType == "PLAN")]
effective_doses = [d for d in rtdoses if (d.dose_kind == "EFFECTIVE" and d.dose_SumType == "PLAN")]
let_doses = [d for d in rtdoses if d.dose_kind == "LET"]
frame_of_ref_uid = struct.frame_of_ref_uid if struct is not None else None
#look for best matching physical dose, at worst, look for effective dose and scale by 10%
dose = _best_dose(physical_doses, plan, effective_doses, "physical",
frame_of_ref_uid, link_warnings)
let = _best_dose(let_doses, plan, [], "LET",
frame_of_ref_uid, link_warnings)
ct_dir = ct_series[series_uid].directory if series_uid in ct_series else None
patient_id = None
for obj in (struct, plan, dose, let):
if obj is not None and obj.patient_id:
patient_id = obj.patient_id
break
if patient_id is None and series_uid in ct_series:
patient_id = ct_series[series_uid].patient_id
linked = None
if struct is not None or ct_series:
linked = verified and len(link_warnings) == 0
found = {
"dose": dose.path if dose else None,
"let": let.path if let else None,
"rtstruct": struct.path if struct else None,
"rtplan": plan.path if plan else None,
"CT": ct_dir,
"Patient_ID": patient_id,
"linked": linked,
"link_warnings": link_warnings
}
for w in link_warnings:
print("⚠", w)
return found
#%%
[docs]
def load_ct_series(ct_folder: Union[str, Path]) -> tuple:
"""
Load a multi-slice CT DICOM series from a folder.
Slices are sorted by ImagePositionPatient z-coordinate.
Spacing is taken from PixelSpacing of the first slice and the
z-step between consecutive slice positions.
Returns
-------
sitk_image : SimpleITK.Image (x, y, z ordering internally)
ct_geometry : dict with keys:
"origin" : [x0, y0, z0] mm (corner of first voxel)
"spacing" : [dx, dy, dz] mm
"shape" : (nz, ny, nx) — numpy (z,y,x) convention
"z_positions" : np.ndarray of slice z-coordinates length nz
"""
ct_folder = Path(ct_folder)
slices = []
for f in ct_folder.glob("*.dcm"):
try:
ds = pydicom.dcmread(str(f), stop_before_pixels=True)
if getattr(ds, "Modality", "") == "CT":
slices.append((float(ds.ImagePositionPatient[2]), str(f), ds))
except Exception:
continue
if not slices:
raise FileNotFoundError(f"No CT DICOM files found in {ct_folder}")
slices.sort(key=lambda t: t[0]) # sort by z-position
z_positions = np.array([s[0] for s in slices])
first_ds = slices[0][2]
pix_sp = [float(v) for v in first_ds.PixelSpacing] # [row_sp=dy, col_sp=dx]
dx, dy = pix_sp[1], pix_sp[0]
dz = float(z_positions[1] - z_positions[0]) if len(z_positions) > 1 else float(first_ds.SliceThickness)
origin = [float(v) for v in first_ds.ImagePositionPatient] # [x0, y0, z0]
# Use SimpleITK series reader for correct pixel data ordering
reader = sitk.ImageSeriesReader()
dicom_names = reader.GetGDCMSeriesFileNames(str(ct_folder))
if not dicom_names:
# fallback: use our sorted file list
dicom_names = [s[1] for s in slices]
reader.SetFileNames(dicom_names)
sitk_image = reader.Execute()
nz = len(slices)
ny = int(first_ds.Rows)
nx = int(first_ds.Columns)
ct_geometry = {
"origin" : origin, # itk [x0, y0, z0]
"spacing" : [dx, dy, dz], # itk [dx, dy, dz]
"shape" : (nz, ny, nx), # numpy (x, y, z)
"z_positions": z_positions,
}
print(f"CT loaded: {nz} slices spacing=({dx:.2f},{dy:.2f},{dz:.2f}) mm "
f"shape={ct_geometry['shape']}")
return sitk_image, ct_geometry
def _np_to_sitk(arr, ds):
"""Wrap a numpy (z,y,x) dose array as a properly georeferenced sitk image."""
img = sitk.GetImageFromArray(arr)
origin, spacing = get_grid_geometry(ds) #[dx, dy, dz]
# SimpleITK spacing order: (x, y, z)
img.SetSpacing(spacing)
img.SetOrigin(tuple(origin))
return img
[docs]
def resample_dose_to_new_grid(
dose_sitk,
dose_ds,
new_spacing,
interpolator = sitk.sitkLinear):
"""
Resample RTDOSE onto a new isotropic/anisotropic grid
Parameters
----------
dose_sitk : sitk.Image Original dose image
dose_ds : pydicom Dataset RTDOSE dataset
new_spacing : tuple/list (sx, sy, sz) in mm
Returns
-------
resampled_dose : sitk.Image
updated_info : dict
dose_ds : updated dataset
"""
old_spacing = np.array(dose_sitk.GetSpacing())
old_size = np.array(dose_sitk.GetSize())
old_origin = dose_sitk.GetOrigin()
old_direction = dose_sitk.GetDirection()
# physical extent
physical_size = old_spacing * old_size
# -----------------------------
new_spacing = np.array(new_spacing)
new_size = np.ceil( physical_size / new_spacing ).astype(int)
# reference image
ref = sitk.Image( [int(v) for v in new_size], dose_sitk.GetPixelID() )
ref.SetSpacing(tuple(new_spacing))
ref.SetOrigin(old_origin)
ref.SetDirection(old_direction)
# -----------------------------
# RESAMPLE
resampled_dose = sitk.Resample( dose_sitk,
ref,
sitk.Transform(),
interpolator,
0.0)
dose_arr_resampled = sitk.GetArrayFromImage( resampled_dose )
# numpy shape = z,y,x
shape_np = dose_arr_resampled.shape
z_spacing = new_spacing[2]
z_positions = (old_origin[2]+ np.arange(shape_np[0]) * z_spacing )
z_offsets = (z_positions - z_positions[0])
# -----------------------------
# UPDATE DICOM RTDOSE
dose_ds.PixelSpacing = [ float(new_spacing[1]), float(new_spacing[0])]
dose_ds.SliceThickness = float(z_spacing)
dose_ds.GridFrameOffsetVector = [float(v) for v in z_offsets]
dose_ds.Rows = shape_np[1]
dose_ds.Columns = shape_np[2]
dose_ds.NumberOfFrames = shape_np[0]
updated_info = {
"spacing": tuple(new_spacing),
"origin": old_origin,
"shape": shape_np,
"z_positions": z_positions,
"z_offsets": z_offsets,
}
return (resampled_dose,
dose_arr_resampled,
updated_info,
dose_ds)
[docs]
def resample_dose_on_ct(sitk_dose: sitk.Image,
sitk_ct: sitk.Image) -> sitk.Image:
"""
Resample a dose (or LET) SimpleITK image onto the CT grid.
The CT image defines the output origin, spacing, direction, and size.
This ensures that the resampled dose array is perfectly aligned with
the CT grid on which contours will be rasterised.
Parameters
----------
sitk_dose : SimpleITK.Image (dose or LET, in dose-grid coordinates)
sitk_ct : SimpleITK.Image (full CT series)
Returns
-------
SimpleITK.Image same grid as sitk_ct
"""
resampler = sitk.ResampleImageFilter()
resampler.SetOutputSpacing(sitk_ct.GetSpacing())
resampler.SetSize(sitk_ct.GetSize())
resampler.SetOutputDirection(sitk_ct.GetDirection())
resampler.SetOutputOrigin(sitk_ct.GetOrigin())
resampler.SetInterpolator(sitk.sitkLinear)
resampler.SetDefaultPixelValue(0.0)
return resampler.Execute(sitk_dose)
# **************************
#=======================================
[docs]
def load_dose_grid(path: Union[str, Path]) -> tuple:
"""
Load a DICOM RT Dose file (dose or LET stored as dose grid).
Returns
-------
array : np.ndarray shape (z, y, x)
ds : pydicom Dataset
"""
ds = pydicom.dcmread(str(path))
scale = float(ds.DoseGridScaling)
if not ds.DoseUnits == "GY":
warnings.warn(f"Dose Units are not correct for: '{path}'.")
array = ds.pixel_array.astype(np.float64) * scale
return array, ds
[docs]
def get_grid_geometry(ds) -> tuple:
"""
Extract (origin, spacing) from an RT Dose dataset.
Returns
-------
origin : [x0, y0, z0] mm
spacing : [dx, dy, dz] mm
"""
origin = [float(v) for v in ds.ImagePositionPatient]
pix_sp = [float(v) for v in ds.PixelSpacing] # [row_spacing=dy, col_spacing=dx]
dz = float(ds.SliceThickness)
# PixelSpacing = [row_spacing, col_spacing] = [dy, dx]
return origin, [pix_sp[1], pix_sp[0], dz] # [dx, dy, dz]
# ============================================================
# Structure mask extraction
def _build_roi_maps(rtstruct_ds) -> tuple:
"""
Parse RT Struct and return two lookup dicts.
Returns
-------
name_to_roi : {roi_name_lower: roi_number}
roi_to_contours : {roi_number: [ np.ndarray shape(N,3) ]}
Each array is one contour polygon with columns [x, y, z] in mm.
"""
# ROI names from StructureSetROISequence
name_to_roi = {}
for item in rtstruct_ds.StructureSetROISequence:
name_to_roi[item.ROIName.strip().lower()] = int(item.ROINumber)
# Contour coordinates from ROIContourSequence
roi_to_contours = {}
for roi_contour in rtstruct_ds.ROIContourSequence:
roi_num = int(roi_contour.ReferencedROINumber)
contours = []
if not hasattr(roi_contour, "ContourSequence"):
roi_to_contours[roi_num] = contours
continue
for contour in roi_contour.ContourSequence:
raw = [float(v) for v in contour.ContourData]
pts = np.array(raw).reshape(-1, 3) # (N, 3) x,y,z
contours.append(pts)
roi_to_contours[roi_num] = contours
return name_to_roi, roi_to_contours
[docs]
def get_all_structure_names(rtstruct_ds) -> list:
"""Return list of all structure names in an RT Struct dataset."""
return [item.ROIName.strip()
for item in rtstruct_ds.StructureSetROISequence]
[docs]
def get_structure_mask_on_grid(struct_name: str,
rtstruct_ds,
origin: list,
spacing: list,
shape: tuple,
z_positions: np.ndarray) -> np.ndarray:
"""
Rasterise RT Struct contours for `struct_name` onto an arbitrary grid.
This is the core function used for both dose-grid and CT-grid masking.
Contour z-values are matched to the nearest z in z_positions.
Parameters
----------
struct_name : str
rtstruct_ds : pydicom Dataset
origin : [x0, y0, z0] mm — physical coordinate of voxel (0,0,0) corner
spacing : [dx, dy, dz] mm
shape : (nz, ny, nx) — numpy array shape
z_positions : 1-D array of z-coordinates for each slice (length nz)
Returns
-------
mask : np.ndarray bool, shape (nz, ny, nx)
"""
name_to_roi, roi_to_contours = _build_roi_maps(rtstruct_ds)
key = struct_name.strip().lower()
if key not in name_to_roi:
raise ValueError(
f"Structure '{struct_name}' not found in RT Struct. "
f"Available: {[item.ROIName for item in rtstruct_ds.StructureSetROISequence]}"
)
roi_number = name_to_roi[key]
contours = roi_to_contours.get(roi_number, [])
x0, y0, z0 = origin
dx, dy, dz = spacing
nz, ny, nx = shape
mask = np.zeros(shape, dtype=bool)
if not contours:
warnings.warn(f"No contour data for '{struct_name}'.")
return mask
# Build grid of voxel-centre x,y coordinates (voxel centres = origin + (i+0.5)*spacing)
# Note: DICOM ImagePositionPatient is the centre of the first voxel, so:
# voxel centre i → x0 + i*dx
xi = np.arange(nx)
yi = np.arange(ny)
XX, YY = np.meshgrid(x0 + xi * dx, y0 + yi * dy)
grid_xy = np.column_stack([XX.ravel(), YY.ravel()])
for pts in contours:
z_val = float(pts[0, 2])
z_idx = int(np.argmin(np.abs(z_positions - z_val)))
poly_xy = pts[:, :2]
if len(poly_xy) < 3:
continue
poly = MplPath(poly_xy)
inside = poly.contains_points(grid_xy).reshape(ny, nx)
mask[z_idx] |= inside
return mask
[docs]
def get_roi_center_of_mass(struct_name: str, rtstruct_ds) -> np.ndarray:
"""
Compute the 3-D center of mass (mm, in the RT Struct's patient
coordinate system) of a structure directly from its contour
polygons — no dose/CT grid required.
Each contour (one per slice) contributes its own 2-D polygon
centroid, weighted by that contour's area, so slices with more
cross-sectional area count more towards the overall COM volume-weighted approximation without needing a full
3-D mask.
Parameters
----------
struct_name : str
Structure name as it appears in the RT Struct.
rtstruct_ds : pydicom Dataset
Returns
-------
np.ndarray, shape (3,), or None
``[x, y, z]`` center of mass in mm. Returns ``None`` if no
usable contour data is available for this structure.
"""
name_to_roi, roi_to_contours = _build_roi_maps(rtstruct_ds)
key = struct_name.strip().lower()
if key not in name_to_roi:
raise ValueError(
f"Structure '{struct_name}' not found in RT Struct. "
f"Available: {[item.ROIName for item in rtstruct_ds.StructureSetROISequence]}"
)
roi_number = name_to_roi[key]
contours = roi_to_contours.get(roi_number, [])
if not contours:
warnings.warn(f"No contour data for '{struct_name}' — cannot compute center of mass.")
return None
centroids = []
weights = []
for pts in contours:
xy = pts[:, :2]
if len(xy) < 3:
continue
area = contour_area_signed(xy)
if area == 0:
continue
x, y = xy[:, 0], xy[:, 1]
x1, y1 = np.roll(x, -1), np.roll(y, -1)
cross = x * y1 - x1 * y
cx = np.sum((x + x1) * cross) / (6 * area)
cy = np.sum((y + y1) * cross) / (6 * area)
z = float(pts[0, 2])
centroids.append([cx, cy, z])
weights.append(abs(area))
if not centroids:
warnings.warn(f"Could not compute a valid centroid for '{struct_name}'.")
return None
centroids = np.array(centroids)
weights = np.array(weights)
return np.average(centroids, axis=0, weights=weights)
[docs]
def contour_area_signed(xy):
"""
Signed polygon area via the shoelace formula.
Parameters
----------
xy : np.ndarray, shape (N, 2)
Polygon vertices [x, y] in mm.
Returns
-------
float
Signed area; positive for counter-clockwise vertex order,
negative for clockwise. Used to detect contour holes.
"""
x, y = xy[:, 0], xy[:, 1]
return 0.5 * (np.dot(x, np.roll(y, -1)) - np.dot(y, np.roll(x, -1)))
[docs]
def smooth_contour(poly_xy, n_pts=300):
"""
Resample a contour polygon onto ``n_pts`` evenly-spaced points using
a periodic B-spline fit. Not used by default in the fractional-mask
pipeline (see inline note in ``get_fractional_mask_on_grid``); kept
available for callers who want smoothed contours.
Parameters
----------
poly_xy : np.ndarray, shape (N, 2)
Polygon vertices [x, y] in mm.
n_pts : int, default 300
Number of points in the resampled output.
Returns
-------
np.ndarray, shape (n_pts, 2)
Smoothed polygon vertices. Falls back to the original polygon,
unchanged, if the spline fit fails.
"""
try:
x, y = poly_xy[:, 0], poly_xy[:, 1]
tck, _ = splprep([x, y], s=0, per=True)
x_s, y_s = splev(np.linspace(0, 1, n_pts), tck)
return np.column_stack([x_s, y_s])
except Exception as e:
warnings.warn(f"Spline failed: {e}")
return poly_xy
[docs]
def rasterize_supersampled(smooth_xy, x0, y0, dx, dy, ny, nx, N):
"""
Rasterise one contour polygon onto a grid at N times finer resolution,
then average back down — giving each output voxel a fractional
[0, 1] membership value instead of a hard binary in/out.
Parameters
----------
smooth_xy : np.ndarray, shape (N, 2)
Polygon vertices [x, y] in mm.
x0, y0 : float
Grid origin (mm) in x and y.
dx, dy : float
Grid voxel spacing (mm) in x and y.
ny, nx : int
Output grid shape.
N : int
Supersampling factor per side (N² sub-samples per voxel).
Returns
-------
np.ndarray, shape (ny, nx), float32
Fractional coverage of each voxel by the polygon, in [0, 1].
"""
xi = (smooth_xy[:, 0] - x0) / dx * N
yi = (smooth_xy[:, 1] - y0) / dy * N
rr, cc = sk_polygon(yi, xi, shape=(ny * N, nx * N))
super_mask = np.zeros((ny * N, nx * N), dtype=np.float32)
super_mask[rr, cc] = 1.0
return super_mask.reshape(ny, N, nx, N).mean(axis=(1, 3))
[docs]
def compute_roi_volume_comparison(frac_mask, dx, dy, dz):
"""
Print and return the total volume of a fractional mask, plus a
per-slice breakdown — useful for sanity-checking
``get_fractional_mask_on_grid`` output against a TPS-reported volume.
Parameters
----------
frac_mask : np.ndarray, shape (nz, ny, nx)
Fractional voxel membership mask, values in [0, 1].
dx, dy, dz : float
Voxel spacing (mm) in x, y, z.
Returns
-------
float
Total volume in cc (sum of fractional weights x voxel volume).
"""
# Your fractional volume
vol_frac = frac_mask.sum() * dx * dy * dz / 1000.0
# Per-slice breakdown
frac_per_slice = frac_mask.sum(axis=(1,2)) * dx * dy * dz / 1000.0
print(f"Fractional volume : {vol_frac:.4f} cc")
print("\nPer-slice :")
for z, f in enumerate(frac_per_slice):
if f > 0 :
print(f" slice {z:3d}: {f:.4f}")
return vol_frac
[docs]
def prismatoid_volume(frac_mask, dx, dy, dz):
"""
Volume of a fractional mask via the prismatoid (Simpson's-rule-like)
formula between consecutive slices, instead of a flat sum-of-slices
approximation. Slightly more accurate for structures with rapidly
changing cross-sectional area between slices.
Parameters
----------
frac_mask : np.ndarray, shape (nz, ny, nx)
Fractional voxel membership mask, values in [0, 1].
dx, dy, dz : float
Voxel spacing (mm) in x, y, z.
Returns
-------
float
Total volume in cc.
"""
areas = frac_mask.sum(axis=(1,2)) * dx * dy
nz = len(areas)
if nz < 2:
return areas.sum() * dz / 1000.0
vol = 0.0
for i in range(nz - 1):
A0, A1 = areas[i], areas[i+1]
Am = (A0 + A1) / 2.0
vol += (dz / 6.0) * (A0 + 4*Am + A1)
return vol / 1000.0
[docs]
def get_fractional_mask_on_grid(struct_name: str,
rtstruct_ds,
origin: list,
spacing: list,
shape: tuple,
z_positions: np.ndarray,
supersample: int = 4) -> np.ndarray:
"""
Compute a fractional voxel membership mask on an arbitrary grid.
Each voxel receives a value in [0,1] — the fraction of its physical
area (in xy) that lies inside the RT Struct contour, estimated by
supersampling (supersample² sub-points per voxel).
Parameters
----------
struct_name : str
rtstruct_ds : pydicom Dataset
origin : [x0, y0, z0] mm
spacing : [dx, dy, dz] mm
shape : (nz, ny, nx)
z_positions : 1-D array length nz
supersample : N subdivisions per side (default 4 → 16 sub-points/voxel)
Returns
-------
frac_mask : np.ndarray float32, shape (nz, ny, nx), values in [0, 1]
"""
name_to_roi, roi_to_contours = _build_roi_maps(rtstruct_ds)
key = struct_name.strip().lower()
if key not in name_to_roi:
raise ValueError(
f"Structure '{struct_name}' not found in RT Struct. "
f"Available: {[item.ROIName for item in rtstruct_ds.StructureSetROISequence]}"
)
roi_number = name_to_roi[key]
contours = roi_to_contours.get(roi_number, [])
x0, y0, z0 = origin
dx, dy, dz = spacing
nz, ny, nx = shape
frac_mask = np.zeros(shape, dtype=np.float32)
if not contours:
warnings.warn(f"No contour data for '{struct_name}'.")
return frac_mask
N = supersample
# Group contours by slice
slice_contours = defaultdict(list)
n = 0
for pts in contours:
n+=1
z_val = float(pts[0, 2])
diffs = np.abs(z_positions - z_val)
z_idx = int(np.argmin(diffs))
if diffs[z_idx] > dz * 0.5:
warnings.warn(f"Contour z={z_val:.2f}mm is {diffs[z_idx]:.2f}mm from nearest slice")
slice_contours[z_idx].append(pts[:, :2])
for z_idx, slice_polys in slice_contours.items():
if not slice_polys:
continue
areas = [abs(contour_area_signed(p)) for p in slice_polys]
slice_polys = [slice_polys[i] for i in np.argsort(areas)[::-1]]
slice_fraction = np.zeros((ny, nx), dtype=np.float32)
for i, poly_xy in enumerate(slice_polys):
if len(poly_xy) < 3:
continue
smooth_xy = poly_xy #smooth_contour(poly_xy) #Avoid soothing, no gain in accuracy compared to RayStation.
# Hole detections
is_hole = False
if i > 0:# assume the biggest area as non-hole
for j in range(i): # handling holes if holes do not have any island inside.
outer_path = MplPath(slice_polys[j], closed=True)
test_pts = poly_xy[::max(1, len(poly_xy) // 5)]
if outer_path.contains_points(test_pts).mean() > 0.5:
is_hole = True
break
fraction = rasterize_supersampled(smooth_xy, x0, y0, dx, dy, ny, nx, N)
if is_hole:
slice_fraction -= fraction
else:
slice_fraction += fraction
frac_mask[z_idx] = np.clip(slice_fraction, 0.0, 1.0) #clipping shouldn#t change anything, it should be already in [0-1]
return frac_mask