Florian Leininger

updating preprocessor for faster computation of gradients

c676516 16 days ago

14.1 kB

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

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