File size: 2,025 Bytes
155dd99 5b483d8 155dd99 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 | """Preprocessor for the ConvBayes_new dose-calc model (single-channel dose output).
Builds a depth-as-channel ``(batch, 1+depth, H, W)`` input tensor from a CT cuboid and
per-bixel energy, and rescales the raw network output back to physical dose [Gy].
"""
import array_api_compat
import torch
from pyRadPlan.ai_models import BasePreprocessor
from pyRadPlan.core.xp_utils import to_namespace
class ConvDoseSinglePreprocessor(BasePreprocessor):
"""CT+energy -> network input; network output -> physical dose."""
def preprocess(self, inputs: dict) -> torch.Tensor:
"""Assemble the depth-as-channel network input from CT cuboid + energy."""
cfg = self.config
ct_cuboid = inputs["ct_cuboid"]
energy = inputs["energy"]
clip_lo, clip_hi = cfg["ct_clip_range"]
xp = array_api_compat.array_namespace(ct_cuboid)
ct_norm = (xp.clip(ct_cuboid, clip_lo, clip_hi) + cfg["ct_offset"]) / cfg["ct_scale"]
energy_norm = energy / cfg["max_energy"]
t = to_namespace("torch", ct_norm).float().detach()
e_t = to_namespace("torch", energy_norm).float().detach()
if e_t.ndim == 0:
e_t = e_t.view(1)
batch_size = e_t.shape[0]
t = t.expand(batch_size, -1, -1, -1)
height, width = t.shape[-2], t.shape[-1]
energy_map = e_t.view(-1, 1, 1, 1).expand(-1, 1, height, width)
return torch.cat((energy_map, t), dim=1)
def postprocess(self, outputs: torch.Tensor) -> dict:
"""Rescale the normalized network output back to physical dose [Gy]."""
return {"physical_dose": outputs * self.config["max_dose"]}
def predict(self, model: torch.nn.Module, model_input: torch.Tensor) -> dict:
"""Deterministic single forward pass. Assumes ``model_input`` is already preprocessed."""
model_input = model_input.to(next(model.parameters()).device)
with torch.inference_mode():
raw_output = model(model_input)
return self.postprocess(raw_output)
|