"""Outcome preprocessor shipped with the model repository. Implements :class:`pyRadPlan.ai.modelhub.OutcomePreprocessor`: CT and masks are resampled onto the model grid once, the dose on every call with differentiable ``grid_sample``. All parameters come from ``model_preprocessing`` in ``model_config.json``. """ from typing import Any, Optional, Union import logging import numpy as np import SimpleITK as sitk from pyRadPlan.ai.modelhub import OutcomePreprocessor try: import torch import torch.nn.functional as F # noqa: N812 except ImportError: torch = None # type: ignore F = None # type: ignore logger = logging.getLogger(__name__) class OutcomeCnnPreprocessor(OutcomePreprocessor): """Differentiable preprocessor mapping planning data onto a fixed model grid. Configuration is read from the ``model_preprocessing`` section of the model's ``model_config.json``: - ``input_dimensions`` : model grid size (X, Y, Z) - ``input_spacing`` : model grid voxel spacing in mm - ``center`` : ``"target"`` centers the model grid on the target center of mass (currently the only supported mode) - ``type_order`` : channel/argument order for :meth:`assemble` (default ``["dose", "ct", "mask"]``) - ``modality.dose`` : ``normalization_value`` (Gy), ``extract`` (mask the dose channel with the structure masks) - ``modality.ct`` : ``window`` HU window (default ``[-1024, 3071]``), mapped linearly to [0, 1] - ``modality.mask`` : ``collapse`` (merge all masks into one channel) Parameters ---------- config : dict, optional The ``model_preprocessing`` dictionary shipped with the model. """ def __init__(self, config: Optional[dict] = None) -> None: if torch is None: raise ImportError( "PyTorch is required for outcome-guided preprocessing. " "Install it e.g. via: pip install torch" ) super().__init__(config) self.input_dimensions: tuple[int, ...] = tuple(self.config["input_dimensions"]) self.input_spacing: tuple[float, ...] = tuple(self.config["input_spacing"]) self.center_mode: str = self.config.get("center", "target") self.type_order: list[str] = list(self.config.get("type_order", ["dose", "ct", "mask"])) modality = self.config.get("modality", {}) self.dose_config: dict = modality.get("dose", {}) self.ct_config: dict = modality.get("ct", {}) self.mask_config: dict = modality.get("mask", {}) self.device: "torch.device" = torch.device("cpu") # Static tensors, set by configure() self._ct_tensor: Optional[torch.Tensor] = None # (1, 1, mZ, mY, mX) self._mask_tensor: Optional[torch.Tensor] = None # (1, C, mZ, mY, mX) # Grid geometry self._model_grid: Optional[dict] = None self._dose_grid: Optional[dict] = None self._coords_dose_to_model: Optional[torch.Tensor] = None # ------------------------------------------------------------------ # Configuration # ------------------------------------------------------------------ def configure( self, ct, cst, cst_masks: Optional[list[str]] = None, device: Optional[Union[str, "torch.device"]] = None, ) -> None: """One-time setup of model-grid geometry, CT and mask tensors. Parameters ---------- ct : CT Planning CT (``ct.cube_hu`` is a SimpleITK image). cst : StructureSet Structure set providing ``target_center_of_mass()`` and the VOIs. cst_masks : list[str], optional VOI names used as mask channels (order matters unless the config collapses them). When *None*, all VOIs are used. device : str or torch.device, optional Compute device; should match the model's device. Defaults to CUDA if available, else CPU. """ if device is not None: self.device = torch.device(device) else: self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu") if self.center_mode == "target": center = np.asarray(cst.target_center_of_mass(), dtype=np.float64) else: raise ValueError(f"Unknown center mode: {self.center_mode}") ct_grid = self._grid_to_dict(ct.grid) self._model_grid = self._centered_model_grid(center, ct_grid["direction"]) coords_ct_to_model = self._compute_sample_coords(ct_grid, self._model_grid).to(self.device) self._prepare_ct(ct, cst, coords_ct_to_model) self._prepare_masks(cst, cst_masks, coords_ct_to_model) logger.info( "Outcome preprocessor configured: model grid %s @ %s mm, device=%s", self._model_grid["size"], self._model_grid["spacing"], self.device, ) def set_dose_grid(self, dose_grid) -> None: """Precompute dose->model sampling coordinates for an arbitrary dose grid. Parameters ---------- dose_grid : Grid The grid the optimization dose vector lives on. """ if self._model_grid is None: raise RuntimeError("Preprocessor not configured - call configure() first.") self._dose_grid = self._grid_to_dict(dose_grid) self._coords_dose_to_model = self._compute_sample_coords( self._dose_grid, self._model_grid ).to(self.device) logger.debug( "Dose grid set: %s @ %s mm", self._dose_grid["size"], self._dose_grid["spacing"] ) # ------------------------------------------------------------------ # Per-iteration API # ------------------------------------------------------------------ def preprocess(self, inputs: "torch.Tensor") -> tuple["torch.Tensor", ...]: """Build the model inputs from a flat (Fortran-order) dose-grid dose tensor.""" if self._coords_dose_to_model is None: raise RuntimeError("Dose grid not set - call set_dose_grid() first.") dx, dy, dz = self._dose_grid["size"] # Fortran-order flat (X, Y, Z) == C-order (Z, Y, X) dose_model = F.grid_sample( inputs.reshape(1, 1, dz, dy, dx), self._coords_dose_to_model, mode="bilinear", padding_mode="zeros", align_corners=True, ) # (1, 1, mZ, mY, mX) norm_value = self.dose_config.get("normalization_value") if norm_value is not None: dose_model = dose_model / norm_value if self.dose_config.get("extract", False): dose_model = dose_model * self._mask_tensor return self.assemble(dose_model, self._ct_tensor, self._mask_tensor) def assemble( self, dose: "torch.Tensor", ct: "torch.Tensor", mask: "torch.Tensor" ) -> tuple["torch.Tensor", ...]: """Arrange the channel tensors into the model's input signature. The default returns a tuple ordered by ``type_order``; the model is then called as ``model(*inputs)``. Override for models expecting e.g. a single channel-stacked tensor. """ channels = {"dose": dose, "ct": ct, "mask": mask} return tuple(channels[key] for key in self.type_order) def postprocess(self, outputs: Any) -> "torch.Tensor": """Map raw model output to a scalar outcome probability. The default assumes the model outputs a logit and applies a sigmoid. Override for models that already output probabilities. """ return torch.sigmoid(outputs).sum() # ------------------------------------------------------------------ # Introspection (mainly for debugging / tests) # ------------------------------------------------------------------ @property def ct_tensor(self) -> Optional["torch.Tensor"]: """CT tensor on the model grid, (1, 1, mZ, mY, mX).""" return self._ct_tensor @property def mask_tensor(self) -> Optional["torch.Tensor"]: """Mask tensor on the model grid, (1, C, mZ, mY, mX).""" return self._mask_tensor @property def model_grid(self) -> Optional[dict]: """Model grid geometry (size, origin, spacing, direction).""" return self._model_grid # ------------------------------------------------------------------ # Internals # ------------------------------------------------------------------ @staticmethod def _grid_to_dict(grid) -> dict: """Reduce a pyRadPlan Grid to the geometry needed for sampling.""" return { "size": tuple(int(d) for d in grid.dimensions), "origin": tuple(float(v) for v in grid.origin), "spacing": tuple(float(v) for v in grid.resolution_vector), "direction": tuple(float(v) for v in np.asarray(grid.direction).flatten()), } def _centered_model_grid(self, center: np.ndarray, direction: tuple) -> dict: """Model grid with the given center at its geometric center.""" size = self.input_dimensions spacing = self.input_spacing dir_mat = np.asarray(direction, dtype=np.float64).reshape(3, 3) center_idx = np.array([(s - 1) / 2.0 for s in size]) offset = dir_mat @ (np.asarray(spacing) * center_idx) origin = tuple((center - offset).tolist()) return { "size": tuple(size), "origin": origin, "spacing": tuple(spacing), "direction": tuple(direction), } @staticmethod def _compute_sample_coords(src_grid: dict, dst_grid: dict) -> "torch.Tensor": """Compute normalized [-1, 1] coords mapping ``dst_grid`` voxel centers into ``src_grid``. This is the grid ``F.grid_sample`` expects when the sampled tensor lives on ``src_grid``. Returns a tensor of shape (1, dZ, dY, dX, 3). """ dx, dy, dz = dst_grid["size"] ix = torch.arange(dx, dtype=torch.float32) iy = torch.arange(dy, dtype=torch.float32) iz = torch.arange(dz, dtype=torch.float32) gz, gy, gx = torch.meshgrid(iz, iy, ix, indexing="ij") # (dZ, dY, dX) indices = torch.stack( [gx.reshape(-1), gy.reshape(-1), gz.reshape(-1)], dim=1 ) # (N, 3) - x, y, z d_origin = torch.tensor(dst_grid["origin"], dtype=torch.float32) d_spacing = torch.tensor(dst_grid["spacing"], dtype=torch.float32) d_dir = torch.tensor(dst_grid["direction"], dtype=torch.float32).reshape(3, 3) phys = d_origin + (indices * d_spacing) @ d_dir.T # (N, 3) s_origin = torch.tensor(src_grid["origin"], dtype=torch.float32) s_spacing = torch.tensor(src_grid["spacing"], dtype=torch.float32) s_dir = torch.tensor(src_grid["direction"], dtype=torch.float32).reshape(3, 3) s_dir_inv = torch.linalg.inv(s_dir) src_idx = ((phys - s_origin) @ s_dir_inv.T) / s_spacing # (N, 3) s_size = torch.tensor(src_grid["size"], dtype=torch.float32) normalized = 2.0 * src_idx / (s_size - 1) - 1.0 # align_corners=True convention # grid_sample 5-D convention: grid[..., 0]=W(X), grid[..., 1]=H(Y), grid[..., 2]=D(Z) return normalized.reshape(1, int(dz), int(dy), int(dx), 3) def _prepare_ct(self, ct, cst, coords: "torch.Tensor") -> None: """Window/normalize the CT and resample it to the model grid.""" window = self.ct_config.get("window", [-1024, 3071]) lo, hi = float(window[0]), float(window[1]) ct_np = sitk.GetArrayFromImage(ct.cube_hu).astype(np.float32) # (Z, Y, X) ct_np = (np.clip(ct_np, lo, hi) - lo) / (hi - lo) # Zero (= window minimum) outside the body so padding and exterior match body = next((v for v in cst.vois if v.name.upper() == "BODY"), None) if body is not None: body_np = sitk.GetArrayViewFromImage(body.mask).astype(np.float32) ct_np = ct_np * (body_np > 0) ct_t = torch.from_numpy(ct_np).unsqueeze(0).unsqueeze(0).to(self.device) self._ct_tensor = F.grid_sample( ct_t, coords, mode="bilinear", padding_mode="zeros", align_corners=True ) def _prepare_masks(self, cst, cst_masks: Optional[list[str]], coords: "torch.Tensor") -> None: """Resample the requested VOI masks to the model grid (nearest neighbor).""" if cst_masks is not None: vois = [] for name in cst_masks: voi = next((v for v in cst.vois if v.name.lower() == name.lower()), None) if voi is None: available = [v.name for v in cst.vois] raise ValueError(f"VOI '{name}' not found. Available: {available}") vois.append(voi) else: vois = list(cst.vois) channels = [] for voi in vois: mask_np = sitk.GetArrayViewFromImage(voi.mask).astype(np.float32) mask_t = torch.from_numpy(mask_np).unsqueeze(0).unsqueeze(0).to(self.device) channels.append( F.grid_sample( mask_t, coords, mode="nearest", padding_mode="zeros", align_corners=True ) ) mask = torch.cat(channels, dim=1) # (1, C, mZ, mY, mX) if self.mask_config.get("collapse", self.mask_config.get("collaps", False)): mask = mask.amax(dim=1, keepdim=True) self._mask_tensor = mask # `extract` multiplies the single-channel dose by the mask; with more than # one mask channel this would broadcast the dose to C channels and break the # single-dose-channel model. Fail early and clearly instead. if self.dose_config.get("extract", False) and mask.shape[1] > 1: raise ValueError( "dose 'extract' requires a single mask channel; got " f"{mask.shape[1]} channels. Set mask 'collapse': true or pass a single VOI." )